Python SAML parser and validator

samlparser.py

from base64 import b64decode

from defusedxml.lxml import fromstring
from signxml import xmldsig


class SAMLParser(object):

    def __init__(self, rawresp):
        self.b64 = rawresp
        self.xml = b64decode(rawresp)
        self.root = fromstring(self.xml)

        # Namespace mappings
        self.samlp_ns = self.root.nsmap['samlp']
        self.saml_ns = self.root.nsmap['saml']
        self.ds_ns = 'http://www.w3.org/2000/09/xmldsig#'

    def validate(self, cert=None):
        """Validate saml signature with xmldsig"""
        assertion = self.root.find("{%s}Assertion" % self.saml_ns)

        # etree can't seem to find the ds:Tag namespace from root so we need the assertion first.
        signature = assertion.find("{%s}Signature" % self.ds_ns)
        if not cert:
            cert = signature.find('.{%s}KeyInfo//{%s}X509Certificate' % (self.ds_ns,
                                                                         self.ds_ns)).text
        try:
            assertion_data = xmldsig(assertion).verify(x509_cert=cert)
        except:
            raise SAMLInvalidError

    def groups(self):
        groups = self.root.findall(".//{%s}Attribute[@Name='groups']/"
                                   "{%s}AttributeValue" % (self.saml_ns,
                                                           self.saml_ns))

        return [group.text for group in groups]

    def username(self):
        return self.root.find(".//{%s}Attribute[@Name='userLoginID']/"
                              "{%s}AttributeValue" % (self.saml_ns, self.saml_ns)).text

    def conditions(self):
        return self.root.find(".//{%s}Conditions" % self.saml_ns).attrib


class SAMLInvalidError(Exception):
    pass


class SAMLParseError(Exception):
    pass