saml/sp/assertion_consumer.py

"""
saml/sp/assertion_consumer.py
ACS endpoint: receive Response, validate full chain, run NAND gate, emit audit.
"""

import base64
import hashlib
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional
from xml.etree import ElementTree as ET

from crypto.xml_dsig import verify_signature, SignatureVerificationError
from nand.gate import NANDGate, GateResult
from audit.worm import WORMAuditChain, AuditRecord, attributes_hash
from audit.replay import ReplayStore, ReplayDetectedError

SAML_NS  = "urn:oasis:names:tc:SAML:2.0:assertion"
SAMLP_NS = "urn:oasis:names:tc:SAML:2.0:protocol"
XENC_NS  = "http://www.w3.org/2001/04/xmlenc#"
_SA  = "{%s}" % SAML_NS
_SAP = "{%s}" % SAMLP_NS
_XE  = "{%s}" % XENC_NS

STATUS_SUCCESS   = "urn:oasis:names:tc:SAML:2.0:status:Success"
CLOCK_SKEW_SECS  = 300


class SAMLParseError(Exception):    pass
class SAMLValidationError(Exception): pass


@dataclass
class ConsumeResult:
    valid:         bool
    subject:       str
    attributes:    dict
    semantic_hash: str
    assertion_id:  str
    rejection_reason: Optional[str] = None


