#!/usr/local/bin/python

import re
import os
import string

from tempfile import mkstemp

def abac_pem_type(cert):
    key_re = re.compile('\s*-----BEGIN RSA PRIVATE KEY-----$')
    cert_re = re.compile('\s*-----BEGIN CERTIFICATE-----$') 
    type = None

    f = open(cert, 'r')
    for line in f:
	if key_re.match(line):
	    if type is None: type = 'key'
	    elif type == 'cert': type = 'both'
	elif cert_re.match(line):
	    if type is None: type = 'cert'
	    elif type == 'key': type = 'both'
	if type == 'both': break
    f.close()
    return type

def abac_split_cert(cert, keyfile=None, certfile=None):
    """
    Split the certificate file in cert into a certificate file and a key file
    in cf and kf respectively.  The ABAC tools generally cannot handle combined
    certificates/keys.  If kf anc cf are given, they are used, otherwise tmp
    files are created.  Created tmp files must be deleted.  Problems opening or
    writing files will cause exceptions.
    """
    class diversion:
	'''
	Wraps up the reqular expression to start and end a diversion, as well as
	the open file that gets the lines.  If fd is passed in, use that system
	file (probably from a mkstemp.  Otherwise open the given filename.
	'''
	def __init__(self, start, end, fn=None, fd=None):
	    self.start = re.compile(start)
	    self.end = re.compile(end)

	    if not fd:
		# Open the file securely with minimal permissions. NB file
		# cannot exist before this call.
		fd = os.open(fn,
		    (os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_EXCL), 0600)

	    self.f = os.fdopen(fd, 'w')

	def close(self):
	    self.f.close()

    if not keyfile:
	kf, rkeyfile = mkstemp(suffix=".pem")
    else:
	kf, rkeyfile = None, keyfile

    if not certfile:
	cf, rcertfile = mkstemp(suffix=".pem")
    else:
	cf, rcertfile = None, certfile

    # Initialize the diversions
    divs = [diversion(s, e, fn=pfn, fd=pfd ) for s, e, pfn, pfd in (
	('\s*-----BEGIN RSA PRIVATE KEY-----$', 
	    '\s*-----END RSA PRIVATE KEY-----$',
	    keyfile, kf),
	('\s*-----BEGIN CERTIFICATE-----$', 
	    '\s*-----END CERTIFICATE-----$',
	    certfile, cf))]

    # walk through the file, beginning a diversion when a start regexp
    # matches until the end regexp matches.  While in the two regexps,
    # print each line to the open diversion file (including the two
    # matches).
    active = None
    f = open(cert, 'r')
    for l in f:
	if active:
	    if active.end.match(l):
		print >>active.f, l,
		active = None
	else:
	    for d in divs:
		if d.start.match(l):
		    active = d
		    break
	if active: print >>active.f, l,

    # This is probably unnecessary.  Close all the diversion files.
    for d in divs: d.close()
    return rkeyfile, rcertfile

def abac_context_to_creds(context):
    """
    Pull all the credentials out of the context and return 2  lists of the
    underlying credentials in an exportable format, IDs and attributes.
    There are no duplicates in the lists.
    """
    ids, attrs = set(), set()
    # This should be a one-iteration loop
    for c in context.credentials():
	ids.add(str(c.issuer_cert()))
	attrs.add(str(c.attribute_cert()))

    return list(ids), list(attrs)



def remove_dirs(dname):
    """
    Remove the directory tree and all files rooted at dir.  Log any errors,
    but continue.
    """
    for path, dirs, files in os.walk(dname, topdown=False):
	for f in files:
	    os.remove(os.path.join(path, f))
	for d in dirs:
	    os.rmdir(os.path.join(path, d))
    os.rmdir(dname)