"""Minimal SAML 2.0 metadata XML parser for the SSO grant flow.

When an employee uploads a vendor-supplied ``EntityDescriptor`` XML
file with their SSO grant request, we use this parser to derive the
three discrete fields the Cloud Identity ``inboundSamlSsoProfiles``
API needs:

* ``entity_id``  — the ``entityID`` attribute on the root
  ``<EntityDescriptor>`` element. This is the SP's own identifier
  the IdP (Google) uses as the audience in the SAML response.
* ``acs_url``    — the first ``<AssertionConsumerService>`` Location
  attribute (preferring HTTP-POST binding). This is where Google
  POSTs the signed SAML response after authentication.
* ``x509_cert``  — the first signing-key X.509 certificate, base64-
  encoded inside ``<X509Certificate>``. We re-wrap it as a standard
  PEM block so the downstream ``create_inbound_saml_sso_profile``
  call gets a format Google accepts.

Parser hardening: we use ``defusedxml.ElementTree.fromstring``
rather than the stdlib ``xml.etree.ElementTree.fromstring`` because
Python's docs explicitly warn the stdlib parser is "not secure
against maliciously constructed data" — it's vulnerable to billion-
laughs (entity expansion DoS) and quadratic blowup (large entity
expansion). The 1 MB size cap upstream limits the input but cannot
prevent a 10 KB carefully-nested entity definition from expanding
to gigabytes during parse. ``defusedxml`` is already a transitive
dep (pyproject.toml); using it costs nothing and closes the entire
class of XML attacks. The ``Element`` instances it returns are
standard ``xml.etree.ElementTree.Element`` so the helpers below
operate unchanged.
"""
from __future__ import annotations

import re
from dataclasses import dataclass
from xml.etree import ElementTree as ET

from defusedxml.ElementTree import ParseError as _DefusedParseError
from defusedxml.ElementTree import fromstring as _defused_fromstring

# SAML 2.0 metadata namespaces. The vendor file may use either the
# explicit ``md:`` prefix or the unprefixed default namespace; both
# resolve to the same URI below.
_NS = {
    "md": "urn:oasis:names:tc:SAML:2.0:metadata",
    "ds": "http://www.w3.org/2000/09/xmldsig#",
}

# HTTP-POST is the binding Google's IdP uses for the SAML response
# (HTTP-Redirect is used for the AuthnRequest in the other direction
# and lives on a different element). Prefer it strictly to avoid
# pinning the wrong endpoint when a vendor lists multiple bindings.
_PREFERRED_BINDING = "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST"

_BEGIN_CERT = "-----BEGIN CERTIFICATE-----"
_END_CERT = "-----END CERTIFICATE-----"


class SamlMetadataParseError(ValueError):
    """Raised when the supplied XML is not a parseable SAML metadata
    descriptor or the required SP fields are missing.

    Caller (the SSO grant approval executor) is expected to mark the
    request as ``execute_failed`` with this message in the audit so
    the requester can re-upload a corrected file.
    """


@dataclass(frozen=True)
class ParsedSamlMetadata:
    entity_id: str
    acs_url: str
    x509_cert_pem: str | None
    # The descriptor type as found in the file — useful for the
    # audit row so we can tell SP-side metadata (the typical case
    # for Google-as-IdP) from IdP-side metadata (the wrong file the
    # vendor sometimes hands over).
    descriptor_kind: str  # "SPSSODescriptor" | "IDPSSODescriptor"


def parse_saml_metadata(xml_text: str) -> ParsedSamlMetadata:
    """Parse a ``<EntityDescriptor>`` XML and pull out SP fields.

    Raises :class:`SamlMetadataParseError` for any of:

    * malformed XML / non-EntityDescriptor root
    * missing ``entityID``
    * no ``AssertionConsumerService`` element (an SP descriptor
      without one is broken)
    """
    if not xml_text or not xml_text.strip():
        raise SamlMetadataParseError("Empty SAML metadata XML.")
    try:
        root = _defused_fromstring(xml_text)
    except _DefusedParseError as e:
        raise SamlMetadataParseError(
            f"Could not parse SAML metadata XML: {e}",
        ) from e
    except ET.ParseError as e:
        # defusedxml wraps ET.ParseError on lower-level XML syntax
        # errors (bad encoding, truncated tag); surface the same
        # error class to callers.
        raise SamlMetadataParseError(
            f"Could not parse SAML metadata XML: {e}",
        ) from e
    # Match against either the namespaced or unprefixed local name.
    local_name = _local_name(root.tag)
    if local_name != "EntityDescriptor":
        raise SamlMetadataParseError(
            f"Root element is <{local_name}>; expected <EntityDescriptor>.",
        )
    entity_id = root.get("entityID") or ""
    if not entity_id:
        raise SamlMetadataParseError(
            "EntityDescriptor missing required 'entityID' attribute.",
        )

    sp_desc = _find_first(root, "SPSSODescriptor")
    idp_desc = _find_first(root, "IDPSSODescriptor")
    # Use ``is not None`` rather than ``sp_desc or idp_desc`` —
    # ``ElementTree.Element`` evaluates to False when it has no
    # children (deprecated in 3.12, removal scheduled), so an empty
    # SPSSODescriptor would silently fall through to IDPSSODescriptor.
    descriptor = sp_desc if sp_desc is not None else idp_desc
    descriptor_kind = "SPSSODescriptor" if sp_desc is not None else (
        "IDPSSODescriptor" if idp_desc is not None else ""
    )
    if descriptor is None:
        raise SamlMetadataParseError(
            "EntityDescriptor has neither SPSSODescriptor nor "
            "IDPSSODescriptor — malformed metadata file.",
        )

    acs_url = _pick_acs_url(descriptor)
    if not acs_url:
        raise SamlMetadataParseError(
            "No usable AssertionConsumerService Location found "
            "(checked HTTP-POST binding first, then any).",
        )

    cert_pem = _pick_signing_cert_pem(descriptor)

    return ParsedSamlMetadata(
        entity_id=entity_id.strip(),
        acs_url=acs_url.strip(),
        x509_cert_pem=cert_pem,
        descriptor_kind=descriptor_kind,
    )


