#!/usr/bin/env python3
"""Verify an Aether grounding response and optional public receipt.

Install the two small cryptography dependencies first:

    python -m pip install blake3 cryptography

The input JSON bundle has this shape:

    {
      "answer": "optional original answer for full private verification",
      "grounding_response": { ...optional authenticated API response... },
      "share_token": "optional 43-character public capability",
      "public_receipt": { ...optional GET /receipts response... }
    }

Signer pins must come from the command line, never from the untrusted bundle:

    python verify-grounding-receipt.py bundle.json \
      --trusted-node-id "$AETHER_SIGNER_NODE_ID" \
      --share-token "$RECEIPT_CAPABILITY"

The verifier intentionally never trusts the response's `verified` booleans or
row-supplied public key by themselves. It recomputes both signatures and
requires every signer NodeId to be explicitly pinned.
"""

from __future__ import annotations

import argparse
import base64
import binascii
import json
import string
import struct
import sys
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Iterable, Optional

import blake3
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey


SOURCE_SET_DOMAIN = b"aether-grounding-source-set/v1\x00"
SOURCE_EVIDENCE_DOMAIN = b"aether-grounding-source-evidence/v1\x00"
BINDING_DOMAIN = b"aether-grounding-binding/v1\x00"
BINDING_ALGORITHM = "blake3-keyed/aether-grounding-binding/v1"
GROUNDING_ATTESTATION_VERSION = "aether-grounding-set-attestation/v1"
PUBLIC_RECEIPT_VERSION = "aether-grounding-receipt/v2"
MAX_RECEIPT_CLOCK_SKEW = timedelta(minutes=5)
HEX = set(string.hexdigits)


class VerificationError(ValueError):
    """Raised when any cryptographic or structural check fails."""


def require(condition: bool, message: str) -> None:
    if not condition:
        raise VerificationError(message)


def require_object(value: Any, name: str) -> dict[str, Any]:
    require(isinstance(value, dict), f"{name} must be a JSON object")
    return value


def require_list(value: Any, name: str) -> list[Any]:
    require(isinstance(value, list), f"{name} must be a JSON array")
    return value


def require_string(value: Any, name: str) -> str:
    require(isinstance(value, str), f"{name} must be a string")
    return value


def fixed_hex(value: Any, size: int, name: str) -> str:
    value = require_string(value, name)
    require(len(value) == size * 2, f"{name} must encode exactly {size} bytes")
    require(all(char in HEX for char in value), f"{name} must be hexadecimal")
    require(value == value.lower(), f"{name} must use canonical lowercase hex")
    return value


def decode_hex(value: Any, size: int, name: str) -> bytes:
    return bytes.fromhex(fixed_hex(value, size, name))


def decode_base64url_32(value: Any, name: str) -> bytes:
    value = require_string(value, name)
    require(len(value) == 43, f"{name} must be canonical unpadded base64url")
    alphabet = string.ascii_letters + string.digits + "-_"
    require(all(char in alphabet for char in value),
            f"{name} contains a non-base64url character")
    try:
        decoded = base64.urlsafe_b64decode(value + "=")
    except (ValueError, binascii.Error) as error:
        raise VerificationError(f"{name} is not valid base64url") from error
    require(len(decoded) == 32, f"{name} must decode to exactly 32 bytes")
    canonical = base64.urlsafe_b64encode(decoded).decode("ascii").rstrip("=")
    require(canonical == value, f"{name} is not canonically encoded")
    return decoded


def u64(value: int, name: str) -> bytes:
    require(isinstance(value, int) and not isinstance(value, bool),
            f"{name} must be an integer")
    require(0 <= value < 2**64, f"{name} is outside unsigned 64-bit range")
    return struct.pack(">Q", value)


def framed(value: bytes) -> bytes:
    return u64(len(value), "field byte length") + value


