"""Fail-closed SES sender authentication (INFRA-107).

SES Email Receiving *prepends* its own trace headers -- including an
``Authentication-Results`` header whose authserv-id is ``amazonses.com`` --
to the top of the raw MIME it writes to S3. Everything below those
prepended headers (the From header, any additional Authentication-Results
copies) is attacker-controlled, so ONLY the topmost Authentication-Results
header is trusted, and only when its authserv-id is ``amazonses.com``.

An email is accepted only when that header carries ``dkim=pass`` for a
domain in the ``ALLOWED_DKIM_DOMAINS`` allowlist (a comma-separated Lambda
environment variable set by the CDK stack). Every other outcome fails
closed and the email is rejected:

- ``ALLOWED_DKIM_DOMAINS`` unset or empty
- no Authentication-Results header at all
- topmost header unparseable or from an authserv-id other than SES
- no ``dkim=pass`` clause
- ``dkim=pass`` only for domains outside the allowlist

Observed SES format (2026-07-15, both ingest buckets)::

    Authentication-Results: amazonses.com;
     spf=pass (spfCheck: ...) client-ip=...; envelope-from=...; helo=...;
     dkim=pass header.i=@seahaven.com;
     dmarc=none header.from=hxgnsmartcloud.com;

Note SES reports the passing DKIM identity as ``header.i=@<domain>``
(RFC 6376 AUID), not ``header.d=``; the parser accepts both, and treats
``header.d`` (the plain signing domain) as authoritative over ``header.i``
when both are present. The ``header.i`` domain is derived per RFC 6376 as
the part after the AUID's *last* ``@`` -- a ``@`` inside a quoted
local-part (e.g. ``i="@seahaven.com"@attacker.com``) is signer-controlled
label text, never the identity domain, and yields the true signer
(``attacker.com``).

**Clause injection defence (INFRA-107 hardening).** SES echoes several
attacker-controlled SMTP-session tokens into its own Authentication-Results
value as their own semicolon-delimited property clauses -- notably
``envelope-from=<MAIL FROM>``, ``helo=<HELO>`` and ``header.from=<From
domain>``. An RFC 5321 quoted-local-part MAIL FROM may legally contain
spaces and semicolons, e.g.::

    MAIL FROM:<"x; dkim=pass header.i=@amazon.coupahost.com"@attacker.com>

which SES renders verbatim as ``envelope-from="x; dkim=pass
header.i=@amazon.coupahost.com"@attacker.com;``. A naive ``split(";")``
would tear the quoted string apart and manufacture a synthetic
``dkim=pass header.i=@amazon.coupahost.com`` clause out of attacker input.
The parser therefore tokenises per RFC 8601 / RFC 5322 structure: CFWS
comments ``(...)`` are stripped first, and the value is split into clauses
only on semicolons that sit *outside* a quoted-string. A ``;`` inside a
quoted ``pvalue`` stays part of that one property clause and can never be
read as the start of a ``dkim=`` methodspec.
"""

import email.parser
import email.policy
import json
import logging
import os
import re

logger = logging.getLogger()

ALLOWED_DKIM_DOMAINS_ENV = "ALLOWED_DKIM_DOMAINS"
SES_AUTHSERV_ID = "amazonses.com"

# One resinfo clause of an Authentication-Results value, e.g.
# "dkim=pass header.i=@seahaven.com". The clause must START with the
# method=result pair; header.d= / header.i= may appear anywhere after it.
# The result token must be terminated by end-of-clause or whitespace so
# "dkim=pass-anything" can never be read as "pass". (Comments are stripped
# before matching, so the pre-strip "dkim=pass(comment)" form arrives here
# as "dkim=pass ..." and still terminates on whitespace.)
_DKIM_RESULT_RE = re.compile(r"^dkim\s*=\s*([a-z0-9]+)(?=$|\s)", re.IGNORECASE)


def get_allowed_dkim_domains() -> frozenset:
    """Read the DKIM-domain allowlist from the environment (may be empty)."""
    raw = os.environ.get(ALLOWED_DKIM_DOMAINS_ENV, "")
    return frozenset(
        d.strip().lower().lstrip("@").rstrip(".") for d in raw.split(",") if d.strip()
    )