def _local_name(tag: str) -> str:
    """Strip the ``{namespace}`` prefix ``ElementTree`` adds."""
    if "}" in tag:
        return tag.rsplit("}", 1)[1]
    return tag


def _find_first(parent: ET.Element, local: str) -> ET.Element | None:
    """Find first child by local name, namespace-agnostic."""
    for child in parent.iter():
        if _local_name(child.tag) == local:
            return child
    return None


def _pick_acs_url(descriptor: ET.Element) -> str | None:
    """Walk ``<AssertionConsumerService>`` children, prefer HTTP-POST.

    Accepts the elements in either namespace prefix style. Returns
    the ``Location`` attribute of the chosen element, or None if no
    ACS element exists. Per SAML core, the index-0 ACS is the
    "default" — we honour HTTP-POST > index=0 > anything-else, in
    that order, which matches what Google and most SaaS vendors
    actually want.
    """
    candidates: list[tuple[int, ET.Element]] = []
    for child in descriptor.iter():
        if _local_name(child.tag) != "AssertionConsumerService":
            continue
        try:
            idx = int(child.get("index", "999"))
        except (TypeError, ValueError):
            idx = 999
        candidates.append((idx, child))
    if not candidates:
        return None
    # Prefer HTTP-POST binding, then by index ascending.
    for _, el in sorted(candidates, key=lambda t: t[0]):
        if el.get("Binding") == _PREFERRED_BINDING:
            loc = el.get("Location")
            if loc:
                return loc
    # Fall back to the lowest-index ACS regardless of binding.
    for _, el in sorted(candidates, key=lambda t: t[0]):
        loc = el.get("Location")
        if loc:
            return loc
    return None


def _pick_signing_cert_pem(descriptor: ET.Element) -> str | None:
    """Find the first signing-use ``X509Certificate`` and rewrap as PEM.

    Returns None when the descriptor has no signing cert (some SP
    metadata only carries an ``encryption`` key, which Google does
    not need for inbound SAML profiles).
    """
    # First pass: explicit ``use="signing"`` KeyDescriptor.
    for child in descriptor.iter():
        if _local_name(child.tag) != "KeyDescriptor":
            continue
        if (child.get("use") or "").lower() == "signing":
            cert_text = _first_x509_text(child)
            if cert_text:
                return _wrap_pem(cert_text)
    # Fallback: any KeyDescriptor without a ``use`` attribute is
    # implicitly both signing and encryption per SAML spec.
    for child in descriptor.iter():
        if _local_name(child.tag) != "KeyDescriptor":
            continue
        if not (child.get("use") or ""):
            cert_text = _first_x509_text(child)
            if cert_text:
                return _wrap_pem(cert_text)
    return None


def _first_x509_text(key_descriptor: ET.Element) -> str | None:
    for child in key_descriptor.iter():
        if _local_name(child.tag) == "X509Certificate":
            return (child.text or "").strip()
    return None


def _wrap_pem(raw_b64_cert: str) -> str:
    """Wrap a bare base64 cert in standard PEM headers + 64-char lines.

    SAML metadata stores the cert as bare base64 (no header / no
    line breaks). Google's ``inboundSamlSsoProfiles`` API wants the
    full PEM envelope — re-wrap rather than trying to skip headers.
    """
    body = re.sub(r"\s+", "", raw_b64_cert)
    if not body:
        return ""
    # 64-char lines are the PEM convention; some downstream parsers
    # (notably older OpenSSL) reject lines > 76 chars.
    lines = [body[i: i + 64] for i in range(0, len(body), 64)]
    return _BEGIN_CERT + "\n" + "\n".join(lines) + "\n" + _END_CERT


__all__ = [
    "ParsedSamlMetadata",
    "SamlMetadataParseError",
    "parse_saml_metadata",
]