def canonical_json_bytes(fields: Iterable[tuple[str, Any]]) -> bytes:
    # Python dictionaries preserve insertion order. Rust's serde_json emits
    # the same compact UTF-8 object for these string/integer/bool/array fields.
    return json.dumps(
        dict(fields),
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")


def parse_rfc3339(value: Any, name: str) -> datetime:
    value = require_string(value, name)
    try:
        parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as error:
        raise VerificationError(f"{name} must be an RFC 3339 timestamp") from error
    require(parsed.tzinfo is not None, f"{name} must include a UTC offset")
    parsed = parsed.astimezone(timezone.utc)
    base = (
        f"{parsed.year:04d}-{parsed.month:02d}-{parsed.day:02d}"
        f"T{parsed.hour:02d}:{parsed.minute:02d}:{parsed.second:02d}"
    )
    if parsed.microsecond == 0:
        fraction = ""
    elif parsed.microsecond % 1_000 == 0:
        # Chrono's RFC 3339 AutoSi form uses milliseconds whenever the
        # PostgreSQL-microsecond value is exactly millisecond-aligned.
        fraction = f".{parsed.microsecond // 1_000:03d}"
    else:
        fraction = f".{parsed.microsecond:06d}"
    canonical = f"{base}{fraction}+00:00"
    require(value == canonical, f"{name} is not in canonical UTC RFC 3339 form")
    return parsed


def canonical_uuid(value: Any, name: str) -> str:
    value = require_string(value, name)
    try:
        parsed = uuid.UUID(value)
    except ValueError as error:
        raise VerificationError(f"{name} must be a UUID") from error
    require(str(parsed) == value, f"{name} must use canonical lowercase UUID form")
    require(
        parsed.version == 4 and parsed.variant == uuid.RFC_4122,
        f"{name} must be an RFC 4122 version-4 UUID",
    )
    return value


def trusted_node_ids(command_line: list[str]) -> set[str]:
    # Trust roots must come from the verifier/operator, never from the receipt
    # bundle. Otherwise an attacker could self-sign a bundle and include their
    # own key as its supposed allowlist.
    values = list(command_line)
    trusted = {
        fixed_hex(value, 32, "trusted signer NodeId")
        for value in values
    }
    require(bool(trusted), "at least one trusted signer NodeId must be pinned")
    return trusted


def verify_ed25519_attestation(
    attestation: dict[str, Any],
    payload: bytes,
    trusted: set[str],
    name: str,
) -> str:
    node_id = fixed_hex(attestation.get("signer_node_id"), 32, f"{name}.signer_node_id")
    public_key = decode_hex(
        attestation.get("signer_public_key"), 32, f"{name}.signer_public_key"
    )
    signature = decode_hex(attestation.get("signature"), 64, f"{name}.signature")
    derived_node_id = blake3.blake3(public_key).hexdigest()
    require(derived_node_id == node_id,
            f"{name} signer NodeId does not match BLAKE3(public key)")
    require(node_id in trusted, f"{name} signer NodeId is not pinned")
    try:
        Ed25519PublicKey.from_public_bytes(public_key).verify(signature, payload)
    except InvalidSignature as error:
        raise VerificationError(f"{name} Ed25519 signature is invalid") from error
    return node_id


def source_set_transcript(sources: list[Any]) -> bytes:
    transcript = bytearray(SOURCE_SET_DOMAIN)
    transcript.extend(u64(len(sources), "source count"))
    for expected_rank, raw_source in enumerate(sources):
        source = require_object(raw_source, f"sources[{expected_rank}]")
        rank = source.get("rank")
        require(rank == expected_rank,
                f"sources[{expected_rank}].rank must preserve declared order")
        document_id = require_string(
            source.get("document_id"), f"sources[{expected_rank}].document_id"
        ).encode("utf-8")
        content_id = require_string(
            source.get("content_id"), f"sources[{expected_rank}].content_id"
        ).encode("utf-8")
        transcript.extend(u64(rank, f"sources[{expected_rank}].rank"))
        transcript.extend(framed(document_id))
        transcript.extend(framed(content_id))
    return bytes(transcript)


def source_evidence_transcript(sources: list[Any], trusted: set[str]) -> bytes:
    transcript = bytearray(SOURCE_EVIDENCE_DOMAIN)
    transcript.extend(u64(len(sources), "source evidence count"))
    for expected_rank, raw_source in enumerate(sources):
        source = require_object(raw_source, f"sources[{expected_rank}]")
        rank = source.get("rank")
        require(rank == expected_rank,
                f"sources[{expected_rank}].rank must preserve declared order")
        document_id = require_string(
            source.get("document_id"), f"sources[{expected_rank}].document_id"
        )
        content_id = require_string(
            source.get("content_id"), f"sources[{expected_rank}].content_id"
        )
        retained_count = source.get("retained_signed_event_count")
        require(isinstance(retained_count, int) and not isinstance(retained_count, bool),
                f"sources[{expected_rank}].retained_signed_event_count must be an integer")
        require(0 <= retained_count < 2**64,
                f"sources[{expected_rank}].retained_signed_event_count is outside u64")
        verified = source.get("current_content_verified")
        require(isinstance(verified, bool),
                f"sources[{expected_rank}].current_content_verified must be boolean")

        transcript.extend(u64(rank, f"sources[{expected_rank}].rank"))
        transcript.extend(framed(document_id.encode("utf-8")))
        transcript.extend(framed(content_id.encode("utf-8")))
        transcript.extend(u64(retained_count, "retained signed event count"))
        transcript.append(int(verified))

        raw_proof = source.get("proof")
        if raw_proof is None:
            require(not verified,
                    f"sources[{expected_rank}] is verified but has no current proof")
            transcript.append(0)
            continue

        proof = require_object(raw_proof, f"sources[{expected_rank}].proof")
        require(verified, f"sources[{expected_rank}] has a proof but is not verified")
        require(proof.get("verified") is True,
                f"sources[{expected_rank}].proof.verified must be true")
        require(retained_count > 0,
                f"sources[{expected_rank}] proof requires retained evidence")
        transcript.append(1)

        proof_content_id = proof.get("content_id")
        require(proof_content_id is not None,
                f"sources[{expected_rank}] current proof must carry a content CID")
        proof_content_id = require_string(
            proof_content_id, f"sources[{expected_rank}].proof.content_id"
        )
        require(proof_content_id == content_id,
                f"sources[{expected_rank}] proof CID does not match its source CID")
        transcript.append(1)
        transcript.extend(framed(proof_content_id.encode("utf-8")))

        lamport = proof.get("lamport")
        transcript.extend(u64(lamport, f"sources[{expected_rank}].proof.lamport"))
        proof_node_id = fixed_hex(
            proof.get("node_id"), 32, f"sources[{expected_rank}].proof.node_id"
        )
        proof_public_key = decode_hex(
            proof.get("public_key"), 32, f"sources[{expected_rank}].proof.public_key"
        )
        proof_signature = fixed_hex(
            proof.get("signature"), 64, f"sources[{expected_rank}].proof.signature"
        )
        require(blake3.blake3(proof_public_key).hexdigest() == proof_node_id,
                f"sources[{expected_rank}] proof NodeId does not match its public key")
        require(proof_node_id in trusted,
                f"sources[{expected_rank}] proof signer NodeId is not pinned")
        transcript.extend(framed(proof_node_id.encode("utf-8")))
        transcript.extend(framed(proof_public_key.hex().encode("utf-8")))
        transcript.extend(framed(proof_signature.encode("utf-8")))
        transcript.append(1)
    return bytes(transcript)


def verify_grounding_response(
    answer: str,
    response: dict[str, Any],
    trusted: set[str],
) -> dict[str, Any]:
    sources = require_list(response.get("sources"), "grounding_response.sources")
    require(bool(sources), "grounding_response.sources must not be empty")
    require(len(sources) <= 50, "grounding_response.sources exceeds the 50-source limit")
    answer_hash = blake3.blake3(answer.encode("utf-8")).digest()
    expected_answer_digest = "blake3:" + answer_hash.hex()
    require(response.get("answer_digest") == expected_answer_digest,
            "answer_digest does not match the original answer bytes")

    source_hash = blake3.blake3(source_set_transcript(sources)).digest()
    evidence_hash = blake3.blake3(source_evidence_transcript(sources, trusted)).digest()
    binding = require_object(response.get("binding"), "grounding_response.binding")
    require(binding.get("algorithm") == BINDING_ALGORITHM,
            "unsupported grounding binding algorithm")
    require(binding.get("source_set_commitment") == "blake3:" + source_hash.hex(),
            "source_set_commitment does not match the ordered source set")
    require(binding.get("source_evidence_commitment") == "blake3:" + evidence_hash.hex(),
            "source_evidence_commitment does not match returned source proof fields")
    binding_key = decode_base64url_32(
        binding.get("verification_salt"), "binding.verification_salt"
    )
    binding_transcript = BINDING_DOMAIN + framed(answer_hash) + framed(source_hash)
    expected_binding = blake3.blake3(binding_transcript, key=binding_key).hexdigest()
    require(binding.get("binding_commitment") == expected_binding,
            "binding_commitment does not match answer and ordered sources")

    verified_ranks: list[int] = []
    for rank, raw_source in enumerate(sources):
        source = require_object(raw_source, f"sources[{rank}]")
        verified = source.get("current_content_verified")
        require(isinstance(verified, bool),
                f"sources[{rank}].current_content_verified must be boolean")
        if verified:
            verified_ranks.append(rank)

    trust = require_object(response.get("trust"), "grounding_response.trust")
    expected_status = "verified" if len(verified_ranks) == len(sources) else "partial"
    require(trust.get("sources_requested") == len(sources),
            "trust.sources_requested is inconsistent")
    require(trust.get("sources_verified") == len(verified_ranks),
            "trust.sources_verified is inconsistent")
    require(trust.get("status") == expected_status, "trust.status is inconsistent")
    require(trust.get("answer_bound") is True, "trust.answer_bound must be true")

    attestation = require_object(
        response.get("attestation"), "grounding_response.attestation"
    )
    require(attestation.get("version") == GROUNDING_ATTESTATION_VERSION,
            "unsupported grounding-set attestation version")
    require(attestation.get("binding_algorithm") == BINDING_ALGORITHM,
            "attestation binding algorithm does not match the binding")
    parse_rfc3339(attestation.get("issued_at"), "attestation.issued_at")
    payload = canonical_json_bytes([
        ("version", attestation["version"]),
        ("issued_at", require_string(attestation.get("issued_at"), "attestation.issued_at")),
        ("answer_digest", expected_answer_digest),
        ("binding_algorithm", BINDING_ALGORITHM),
        ("source_set_commitment", binding["source_set_commitment"]),
        ("source_evidence_commitment", binding["source_evidence_commitment"]),
        ("binding_commitment", expected_binding),
        ("answer_bound", True),
        ("source_count", len(sources)),
        ("verified_source_ranks", verified_ranks),
        ("verified_source_count", len(verified_ranks)),
        ("status", expected_status),
    ])
    signer = verify_ed25519_attestation(
        attestation, payload, trusted, "grounding_response.attestation"
    )
    return {
        "source_count": len(sources),
        "verified_source_count": len(verified_ranks),
        "status": expected_status,
        "binding_commitment": expected_binding,
        "signer_node_id": signer,
    }


def verify_public_receipt(
    receipt: dict[str, Any],
    share_token: str,
    grounding: Optional[dict[str, Any]],
    trusted: set[str],
    allow_expired: bool = False,
) -> str:
    require(receipt.get("version") == PUBLIC_RECEIPT_VERSION,
            "unsupported public receipt version")
    capability = decode_base64url_32(share_token, "share token")
    capability_commitment = "blake3:" + blake3.blake3(capability).hexdigest()
    require(receipt.get("capability_commitment") == capability_commitment,
            "public capability does not match the signed receipt")
    fixed_hex(receipt.get("binding_commitment"), 32, "receipt.binding_commitment")
    fixed_hex(receipt.get("owner_commitment"), 32, "receipt.owner_commitment")
    source_count = receipt.get("source_count")
    verified_source_count = receipt.get("verified_source_count")
    require(isinstance(source_count, int) and not isinstance(source_count, bool),
            "receipt.source_count must be an integer")
    require(0 < source_count <= 50,
            "receipt.source_count must be between 1 and 50")
    require(isinstance(verified_source_count, int)
            and not isinstance(verified_source_count, bool),
            "receipt.verified_source_count must be an integer")
    require(0 <= verified_source_count <= source_count,
            "receipt.verified_source_count is outside its source count")
    expected_status = (
        "verified" if verified_source_count == source_count else "partial"
    )
    require(receipt.get("status") == expected_status,
            "public receipt aggregate status is inconsistent")
    if grounding is not None:
        require(receipt.get("binding_commitment") == grounding["binding_commitment"],
                "public receipt is not bound to the authenticated grounding result")
        require(source_count == grounding["source_count"],
                "public receipt source count is inconsistent")
        require(verified_source_count == grounding["verified_source_count"],
                "public receipt verified source count is inconsistent")
        require(expected_status == grounding["status"],
                "public receipt status is inconsistent")

    issued_at = parse_rfc3339(receipt.get("issued_at"), "receipt.issued_at")
    expires_at = parse_rfc3339(receipt.get("expires_at"), "receipt.expires_at")
    now = datetime.now(timezone.utc)
    require(
        issued_at <= now + MAX_RECEIPT_CLOCK_SKEW,
        "public receipt issuance is too far in the future",
    )
    require(expires_at > issued_at, "public receipt expiry must follow issuance")
    require(expires_at - issued_at == timedelta(days=30),
            "public receipt must use the version-2 30-day lifetime")
    if not allow_expired:
        require(expires_at > now, "public receipt has expired")

    payload = canonical_json_bytes([
        ("version", receipt["version"]),
        ("receipt_id", canonical_uuid(receipt.get("receipt_id"), "receipt.receipt_id")),
        ("issued_at", require_string(receipt.get("issued_at"), "receipt.issued_at")),
        ("expires_at", require_string(receipt.get("expires_at"), "receipt.expires_at")),
        ("source_count", receipt["source_count"]),
        ("verified_source_count", receipt["verified_source_count"]),
        ("status", receipt["status"]),
        ("binding_commitment", receipt["binding_commitment"]),
        ("capability_commitment", capability_commitment),
        ("owner_commitment", receipt["owner_commitment"]),
    ])
    attestation = require_object(receipt.get("attestation"), "receipt.attestation")
    return verify_ed25519_attestation(attestation, payload, trusted, "receipt.attestation")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("bundle", type=Path, help="JSON verification bundle")
    parser.add_argument(
        "--trusted-node-id",
        action="append",
        default=[],
        help="trusted 64-character signer NodeId hex; repeat for a fleet",
    )
    parser.add_argument(
        "--share-token",
        help="43-character public receipt capability; overrides bundle value",
    )
    parser.add_argument(
        "--allow-expired",
        action="store_true",
        help="verify a historical test vector's signature despite expiry",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        loaded = require_object(
            json.loads(args.bundle.read_text(encoding="utf-8")), "bundle"
        )
        # A freshly fetched public-receipt JSON object is accepted directly;
        # full private verification uses the wrapper bundle documented above.
        bundle = (
            {"public_receipt": loaded}
            if loaded.get("version") == PUBLIC_RECEIPT_VERSION
            and isinstance(loaded.get("attestation"), dict)
            else loaded
        )
        trusted = trusted_node_ids(args.trusted_node_id)
        raw_answer = bundle.get("answer")
        raw_response = bundle.get("grounding_response")
        require((raw_answer is None) == (raw_response is None),
                "answer and grounding_response must be supplied together")
        grounding = None
        response = None
        if raw_response is not None:
            answer = require_string(raw_answer, "bundle.answer")
            response = require_object(raw_response, "bundle.grounding_response")
            grounding = verify_grounding_response(answer, response, trusted)

        raw_receipt = bundle.get("public_receipt")
        if raw_receipt is None and response is not None:
            raw_receipt = response.get("receipt")
        if raw_receipt is not None:
            receipt = require_object(raw_receipt, "public_receipt")
            share_token = args.share_token or bundle.get("share_token")
            require(isinstance(share_token, str),
                    "a share token is required when a public receipt is present")
            receipt_signer = verify_public_receipt(
                receipt, share_token, grounding, trusted, args.allow_expired
            )
            if grounding is None:
                print(
                    "verified public receipt capability, aggregate structure, expiry, and "
                    f"signature; signer={receipt_signer}. The answer/source binding is "
                    "opaque in public-only mode. A fresh successful GET is still required "
                    "to establish live revocation status."
                )
            else:
                print(
                    "verified grounding set "
                    f"({grounding['verified_source_count']}/{grounding['source_count']} sources) "
                    "and public receipt cryptography; "
                    f"signers={grounding['signer_node_id']},{receipt_signer}. "
                    "A fresh successful GET is still required to establish live revocation status."
                )
        else:
            require(grounding is not None,
                    "bundle must contain a grounding response or public receipt")
            print(
                "verified grounding set "
                f"({grounding['verified_source_count']}/{grounding['source_count']} sources); "
                f"signer={grounding['signer_node_id']}"
            )
        return 0
    except (
        OSError,
        UnicodeError,
        json.JSONDecodeError,
        VerificationError,
        KeyError,
    ) as error:
        print(f"verification failed: {error}", file=sys.stderr)
        return 1


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