def _unfold(value: str) -> str:
    """Collapse RFC 5322 folding whitespace into single spaces."""
    return re.sub(r"[\r\n\t ]+", " ", value).strip()


def _strip_comments(value: str) -> tuple:
    """Remove RFC 5322 CFWS comments ``(...)`` from an unfolded header value.

    Comments may nest and may contain quoted pairs (``\\)``). A ``(`` that
    appears *inside* a quoted-string is literal text, not a comment start,
    so quoted-strings (which carry attacker-controlled ``pvalue`` content
    such as a quoted MAIL FROM local part) are passed through untouched.
    Dropping comments first means comment-embedded ``header.i=`` / ``;``
    fakes can never influence clause splitting or domain extraction.

    Every removed comment is replaced by a single space so comments act as
    folding whitespace (RFC 5322 CFWS semantics). A ``)`` that appears
    outside any comment (depth 0) is an unmatched close and makes the value
    not well-formed, closing the ``)(`` clause-injection gap.

    CFWS also means a comment abutting a token *ends* that token:
    ``header.d=seahaven.com(note)`` reads as domain ``seahaven.com``, where
    the pre-CFWS gluing behaviour read ``seahaven.comnote`` (rejected).
    This is deliberate -- ``header.d`` is written verbatim by SES from a
    *verified* signature's ``d=`` tag, so a comment in that position is
    never attacker-controlled.

    Returns ``(stripped_text, well_formed)``. ``well_formed`` is False when
    the value ends inside an unterminated comment or quoted-string, or when
    an unmatched ``)`` appears at the top level. Callers reject on
    ``not well_formed`` so a malformed header (which could otherwise be
    mis-tokenised) fails closed rather than being partially parsed.
    """
    out = []
    depth = 0  # comment nesting depth
    in_quote = False  # inside a quoted-string (only tracked at depth 0)
    extra_close = False  # SES-AR-01: ")" at depth 0 (no matching open)
    i = 0
    n = len(value)
    while i < n:
        c = value[i]
        if depth > 0:
            # Inside a comment: only quoted-pairs and nested parens matter.
            if c == "\\":
                i += 2
                continue
            if c == "(":
                depth += 1
            elif c == ")":
                depth -= 1
                if depth == 0:
                    # SES-AR-02: comment acts as folding whitespace (RFC 5322).
                    out.append(" ")
            i += 1
            continue
        if in_quote:
            out.append(c)
            if c == "\\" and i + 1 < n:
                out.append(value[i + 1])
                i += 2
                continue
            if c == '"':
                in_quote = False
            i += 1
            continue
        # Normal context (outside any comment or quoted-string).
        if c == "(":
            depth += 1
            i += 1
            continue
        if c == ")":
            # SES-AR-01: ")" at depth 0 is an unmatched close -- refuse to
            # tokenise rather than letting a ")(...)" pair manufacture a
            # gap where attacker text leaks out at depth 0.
            extra_close = True
            i += 1
            continue
        if c == '"':
            in_quote = True
        out.append(c)
        i += 1
    well_formed = depth == 0 and not in_quote and not extra_close
    return "".join(out), well_formed


def _split_clauses(value: str) -> list:
    """Split a comment-free header value into clauses on top-level ``;``.

    A semicolon inside a quoted-string is preserved as part of the clause,
    so an attacker-controlled quoted ``pvalue`` (e.g. a quoted MAIL FROM
    echoed into ``envelope-from=``) cannot smuggle in a fake ``dkim=pass``
    clause. Callers must run :func:`_strip_comments` first.
    """
    clauses = []
    buf = []
    in_quote = False
    i = 0
    n = len(value)
    while i < n:
        c = value[i]
        if in_quote:
            buf.append(c)
            if c == "\\" and i + 1 < n:
                buf.append(value[i + 1])
                i += 2
                continue
            if c == '"':
                in_quote = False
            i += 1
            continue
        if c == '"':
            in_quote = True
            buf.append(c)
            i += 1
            continue
        if c == ";":
            clauses.append("".join(buf))
            buf = []
            i += 1
            continue
        buf.append(c)
        i += 1
    clauses.append("".join(buf))
    return clauses


