#/usr/local/bin/python


import ABAC
import sys
import os, os.path

from string import join
from base64 import b64encode, b64decode
from xml.parsers.expat import ParserCreate

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])
	return list(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 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', []):
		    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