source: fedd/federation/proof.py @ 5ae9d94

#/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