#/usr/local/bin/python


import ABAC
import sys
import os, os.path
import re

from string import join
from base64 import b64encode, b64decode
from xml.parsers.expat import ParserCreate
from tempfile import mkdtemp
from shutil import rmtree

class proof:
    """
    Class to make manipulating proofs a little simpler.  Keeps track of the
    target and does some of the credential formatting.
    """
    def __init__(self, me, p, a, c):
	"""
	P is the principal (string), a is the attribute to prove (string)
	and c is the list of abac certificates that prove (or partially
	prove) the connection.
	"""
	self.prover = "%s" % me
	self.principal = "%s" % p
	self.attribute = a
	self.creds = c or []

    def creds_to_certs(self):
	ids = set([c.issuer_cert() for c in self.creds])
	attrs = set([c.attribute_cert() for c in self.creds])
	out_ids = []
	d = mkdtemp()
	for i, c in enumerate(ids):
	    try:
		der_name = os.path.join(d, 'file%04d.der' % i)
		pem_name = os.path.join(d, 'file%04d.pem' % i)
		f = open(der_name, 'w')
		f.write(c)
		f.close()

		cid = ABAC.ID(der_name)

		tf = open(pem_name, 'w')
		cid.write_cert(tf)
		tf.close()
		tf = open(pem_name, 'r')
		out_ids.append(tf.read())
		tf.close()
	    except EnvironmentError, e:
		print "Certs to creds Error: %s" % e
		pass
	rmtree(d)

	return out_ids + list(attrs)

    def __str__(self):
	"""
	User friendly output format
	"""
	rv = "prover: %s\nprincipal: %s\nattribute: %s" % \
		(self.prover, self.principal, self.attribute)
	if self.creds:
	    rv += '\n'
	    rv += join(["%s <- %s"%(c.head().string(), c.tail().string()) \
		for c in self.creds], '\n')
	return rv

    def to_dict(self):
	"""
	Dict with the credentials encoded as X.509 certificates
	"""

	return { 
		'prover': self.prover,
		'principal': self.principal,
		'attribute': self.attribute,
		'credential': self.creds_to_certs(),
		}
    def to_context(self):
	"""
	Return an ABAC context containing just the credentials in the proof
	"""
	ctxt = ABAC.Context()
	for c in self.creds_to_certs():
	    if ctxt.load_id_chunk(c) != ABAC.ABAC_CERT_SUCCESS:
		ctxt.load_attribute_chunk(c)
	return ctxt

    def to_xml(self):
	"""
	Convert the proof to a simple XML format
	"""
	rv = "<proof>\n"
	rv += "<prover>%s</prover>\n" % self.prover
	rv += "<principal>%s</principal>\n" % self.principal
	rv += "<attribute>%s</attribute>\n" % self.attribute
	for c in self.creds_to_certs(): 
	    rv += "<credential>%s</credential>\n" % b64encode(c)
	rv += "</proof>"
	return rv

    def save_to_xml(self, fn, mode='w'):
	"""
	Write an XML reprsentation to fn.  Mode is 'w' or 'a' to overwrite or
	append to the file.
	"""
	f = open(fn, mode)
	print >>f, self.to_xml()
	f.close()

    @staticmethod
    def from_xml(filename=None, file=None, string=None):
	"""
	Create a proof from the XML representation
	"""
	class parse_proof:
	    def element_end(self, name):
		if name =='proof':
		    self.proof = proof(self.me, self.pr, self.attr, 
			    self.context.credentials())
		elif name =='prover':
		    self.me = self.chars.strip()
		elif name =='principal':
		    self.pr = self.chars.strip()
		elif name == 'attribute':
		    self.attr = self.chars.strip()
		elif name == 'credential':
		    c = b64decode(self.chars)
		    if self.context.load_id_chunk(c) != ABAC.ABAC_CERT_SUCCESS:
			self.context.load_attribute_chunk(c)
		else:
		    print "Doh!"

	    def element_start(self, name, attrs):
		self.chars = ''

	    def cdata(self, data):
		self.chars += data


	    def __init__(self):
		self.p = ParserCreate()
		self.me = None
		self.pr = None
		self.attr = None
		self.proof = None
		self.context = ABAC.Context()
		self.chars = ''

		self.p.CharacterDataHandler = self.cdata
		self.p.EndElementHandler = self.element_end
		self.p.StartElementHandler = self.element_start

	pp = parse_proof()
	if filename: pp.p.ParseFile(open(filename, 'r'))
	elif file: pp.p.ParseFile(file)
	elif string: pp.p.Parse(string, True)
	else: print "Doh1"

	return pp.proof
    @staticmethod
    def pem_to_der(c):
	pat = '-----BEGIN CERTIFICATE-----(.*)-----END CERTIFICATE'
	m = re.match(pat, c, re.DOTALL)
	if m: return b64decode(m.group(1))
	else: return c


    @staticmethod
    def from_dict(d):
	"""
	Reverse the to_dict process
	"""
	if d: 
	    if 'prover' in d and 'principal' in d and 'attribute' in d:
		ac = ABAC.Context()
		self = proof(d['prover'], d['principal'], d['attribute'], None)
		for c in d.get('credential', []):
		    c = self.pem_to_der(c)
		    if ac.load_id_chunk(c)!= ABAC.ABAC_CERT_SUCCESS:
			ac.load_attribute_chunk(c)
		self.creds = ac.credentials()
		return self
	    else:
		return None
	else:
	    return None