vignette/apps/api/app/saml.py

"""Minimal SAML SP helpers for local fixture authentication tests.

This module intentionally implements only the Redirect-binding AuthnRequest and
unsigned fixture ACS parsing needed for backend proof. Signed production SAML
assertion verification is not implemented here.
"""

from __future__ import annotations

import base64
import html
import uuid
import zlib
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Iterable
from urllib.parse import urlsplit, urlunsplit, urlencode
from xml.etree import ElementTree


SAML_PROTOCOL_NS = "urn:oasis:names:tc:SAML:2.0:protocol"
SAML_ASSERTION_NS = "urn:oasis:names:tc:SAML:2.0:assertion"
SAML_ATTRIBUTE_ROLE_NAMES = {
    "role",
    "roles",
    "groups",
    "memberOf",
    "http://schemas.microsoft.com/ws/2008/06/identity/claims/role",
}
SAML_ATTRIBUTE_EMAIL_NAMES = {
    "email",
    "mail",
    "emailaddress",
    "EmailAddress",
    "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress",
}
SAML_ATTRIBUTE_DISPLAY_NAME_NAMES = {
    "display_name",
    "displayName",
    "name",
    "cn",
    "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name",
}
SAML_ATTRIBUTE_COHORT_NAMES = {
    "cohort",
    "cohort_id",
    "cohortIds",
    "cohorts",
    "class",
    "group",
    "groups",
}


@dataclass(frozen=True, slots=True)
class SamlIdentity:
    subject: str
    email: str
    display_name: str
    role_hint: str | None = None
    cohort_hint: str | None = None


def acs_url_for_entity_id(entity_id: str) -> str:
    parsed = urlsplit(entity_id.strip())
    if parsed.scheme and parsed.netloc:
        path = parsed.path.rstrip("/")
        if path.endswith("/metadata"):
            path = path[: -len("/metadata")]
        return urlunsplit((parsed.scheme, parsed.netloc, f"{path}/acs", "", ""))
    value = entity_id.strip().rstrip("/")
    if value.endswith("/metadata"):
        value = value[: -len("/metadata")]
    return value + "/acs"


def build_authn_request(
    *,
    sp_entity_id: str,
    sso_url: str,
    acs_url: str,
) -> tuple[str, str]:
    request_id = "_" + uuid.uuid4().hex
    issued_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    xml = (
        f'<samlp:AuthnRequest xmlns:samlp="{SAML_PROTOCOL_NS}" '
        f'xmlns:saml="{SAML_ASSERTION_NS}" ID="{request_id}" Version="2.0" '
        f'IssueInstant="{issued_at}" Destination="{html.escape(sso_url, quote=True)}" '
        f'AssertionConsumerServiceURL="{html.escape(acs_url, quote=True)}" '
        f'ProtocolBinding="urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST">'
        f"<saml:Issuer>{html.escape(sp_entity_id)}</saml:Issuer>"
        "</samlp:AuthnRequest>"
    )
    return request_id, xml


def redirect_binding_url(*, sso_url: str, authn_request_xml: str, relay_state: str) -> str:
    compressor = zlib.compressobj(wbits=-15)
    deflated = compressor.compress(authn_request_xml.encode("utf-8")) + compressor.flush()
    params = urlencode(
        {
            "SAMLRequest": base64.b64encode(deflated).decode("ascii"),
            "RelayState": relay_state,
        }
    )
    separator = "&" if "?" in sso_url else "?"
    return f"{sso_url}{separator}{params}"


def inflate_redirect_request(encoded_request: str) -> str:
    payload = base64.b64decode(encoded_request)
    return zlib.decompress(payload, wbits=-15).decode("utf-8")


def parse_fixture_response(encoded_response: str) -> SamlIdentity:
    try:
        xml = base64.b64decode(encoded_response).decode("utf-8")
        root = ElementTree.fromstring(xml)
    except Exception as exc:
        raise ValueError("invalid SAMLResponse") from exc

    name_id = _first_text(root, f".//{{{SAML_ASSERTION_NS}}}NameID")
    attributes = _attributes(root)
    email = _first_attribute(attributes, SAML_ATTRIBUTE_EMAIL_NAMES) or name_id
    if not email:
        raise ValueError("email claim is required")

    display_name = (
        _first_attribute(attributes, SAML_ATTRIBUTE_DISPLAY_NAME_NAMES)
        or name_id
        or email
    )
    role_hint = _first_attribute(attributes, SAML_ATTRIBUTE_ROLE_NAMES)
    cohort_hint = _first_attribute(attributes, SAML_ATTRIBUTE_COHORT_NAMES)
    return SamlIdentity(
        subject=name_id or email,
        email=email,
        display_name=display_name or email,
        role_hint=role_hint,
        cohort_hint=cohort_hint or None,
    )


def _first_text(root: ElementTree.Element, selector: str) -> str:
    node = root.find(selector)
    return (node.text or "").strip() if node is not None else ""


def _attributes(root: ElementTree.Element) -> dict[str, list[str]]:
    values: dict[str, list[str]] = {}
    for attribute in root.findall(f".//{{{SAML_ASSERTION_NS}}}Attribute"):
        name = (attribute.attrib.get("Name") or "").strip()
        if not name:
            continue
        collected: list[str] = []
        for value in attribute.findall(f".//{{{SAML_ASSERTION_NS}}}AttributeValue"):
            text = (value.text or "").strip()
            if text:
                collected.append(text)
        if collected:
            values[name] = collected
    return values


def _first_attribute(attributes: dict[str, list[str]], names: Iterable[str]) -> str:
    for name in names:
        values = attributes.get(name)
        if values:
            return values[0]
    lowered = {key.lower(): value for key, value in attributes.items()}
    for name in names:
        values = lowered.get(name.lower())
        if values:
            return values[0]
    return ""