def _split_properties(clause: str) -> list:
    """Split a clause into whitespace-separated ``name=pvalue`` tokens.

    Whitespace *inside* a quoted-string does not split, so an RFC 8601 pvalue
    that embeds a quoted-string (e.g. a DKIM AUID with a quoted local-part that
    legally contains spaces) survives as a single token. This is the same
    quoted-string discipline :func:`_split_clauses` applies at the ``;`` level,
    reused here at the token level so ``header.d=`` / ``header.i=`` extraction
    is quoted-string-aware rather than a naive regex grab.
    """
    tokens = []
    buf = []
    in_quote = False
    i = 0
    n = len(clause)
    while i < n:
        c = clause[i]
        if in_quote:
            buf.append(c)
            if c == "\\" and i + 1 < n:
                buf.append(clause[i + 1])
                i += 2
                continue
            if c == '"':
                in_quote = False
            i += 1
            continue
        if c == '"':
            in_quote = True
            buf.append(c)
            i += 1
            continue
        if c.isspace():
            if buf:
                tokens.append("".join(buf))
                buf = []
            i += 1
            continue
        buf.append(c)
        i += 1
    if buf:
        tokens.append("".join(buf))
    return tokens


def _auid_domain(pvalue: str) -> str:
    """Domain of a DKIM AUID (``header.i``) per RFC 6376.

    The identity domain is the part after the *last* ``@`` of the AUID -- but a
    ``@`` inside a quoted local-part is NOT the identity separator. So
    ``"@seahaven.com"@attacker.com`` yields ``attacker.com`` (the real signer),
    not ``seahaven.com``: the ``@seahaven.com`` sits inside the quoted
    local-part and is signer-controlled label text, never the domain. A bare
    unquoted ``@seahaven.com`` still yields ``seahaven.com``. Returns "" when
    there is no top-level ``@`` (no valid domain) or the domain looks malformed.
    """
    last_at = -1
    in_quote = False
    i = 0
    n = len(pvalue)
    while i < n:
        c = pvalue[i]
        if in_quote:
            if c == "\\":
                i += 2
                continue
            if c == '"':
                in_quote = False
            i += 1
            continue
        if c == '"':
            in_quote = True
        elif c == "@":
            last_at = i
        i += 1
    if last_at < 0:
        return ""
    domain = pvalue[last_at + 1 :].lower().rstrip(".")
    # A DKIM domain-name is a plain dot-atom; anything with a residual quote is
    # malformed (or a smuggling attempt) and must not be trusted.
    if not domain or '"' in domain:
        return ""
    return domain


def _plain_domain(pvalue: str) -> str:
    """Domain of ``header.d`` -- the DKIM signing domain, a plain dot-atom.

    SES writes ``header.d`` verbatim from the signature's ``d=`` tag, which is
    never a quoted-string. A residual quote means malformed/smuggled input and
    fails closed.
    """
    domain = pvalue.lower().rstrip(".")
    if not domain or '"' in domain:
        return ""
    return domain


def _clause_signer_domain(clause: str) -> str:
    """Authoritative DKIM signer domain of a ``dkim=pass`` clause, or "".

    ``header.d`` (the signing domain) is authoritative and is preferred when
    present; only when it is absent does this fall back to ``header.i`` and
    derive the domain from the AUID's post-final-``@`` part. Both lookups run
    over :func:`_split_properties` tokens, so a ``header.d=``/``header.i=``
    literal smuggled *inside* another property's quoted pvalue is confined to
    that one token and can never be read as a top-level property.
    """
    header_d = None
    header_i = None
    for token in _split_properties(clause):
        name, sep, val = token.partition("=")
        if not sep:
            continue
        key = name.strip().lower()
        if key == "header.d" and header_d is None:
            header_d = val
        elif key == "header.i" and header_i is None:
            header_i = val
    if header_d is not None:
        return _plain_domain(header_d)
    if header_i is not None:
        return _auid_domain(header_i)
    return ""


