"""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.

**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)
_HEADER_D_RE = re.compile(r"header\.d\s*=\s*\"?([^\s\";]+)", re.IGNORECASE)
_HEADER_I_RE = re.compile(r"header\.i\s*=\s*\"?([^\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.

    Returns ``(stripped_text, well_formed)``. ``well_formed`` is False when
    the value ends inside an unterminated comment or quoted-string, i.e. the
    parens/quotes are unbalanced. 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)
    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
            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 == '"':
            in_quote = True
        out.append(c)
        i += 1
    well_formed = depth == 0 and not in_quote
    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 parse_authentication_results(value: str) -> tuple:
    """Parse one Authentication-Results header value.

    Returns ``(authserv_id, passing_dkim_domains)`` where the domains are
    the lowercased d=/i= domains of every ``dkim=pass`` clause. Malformed
    input yields ``("", frozenset())``, which callers treat as a rejection.

    Tokenisation is comment- and quoted-string-aware (RFC 8601 / RFC 5322)
    so attacker-controlled tokens SES echoes into its header (envelope-from,
    helo, header.from) cannot be split into a forged ``dkim=pass`` clause.
    """
    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
        d_match = _HEADER_D_RE.search(clause)
        if d_match:
            passing.add(d_match.group(1).lower().rstrip("."))
            continue
        i_match = _HEADER_I_RE.search(clause)
        if i_match:
            # AUID (header.i) is "local-part@domain" or "@domain";
            # the signing domain is everything after the last "@".
            identity = i_match.group(1).lower().rstrip(".")
            domain = identity.rsplit("@", 1)[-1]
            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.
    authserv_id, passing = parse_authentication_results(str(ar_headers[0]))
    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