class AssertionConsumer:
    def __init__(
        self,
        idp_certificate,
        sp_private_key,
        sp_entity_id:  str,
        acs_url:       str,
        nand_gate:     NANDGate,
        replay_store:  ReplayStore,
        audit_chain:   WORMAuditChain,
        clock_skew:    int = CLOCK_SKEW_SECS,
    ):
        self.idp_cert     = idp_certificate
        self.sp_key       = sp_private_key
        self.sp_entity_id = sp_entity_id
        self.acs_url      = acs_url
        self.gate         = nand_gate
        self.replay       = replay_store
        self.audit        = audit_chain
        self.clock_skew   = clock_skew

    def receive(self, saml_response_b64: str) -> ConsumeResult:
        raw_xml = base64.b64decode(saml_response_b64)
        root    = self._parse_response(raw_xml)

        self._validate_response(root)
        assertion = self._extract_assertion(root)
        self._validate_assertion(assertion)

        subject = self._extract_subject(assertion)
        attrs   = self._extract_attributes(assertion)
        assertion_id   = assertion.get('ID', '_unknown')
        issuer_el      = assertion.find(_SA + 'Issuer')
        issuer         = issuer_el.text.strip() if issuer_el is not None else ''
        issue_instant  = assertion.get('IssueInstant', '')

        # Replay check — use SubjectConfirmationData.NotOnOrAfter
        not_on_or_after = self._get_not_on_or_after(assertion)
        self.replay.check_and_store(assertion_id, not_on_or_after)

        gate_result: GateResult = self.gate.evaluate(
            attrs, assertion_id, issuer, issue_instant
        )

        event = 'ASSERTION_ACCEPTED' if gate_result.valid else 'ASSERTION_REJECTED'
        record = AuditRecord(
            timestamp=datetime.now(timezone.utc).isoformat(),
            semantic_hash=gate_result.semantic_hash,
            assertion_id=assertion_id,
            issuer=issuer,
            subject=subject,
            gate_result=gate_result.valid,
            entropy=gate_result.entropy,
            attributes_hash=attributes_hash(attrs),
            event=event,
        )
        self.audit.append(record)

        return ConsumeResult(
            valid=gate_result.valid,
            subject=subject,
            attributes=attrs,
            semantic_hash=gate_result.semantic_hash,
            assertion_id=assertion_id,
            rejection_reason=gate_result.rejection_reason,
        )

    # ── Internal validation ───────────────────────────────────────────────────

    def _parse_response(self, raw_xml: bytes) -> ET.Element:
        try:
            return ET.fromstring(raw_xml)
        except ET.ParseError as e:
            raise SAMLParseError(f"XML parse error: {e}")

    def _validate_response(self, root: ET.Element) -> None:
        # Status
        status_code = root.find(f".//{_SAP}StatusCode")
        if status_code is None or status_code.get('Value') != STATUS_SUCCESS:
            raise SAMLValidationError("Response status is not Success")

        # Destination
        dest = root.get('Destination', '')
        if dest and dest != self.acs_url:
            raise SAMLValidationError(f"Destination mismatch: {dest} != {self.acs_url}")

        # IssueInstant clock skew
        issue_instant = root.get('IssueInstant', '')
        if issue_instant:
            self._check_clock_skew(issue_instant)

    def _extract_assertion(self, root: ET.Element) -> ET.Element:
        # Try encrypted first
        enc = root.find(f".//{_SA}EncryptedAssertion")
        if enc is not None:
            from crypto.xml_encrypt import decrypt_assertion
            return decrypt_assertion(enc, self.sp_key)

        assertion = root.find(_SA + 'Assertion')
        if assertion is None:
            assertion = root.find(f".//{_SA}Assertion")
        if assertion is None:
            raise SAMLValidationError("No saml:Assertion found in Response")
        return assertion

    def _validate_assertion(self, assertion: ET.Element) -> None:
        now = datetime.now(timezone.utc)
        skew = timedelta(seconds=self.clock_skew)

        # Conditions timing
        conditions = assertion.find(_SA + 'Conditions')
        if conditions is not None:
            not_before_str = conditions.get('NotBefore', '')
            not_after_str  = conditions.get('NotOnOrAfter', '')
            if not_before_str:
                nb = self._parse_saml_datetime(not_before_str)
                if now < nb - skew:
                    raise SAMLValidationError(f"Assertion not yet valid: NotBefore={not_before_str}")
            if not_after_str:
                na = self._parse_saml_datetime(not_after_str)
                if now > na + skew:
                    raise SAMLValidationError(f"Assertion expired: NotOnOrAfter={not_after_str}")

            # Audience
            audience = assertion.findtext(f".//{_SA}Audience", '')
            if audience and audience.strip() != self.sp_entity_id:
                raise SAMLValidationError(f"Audience mismatch: {audience}")

        # Signature
        verify_signature(assertion, self.idp_cert)

    def _extract_subject(self, assertion: ET.Element) -> str:
        nameid = assertion.find(f".//{_SA}NameID")
        return nameid.text.strip() if nameid is not None and nameid.text else ''

    def _extract_attributes(self, assertion: ET.Element) -> dict:
        attrs = {}
        for attr in assertion.findall(f".//{_SA}Attribute"):
            name   = attr.get('Name', '')
            values = [
                v.text.strip() if v.text else ''
                for v in attr.findall(_SA + 'AttributeValue')
            ]
            if name:
                attrs[name] = values
        return attrs

    def _get_not_on_or_after(self, assertion: ET.Element) -> str:
        scd = assertion.find(f".//{_SA}SubjectConfirmationData")
        if scd is not None:
            return scd.get('NotOnOrAfter', '')
        conditions = assertion.find(_SA + 'Conditions')
        if conditions is not None:
            return conditions.get('NotOnOrAfter', '')
        return ''

    def _check_clock_skew(self, issue_instant: str) -> None:
        now  = datetime.now(timezone.utc)
        skew = timedelta(seconds=self.clock_skew)
        try:
            dt = self._parse_saml_datetime(issue_instant)
        except ValueError:
            return
        if dt > now + skew:
            raise SAMLValidationError(f"IssueInstant too far in future: {issue_instant}")
        if dt < now - timedelta(hours=8):
            raise SAMLValidationError(f"IssueInstant too old: {issue_instant}")

    @staticmethod
    def _parse_saml_datetime(s: str) -> datetime:
        s = s.replace('Z', '+00:00')
        return datetime.fromisoformat(s)