source: fedd/federation/proof.py @ 51d3aa0

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