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