#!/usr/bin/env python3
"""Verify an Analog Seed signed entropy receipt (AS-RECEIPT-V1).

The signature is an Ed25519 origin attestation over physical entropy.
It is not a VRF: a VRF would recompute the output from the key and an input.

Depends only on the Python standard library and ``cryptography``.

  python tools/verify_receipt.py receipt.json --pubkey-hex <64 hex chars>
  python tools/verify_receipt.py --pubkey-file pubkey.json < receipt.json
  python tools/verify_receipt.py receipt.json --pubkey-hex <hex> --expect-nonce "ticket 1"

Exit status is 0 when the receipt is valid and 1 otherwise.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import re
import sys
from pathlib import Path
from typing import Any

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

PREFIX = b"AS-RECEIPT-V1\n"
VERSION = "AS-RECEIPT-V1"
_DRAW_ENDPOINTS = {"/v1/draw", "/v1/sandbox/draw"}
_HEX_RE = re.compile(r"^[0-9a-fA-F]+$")


class VerifyError(Exception):
    """The receipt failed one of the AS-RECEIPT-V1 checks."""


def canonical_json(obj: Any) -> bytes:
    """Match the server: sort keys, no whitespace, reject NaN."""
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8")


def verify_document(
    document: dict[str, Any],
    public_key: bytes,
    *,
    expect_nonce: str | None = None,
) -> None:
    """Raise ``VerifyError`` when any AS-RECEIPT-V1 check fails."""
    if len(public_key) != 32:
        raise VerifyError("public key must be 32 bytes")
    receipt = document.get("receipt")
    if not isinstance(receipt, dict):
        raise VerifyError("response has no receipt object")
    if receipt.get("version") != VERSION:
        raise VerifyError("version is not AS-RECEIPT-V1")
    payload = receipt.get("payload")
    if not isinstance(payload, dict):
        raise VerifyError("receipt payload is missing")
    if payload.get("source") != "avalanche":
        raise VerifyError("payload.source is not avalanche")
    consumed = payload.get("bytes_consumed")
    if isinstance(consumed, bool) or not isinstance(consumed, int) or consumed < 1:
        raise VerifyError("payload.bytes_consumed must be a positive integer")
    signature_hex = receipt.get("signature_hex")
    if not isinstance(signature_hex, str):
        raise VerifyError("signature_hex is missing")
    try:
        signature = bytes.fromhex(signature_hex)
    except ValueError as exc:
        raise VerifyError("signature_hex is not hex") from exc
    if len(signature) != 64:
        raise VerifyError("Ed25519 signature must be 64 bytes")

    key_id = hashlib.sha256(public_key).hexdigest()[:16]
    claimed = payload.get("key_id")
    if not isinstance(claimed, str) or not hmac.compare_digest(claimed, key_id):
        raise VerifyError("key_id does not match sha256(public key)[:16]")

    message = PREFIX + canonical_json(payload)
    try:
        Ed25519PublicKey.from_public_bytes(public_key).verify(signature, message)
    except Exception as exc:
        raise VerifyError("Ed25519 signature does not verify") from exc

    if expect_nonce is not None and payload.get("nonce") != expect_nonce:
        raise VerifyError("nonce does not match --expect-nonce")

    endpoint = payload.get("endpoint")
    digest = payload.get("output_sha256")
    if not isinstance(digest, str) or len(digest) != 64:
        raise VerifyError("output_sha256 is missing")
    if endpoint == "/v1/random":
        encoded = document.get("bytes")
        if not isinstance(encoded, str):
            raise VerifyError("/v1/random response has no bytes field")
        raw = _decode_matching(encoded, digest)
        if raw is None:
            raise VerifyError("output_sha256 does not match the served bytes")
    elif endpoint in _DRAW_ENDPOINTS:
        if "data" not in document:
            raise VerifyError("draw response has no data field")
        got = hashlib.sha256(canonical_json(document["data"])).hexdigest()
        if not hmac.compare_digest(got, digest):
            raise VerifyError("output_sha256 does not match canonical JSON of data")
    else:
        raise VerifyError("endpoint is not /v1/random, /v1/draw, or /v1/sandbox/draw")


def _decode_matching(encoded: str, digest: str) -> bytes | None:
    if _HEX_RE.fullmatch(encoded) and len(encoded) % 2 == 0:
        raw = bytes.fromhex(encoded)
        if hmac.compare_digest(hashlib.sha256(raw).hexdigest(), digest):
            return raw
    try:
        raw = base64.b64decode(encoded, validate=True)
    except Exception:
        return None
    if hmac.compare_digest(hashlib.sha256(raw).hexdigest(), digest):
        return raw
    return None


def load_public_key(hex_key: str | None, path: str | None) -> bytes:
    if (hex_key is None) == (path is None):
        raise VerifyError("pass exactly one of --pubkey-hex or --pubkey-file")
    if hex_key is None:
        document = json.loads(Path(path or "").read_text(encoding="utf-8-sig"))
        if not isinstance(document, dict) or not isinstance(document.get("public_key_hex"), str):
            raise VerifyError("pubkey file has no public_key_hex")
        hex_key = document["public_key_hex"]
        claimed = document.get("key_id")
    else:
        claimed = None
    text = hex_key.strip().lower()
    if len(text) != 64 or _HEX_RE.fullmatch(text) is None:
        raise VerifyError("public key hex must be 64 hex characters")
    public_key = bytes.fromhex(text)
    if claimed is not None:
        key_id = hashlib.sha256(public_key).hexdigest()[:16]
        if not hmac.compare_digest(str(claimed), key_id):
            raise VerifyError("pubkey file key_id does not match public_key_hex")
    return public_key


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Verify an AS-RECEIPT-V1 signed entropy receipt.")
    parser.add_argument("file", nargs="?", help="API response JSON. Reads stdin when omitted.")
    parser.add_argument("--pubkey-hex", default=None, help="Ed25519 public key, 32 bytes hex.")
    parser.add_argument("--pubkey-file", default=None, help="JSON document from GET /v1/pubkey.")
    parser.add_argument("--expect-nonce", default=None, help="Require payload.nonce to equal this string.")
    args = parser.parse_args(argv)
    try:
        public_key = load_public_key(args.pubkey_hex, args.pubkey_file)
        if args.file:
            raw = Path(args.file).read_text(encoding="utf-8-sig")
        else:
            raw = sys.stdin.read().lstrip("\ufeff")
        document = json.loads(raw)
        if not isinstance(document, dict):
            raise VerifyError("response JSON must be an object")
        verify_document(document, public_key, expect_nonce=args.expect_nonce)
    except VerifyError as exc:
        print(f"invalid: {exc}", file=sys.stderr)
        return 1
    except (OSError, json.JSONDecodeError) as exc:
        print(f"invalid: {exc}", file=sys.stderr)
        return 1
    print("valid")
    return 0


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