def parse_authentication_results(value: str) -> tuple:
    """Parse one Authentication-Results header value.

    Returns ``(authserv_id, passing_dkim_domains)`` where the domains are the
    authoritative signer domains of every ``dkim=pass`` clause -- ``header.d``
    when present, else the ``header.i`` AUID's post-final-``@`` domain (see
    :func:`_clause_signer_domain`). Malformed input yields ``("", frozenset())``,
    which callers treat as a rejection.

    Tokenisation is comment- and quoted-string-aware (RFC 8601 / RFC 5322) at
    the ``;`` (clause), whitespace (property) and ``@`` (AUID domain) levels, so
    attacker-controlled tokens SES echoes into its header (envelope-from, helo,
    header.from) cannot be split into a forged ``dkim=pass`` clause and a quoted
    AUID local-part cannot masquerade as the identity domain.
    """
    text, well_formed = _strip_comments(_unfold(value))
    if not well_formed:
        # Unbalanced quotes/comments: refuse to guess how to tokenise it.
        return "", frozenset()
    clauses = [c.strip() for c in _split_clauses(text)]
    if not clauses or not clauses[0]:
        return "", frozenset()

    # First clause is the authserv-id, optionally followed by a version
    # token ("amazonses.com 1"); take only the first token.
    authserv_id = clauses[0].split()[0].strip('"').lower()

    passing = set()
    for clause in clauses[1:]:
        match = _DKIM_RESULT_RE.match(clause)
        if not match or match.group(1).lower() != "pass":
            continue
        domain = _clause_signer_domain(clause)
        if domain:
            passing.add(domain)
    return authserv_id, frozenset(passing)


def evaluate_sender_authentication(raw_email: bytes, allowed_domains) -> tuple:
    """Evaluate the SES-stamped verdicts in a raw MIME message.

    Returns ``(accepted, reason, detail)``. Pure function of its inputs so
    it can be unit-tested without touching the environment.
    """
    if not allowed_domains:
        return False, "allowlist_not_configured", {}

    try:
        # compat32 keeps header values as raw strings (we unfold ourselves)
        # and never raises on structurally odd headers; headersonly avoids
        # parsing the body at all.
        msg = email.parser.BytesParser(policy=email.policy.compat32).parsebytes(
            raw_email, headersonly=True
        )
    except Exception:
        return False, "unparseable_message", {}

    ar_headers = msg.get_all("Authentication-Results") or []
    if not ar_headers:
        return False, "authentication_results_missing", {}

    # SES prepends its trace headers, so index 0 is the SES-stamped copy.
    # Any Authentication-Results header further down arrived inside the
    # message (attacker-suppliable) and is deliberately ignored.
    #
    # The parse is wrapped so an unexpected parser exception fails CLOSED
    # (rejected, structured reason) instead of propagating out of the handler
    # into Lambda's async retries / DLQ on attacker-crafted input.
    try:
        authserv_id, passing = parse_authentication_results(str(ar_headers[0]))
    except Exception:
        return False, "authentication_results_unparseable", {}
    detail = {
        "authserv_id": authserv_id,
        "passing_dkim_domains": sorted(passing),
    }

    if authserv_id != SES_AUTHSERV_ID:
        return False, "untrusted_authserv_id", detail
    if not passing:
        return False, "no_passing_dkim_signature", detail

    matched = passing & set(allowed_domains)
    if not matched:
        return False, "dkim_domain_not_allowlisted", detail

    detail["matched_domains"] = sorted(matched)
    return True, "authenticated", detail


def authenticate_inbound_email(raw_email: bytes, s3_key: str) -> bool:
    """Fail-closed gate used by the S3-triggered handlers.

    On rejection: logs a structured warning with the reason and S3 key and
    returns False. Callers skip the message and return normally, so
    rejected mail never errors the invocation (no retries, no DLQ spam).
    """
    allowed = get_allowed_dkim_domains()
    accepted, reason, detail = evaluate_sender_authentication(raw_email, allowed)
    if accepted:
        logger.info(json.dumps({"event": "sender_auth_ok", "s3_key": s3_key, **detail}))
        return True
    logger.warning(
        json.dumps(
            {
                "event": "sender_auth_rejected",
                "reason": reason,
                "s3_key": s3_key,
                "allowed_dkim_domains": sorted(allowed),
                **detail,
            }
        )
    )
    return False