#!/usr/bin/env python3
"""Offline RFC 8617 verifier for the single Google ARC set in the December EML.

This script deliberately supports the narrow evidence object under review:

* one ARC set (i=1);
* RSA-SHA256;
* relaxed/relaxed ARC-Message-Signature canonicalisation; and
* relaxed ARC-Seal canonicalisation.

The embedded ARC public key is the value returned for
``arc-20240605._domainkey.google.com`` by two distinct resolver services in the
retained parsed current-DNS capture made on 2 August 2026.  A successful result
establishes the signature mathematics relative to that audit-date selector
binding.  It is not, by itself, raw resolver evidence, a historical DNS archive,
a finding about exclusive private-key custody or a live-mailbox acquisition
record.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import re
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa


GOOGLE_ARC_PUBLIC_KEY = (
    "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAl3xpv8nzv48sGAt6yT3B"
    "hDpQcLv9Cgnot1arQefERmoFpgygBi73YM3lllM3fb435QBz8Qsu3ZM+xkIM5+UM"
    "9p286xJXI/26vNSbCFEAHeYJmUCJMUP1sX1b5oKzW5NJgojoJNLRulyo6GC6nFpB"
    "Lv6j2hL41GOjg+RdYDxrTcQavsP2DMQLFP/mdsg7zBlfYdg0occfvpJXHYgxN3M1"
    "mYvBvbmnmPIrMtfoMFDdmJCiri3c9X9je2yz/gcpKXoz0n+O9Clk8jVQkrPEcahk"
    "CHPpbe662VLndAYsiD36IyF0+shLOJn0duf2O5eMNlUBwMPX7N589dz3Xy32B0eo"
    "PQIDAQAB"
)


@dataclass
class ArcResult:
    path: str
    bytes: int
    sha256: str
    instance: int
    domain: str
    selector: str
    arc_key_bits: int
    arc_key_spki_sha256: str
    ams_body_hash_expected: str
    ams_body_hash_computed: str
    ams_body_hash_ok: bool
    ams_signed_input_sha256: str
    ams_signature_math_ok: bool
    arc_seal_signed_input_sha256: str
    arc_seal_signature_math_ok: bool
    arc_set_complete: bool
    overall_arc_set_ok: bool
    scope_note: str


def split_message(raw: bytes) -> Tuple[bytes, bytes, bytes]:
    for separator in (b"\r\n\r\n", b"\n\n"):
        offset = raw.find(separator)
        if offset >= 0:
            return raw[:offset], raw[offset + len(separator) :], separator
    raise ValueError("message has no header/body separator")


def parse_headers(header_block: bytes) -> List[Tuple[bytes, bytes]]:
    normalised = header_block.replace(b"\r\n", b"\n")
    parsed: List[Tuple[bytes, bytes]] = []
    current: Optional[List[bytes]] = None
    for line in normalised.split(b"\n"):
        if line[:1] in (b" ", b"\t") and current is not None:
            current[1] += b"\r\n" + line
            continue
        if current is not None:
            parsed.append((current[0], current[1]))
            current = None
        if b":" in line:
            name, value = line.split(b":", 1)
            current = [name, value]
    if current is not None:
        parsed.append((current[0], current[1]))
    return parsed


def relaxed_header(name: bytes, value: bytes) -> bytes:
    unfolded = value.replace(b"\r\n", b"").replace(b"\n", b"")
    unfolded = re.sub(rb"[ \t]+", b" ", unfolded).strip(b" \t")
    return name.lower().strip() + b":" + unfolded + b"\r\n"


def relaxed_body(body: bytes) -> bytes:
    lines = body.replace(b"\r\n", b"\n").split(b"\n")
    canonical = [re.sub(rb"[ \t]+", b" ", line).rstrip(b" \t") for line in lines]
    while canonical and canonical[-1] == b"":
        canonical.pop()
    # RFC 6376 section 3.4.4: an empty relaxed body canonicalises to the null
    # input, not a lone CRLF.
    return b"\r\n".join(canonical) + (b"\r\n" if canonical else b"")


def tags(value: bytes) -> Dict[str, bytes]:
    compact = re.sub(rb"\s+", b"", value)
    result: Dict[str, bytes] = {}
    for part in compact.split(b";"):
        if b"=" in part:
            key, item = part.split(b"=", 1)
            key_text = key.decode("ascii", "strict").lower()
            if key_text in result:
                raise ValueError(f"duplicate ARC tag: {key_text}")
            result[key_text] = item
    return result


def empty_b(value: bytes) -> bytes:
    return re.sub(rb"([;\s]b=)[^;]*", rb"\1", value)


def aar_instance(value: bytes) -> int:
    """Parse the required leading ARC instance from an AAR field value."""
    unfolded = value.replace(b"\r\n", b"").replace(b"\n", b"")
    match = re.match(rb"[ \t]*i[ \t]*=[ \t]*(\d{1,2})[ \t]*;", unfolded, re.I)
    if match is None:
        raise ValueError("ARC-Authentication-Results lacks a valid leading i= instance")
    return int(match.group(1))


def load_key() -> Tuple[rsa.RSAPublicKey, str]:
    der = base64.b64decode(GOOGLE_ARC_PUBLIC_KEY, validate=True)
    key = serialization.load_der_public_key(der)
    if not isinstance(key, rsa.RSAPublicKey):
        raise TypeError("Google ARC key is not RSA")
    return key, hashlib.sha256(der).hexdigest()


def verify_signature(key: rsa.RSAPublicKey, signature: bytes, data: bytes) -> bool:
    try:
        key.verify(signature, data, padding.PKCS1v15(), hashes.SHA256())
        return True
    except InvalidSignature:
        return False


def verify(raw: bytes, label: str = "<memory>") -> ArcResult:
    header_block, body, _ = split_message(raw)
    headers = parse_headers(header_block)
    by_name: Dict[str, List[Tuple[bytes, bytes]]] = {}
    for name, value in headers:
        by_name.setdefault(name.lower().decode("ascii", "strict"), []).append((name, value))

    aar_rows = by_name.get("arc-authentication-results", [])
    ams_rows = by_name.get("arc-message-signature", [])
    seal_rows = by_name.get("arc-seal", [])
    complete = len(aar_rows) == len(ams_rows) == len(seal_rows) == 1
    if not complete:
        raise ValueError("narrow verifier requires exactly one complete ARC set")
    aar, ams, seal = aar_rows[0], ams_rows[0], seal_rows[0]
    at, st = tags(ams[1]), tags(seal[1])
    aar_i = aar_instance(aar[1])
    required_ams = {"i", "a", "c", "d", "s", "h", "bh", "b"}
    required_seal = {"i", "a", "d", "s", "t", "cv", "b"}
    if missing := sorted(required_ams - set(at)):
        raise ValueError("AMS missing tags: " + ", ".join(missing))
    if missing := sorted(required_seal - set(st)):
        raise ValueError("ARC-Seal missing tags: " + ", ".join(missing))
    if extras := sorted(set(st) - required_seal):
        raise ValueError("ARC-Seal has unsupported tags: " + ", ".join(extras))
    if at["a"].lower() != b"rsa-sha256" or st["a"].lower() != b"rsa-sha256":
        raise ValueError("only rsa-sha256 is supported")
    if at["c"].lower() != b"relaxed/relaxed":
        raise ValueError("only relaxed/relaxed AMS is supported")
    if aar_i != 1 or at["i"] != b"1" or st["i"] != b"1" or st["cv"].lower() != b"none":
        raise ValueError("narrow verifier requires i=1 and cv=none")
    if at["d"] != st["d"] or at["s"] != st["s"]:
        raise ValueError("AMS and ARC-Seal domain/selector differ")
    if at["d"].lower() != b"google.com" or at["s"].lower() != b"arc-20240605":
        raise ValueError(
            "embedded key is authorised only for google.com/arc-20240605 in this narrow verifier"
        )

    body_bytes = relaxed_body(body)
    if "l" in at:
        length_tag = at["l"]
        if re.fullmatch(rb"[0-9]+", length_tag) is None:
            raise ValueError("AMS l= tag must be an unsigned decimal integer")
        signed_length = int(length_tag)
        if signed_length > len(body_bytes):
            raise ValueError("AMS l= tag exceeds the canonicalised body length")
        body_bytes = body_bytes[:signed_length]
    body_hash = base64.b64encode(hashlib.sha256(body_bytes).digest()).decode()
    expected_body_hash = at["bh"].decode("ascii")

    signed_names = [n.strip().lower() for n in at["h"].decode("ascii").split(":") if n.strip()]
    forbidden_ams_headers = {
        "authentication-results",
        "arc-authentication-results",
        "arc-message-signature",
        "arc-seal",
    }
    if forbidden := sorted(set(signed_names) & forbidden_ams_headers):
        raise ValueError("AMS h= includes forbidden authentication/ARC headers: " + ", ".join(forbidden))
    used: Dict[str, int] = {}
    ams_input = b""
    for header_name in signed_names:
        rows = by_name.get(header_name, [])
        index = used.get(header_name, 0)
        if index < len(rows):
            selected = rows[len(rows) - 1 - index]
            ams_input += relaxed_header(*selected)
            used[header_name] = index + 1
    ams_input += relaxed_header(ams[0], empty_b(ams[1])).removesuffix(b"\r\n")

    # RFC 8617 section 5.1.1: increasing instance order; within each set,
    # AAR, AMS, then ARC-Seal.  With i=1 there is exactly one such set.
    seal_input = relaxed_header(*aar) + relaxed_header(*ams)
    seal_input += relaxed_header(seal[0], empty_b(seal[1])).removesuffix(b"\r\n")

    key, fingerprint = load_key()
    ams_ok = verify_signature(key, base64.b64decode(at["b"], validate=True), ams_input)
    seal_ok = verify_signature(key, base64.b64decode(st["b"], validate=True), seal_input)
    body_ok = body_hash == expected_body_hash
    return ArcResult(
        path=label,
        bytes=len(raw),
        sha256=hashlib.sha256(raw).hexdigest(),
        instance=1,
        domain=at["d"].decode("ascii"),
        selector=at["s"].decode("ascii"),
        arc_key_bits=key.key_size,
        arc_key_spki_sha256=fingerprint,
        ams_body_hash_expected=expected_body_hash,
        ams_body_hash_computed=body_hash,
        ams_body_hash_ok=body_ok,
        ams_signed_input_sha256=hashlib.sha256(ams_input).hexdigest(),
        ams_signature_math_ok=ams_ok,
        arc_seal_signed_input_sha256=hashlib.sha256(seal_input).hexdigest(),
        arc_seal_signature_math_ok=seal_ok,
        arc_set_complete=complete,
        overall_arc_set_ok=complete and body_ok and ams_ok and seal_ok,
        scope_note=(
            "Verified relative to the ARC public key returned at the Google-controlled selector "
            "by two distinct resolver services in the retained parsed current-DNS capture on "
            "2026-08-02; this is not raw resolver evidence, a historical DNS archive, a private-key "
            "custody finding or a live-mailbox acquisition."
        ),
    )


def mutate_body(raw: bytes) -> bytes:
    head, body, separator = split_message(raw)
    changed = bytearray(body)
    index = next(i for i, value in enumerate(changed) if value not in b"\r\n\t ")
    changed[index] ^= 1
    return head + separator + bytes(changed)


def header_span(header_block: bytes, name: bytes) -> re.Match[bytes]:
    """Return one complete, possibly folded header and refuse ambiguity."""
    matches = list(
        re.finditer(
            rb"(?im)^" + re.escape(name) + rb":[^\r\n]*(?:\r?\n[ \t][^\r\n]*)*",
            header_block,
        )
    )
    if len(matches) != 1:
        raise ValueError(f"expected exactly one header {name.decode()}; found {len(matches)}")
    return matches[0]


def mutate_header(raw: bytes, name: bytes) -> bytes:
    head, body, separator = split_message(raw)
    match = header_span(head, name)
    changed = bytearray(match.group(0))
    colon = changed.index(58)
    index = next(i for i in range(colon + 1, len(changed)) if 33 <= changed[i] <= 126)
    changed[index] = 65 if changed[index] != 65 else 66
    return head[: match.start()] + bytes(changed) + head[match.end() :] + separator + body


def mutate_header_token(raw: bytes, name: bytes, old: bytes, new: bytes) -> bytes:
    """Replace one same-length token inside one named header, including folds."""
    if len(old) != len(new):
        raise ValueError("control replacement must preserve message length")
    head, body, separator = split_message(raw)
    match = header_span(head, name)
    original = match.group(0)
    if original.count(old) != 1:
        raise ValueError(
            f"expected token {old!r} exactly once in {name.decode()}; found {original.count(old)}"
        )
    changed = original.replace(old, new, 1)
    if changed == original:
        raise AssertionError("control mutation was a no-op")
    return head[: match.start()] + changed + head[match.end() :] + separator + body


def mutate_ams_fh(raw: bytes) -> bytes:
    head, body, separator = split_message(raw)
    match = header_span(head, b"ARC-Message-Signature")
    original = match.group(0)
    token = re.search(rb"(?i)([;\s]fh=)([A-Za-z0-9+/])", original)
    if token is None:
        raise ValueError("ARC-Message-Signature has no mutable fh= value")
    changed = bytearray(original)
    index = token.start(2)
    changed[index] = 65 if changed[index] != 65 else 66
    return head[: match.start()] + bytes(changed) + head[match.end() :] + separator + body


def mutate_seal_timestamp(raw: bytes) -> bytes:
    head, body, separator = split_message(raw)
    match = header_span(head, b"ARC-Seal")
    original = match.group(0)
    token = re.search(rb"(?i)([;\s]t=)(\d+)", original)
    if token is None:
        raise ValueError("ARC-Seal has no t= value")
    old = token.group(2)
    new = str(int(old) + 1).encode("ascii")
    if len(new) != len(old):
        raise ValueError("timestamp increment changed field length")
    changed = original[: token.start(2)] + new + original[token.end(2) :]
    return head[: match.start()] + changed + head[match.end() :] + separator + body


def changed_byte_count(original: bytes, mutated: bytes) -> int:
    return sum(left != right for left, right in zip(original, mutated)) + abs(
        len(original) - len(mutated)
    )


def checked_control(
    raw: bytes,
    mutated: bytes,
    label: str,
    expected: Tuple[bool, bool, bool],
) -> Dict[str, object]:
    count = changed_byte_count(raw, mutated)
    if count == 0:
        raise AssertionError(f"{label} mutation changed zero bytes")
    result = asdict(verify(mutated, label))
    observed = (
        result["ams_body_hash_ok"],
        result["ams_signature_math_ok"],
        result["arc_seal_signature_math_ok"],
    )
    result["changed_bytes"] = count
    result["expected_components"] = {
        "ams_body_hash_ok": expected[0],
        "ams_signature_math_ok": expected[1],
        "arc_seal_signature_math_ok": expected[2],
    }
    result["expectation_ok"] = observed == expected
    if not result["expectation_ok"]:
        raise AssertionError(f"{label} produced {observed}, expected {expected}")
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("eml", type=Path)
    parser.add_argument("--negative-controls", action="store_true")
    args = parser.parse_args()
    raw = args.eml.read_bytes()
    baseline = asdict(verify(raw, str(args.eml)))
    baseline["changed_bytes"] = 0
    output = {"baseline": baseline}
    if args.negative_controls:
        body = mutate_body(raw)
        date = mutate_header(raw, b"Date")
        ams_fh = mutate_ams_fh(raw)
        aar_dkim = mutate_header_token(
            raw, b"ARC-Authentication-Results", b"dkim=pass", b"dkim=fail"
        )
        aar_spf = mutate_header_token(
            raw, b"ARC-Authentication-Results", b"spf=pass", b"spf=fail"
        )
        seal_t = mutate_seal_timestamp(raw)
        unsealed_dkim = mutate_header_token(
            raw, b"Authentication-Results", b"dkim=pass", b"dkim=fail"
        )
        output["negative_controls"] = {
            "body_mutation": checked_control(
                raw, body, "body-mutation", (False, True, True)
            ),
            "signed_date_mutation": checked_control(
                raw, date, "date-mutation", (True, False, True)
            ),
            "ams_fh_mutation": checked_control(
                raw, ams_fh, "ams-fh-mutation", (True, False, False)
            ),
            "aar_dkim_verdict_mutation": checked_control(
                raw, aar_dkim, "aar-dkim-mutation", (True, True, False)
            ),
            "aar_spf_verdict_mutation": checked_control(
                raw, aar_spf, "aar-spf-mutation", (True, True, False)
            ),
            "seal_timestamp_mutation": checked_control(
                raw, seal_t, "seal-timestamp-mutation", (True, True, False)
            ),
            "unsealed_authentication_results_mutation": checked_control(
                raw, unsealed_dkim, "unsealed-auth-results-mutation", (True, True, True)
            ),
        }
    print(json.dumps(output, indent=2, sort_keys=True))
    controls_ok = all(
        control.get("expectation_ok", False)
        for control in output.get("negative_controls", {}).values()
    )
    return 0 if output["baseline"]["overall_arc_set_ok"] and (
        not args.negative_controls or controls_ok
    ) else 1


if __name__ == "__main__":
    raise SystemExit(main())