wp-ak-saml/SAML User Mapping.py
2025-04-15 21:29:38 -05:00

59 lines
1.9 KiB
Python

from lxml import etree
NS_SAML_ASSERTION = "urn:oasis:names:tc:SAML:2.0:assertion"
assertion = root.find(f"{{{NS_SAML_ASSERTION}}}Assertion")
if assertion is None:
raise ValueError("Assertion element not found")
attribute_statement = assertion.find(f"{{{NS_SAML_ASSERTION}}}AttributeStatement")
if attribute_statement is None:
raise ValueError("AttributeStatement element not found")
def _name_id():
name_id = element.find(f"{{{NS_SAML_ASSERTION}}}NameID")
name_qualifier = name_id.attrib.get("NameQualifier", "")
spname_qualifier = name_id.attrib.get("SPNameQualifier", "")
return f"{name_qualifier}!{spname_qualifier}!{name_id.text}"
_extractors = {"urn:oid:1.3.6.1.4.1.5923.1.1.1.10": _name_id}
# extract attribute values from attribute statements into attributes dictionary
attributes = {}
for element in attribute_statement.findall(f"{{{NS_SAML_ASSERTION}}}Attribute"):
name = element.attrib["Name"]
element_value = element.find(f"{{{NS_SAML_ASSERTION}}}AttributeValue")
extractor_fun = _extractors.get(name, lambda el: el.text)
attributes[name] = extractor_fun(element_value)
def extract(lookup_list: tuple, is_mandatory: bool = False, def_value=None):
for attrib_name in lookup_list:
val = attributes.get(attrib_name, None)
if val:
return val
if is_mandatory:
raise ValueError(
f'No SAML source attribute found for mandatory value: "{name}"'
)
return def_value
uid = extract(("urn:oid:0.9.2342.19200300.100.1.1", "urn:mace:dir:attribute-def:uid"), is_mandatory=True)
name = extract(("urn:oid:2.5.4.41", "urn:mace:dir:attribute-def:name"))
displayName = extract(("urn:oid:2.16.840.1.113730.3.1.241", "urn:mace:dir:attribute-def:displayName"))
email = extract(("urn:oid:1.2.840.113549.1.9.1", "urn:mace:dir:attribute-def:email"))
# User fields.
results = {
"username": uid,
"email": email,
"uid": uid,
"name": displayName,
"is_active": True,
"is_staff": False,
}
return results