#!/usr/bin/env python3
"""
sebbi_pq_verify.py - check a sebbi.pro Quantum Seal receipt yourself, offline.

    pip install "cryptography>=50"
    python3 sebbi_pq_verify.py receipt.json keys.json

receipt.json : https://sebbi.pro/quantum/receipt/<block>.json
keys.json    : https://sebbi.pro/.well-known/sebbi-pq-keys.json

Checks: the leaf from the block number and audit hash, the RFC 6962 Merkle
path to the batch root, the signed message rebuilt from the batch fields, and
the ML-DSA (FIPS 204) and Ed448/Ed25519 (RFC 8032) signatures against the
published keys. VALID needs both a post-quantum and a classical signature of
the same policy to verify. If the receipt carries the block record (owners
only), the block hash is recomputed too. Nothing here talks to sebbi.pro.
"""
import base64, hashlib, json, sys
from cryptography.hazmat.primitives.asymmetric import ed25519, ed448, mldsa

CONTEXT = b"sebbi-pq-batch/1"
FAMILY = {"ed25519": "classical", "ed448": "classical", "ml-dsa-44": "post-quantum", "ml-dsa-65": "post-quantum", "ml-dsa-87": "post-quantum"}
PUB = {"ed25519": ed25519.Ed25519PublicKey, "ed448": ed448.Ed448PublicKey, "ml-dsa-44": mldsa.MLDSA44PublicKey,
       "ml-dsa-65": mldsa.MLDSA65PublicKey, "ml-dsa-87": mldsa.MLDSA87PublicKey}

def leaf(block, h): return hashlib.sha256(b"\x00" + int(block).to_bytes(8, "big") + bytes.fromhex(h)).digest()
def node(a, b): return hashlib.sha256(b"\x01" + a + b).digest()

def root_from_path(m, n, r, path):
    if m >= n: return None
    fn, sn = m, n - 1
    for p in path:
        if sn == 0: return None
        if fn & 1 or fn == sn:
            r = node(p, r)
            while not (fn & 1) and fn != 0:
                fn >>= 1; sn >>= 1
        else:
            r = node(r, p)
        fn >>= 1; sn >>= 1
    return r if sn == 0 else None

def message(b):
    body = {"format": b["format"], "batch": b["batch"], "first_block": b["first_block"], "last_block": b["last_block"], "count": b["count"],
            "merkle_root": b["merkle_root"], "prev_batch_root": b["prev_batch_root"], "chain_tip_hash": b["chain_tip_hash"], "sealed_at": b["sealed_at"]}
    return CONTEXT + b"\n" + json.dumps(body, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")

def sig_ok(alg, pub, sig, msg):
    try:
        k = PUB[alg].from_public_bytes(pub)
        k.verify(sig, msg, CONTEXT) if alg.startswith("ml-dsa") else k.verify(sig, msg)
        return True
    except Exception:
        return False

def main(rp, kp):
    rc, kd = json.load(open(rp)), json.load(open(kp))
    keys = {k["kid"]: k for k in kd["keys"]}
    res = []
    def check(name, ok): res.append(ok); print(("PASS  " if ok else "FAIL  ") + name)
    rec = rc.get("record")
    if isinstance(rec, dict):
        h = hashlib.sha256(json.dumps({"prev_hash": rec["prev_hash"], "ts": rec["ts"], "event": rec["event"], "result": rec["result"]}, sort_keys=True).encode()).hexdigest()
        check("block record hashes to audit_hash", h == rc["audit_hash"])
    lf = leaf(rc["block"], rc["audit_hash"]); m = rc["merkle"]; b = rc["batch"]
    check("leaf from block number and audit hash", lf.hex() == m["leaf"])
    r = root_from_path(m["index"], m["size"], lf, [bytes.fromhex(p) for p in m["path"]])
    check("Merkle path leads to the batch root", bool(r) and r.hex() == m["root"] == b["merkle_root"])
    check("block inside the batch range", b["first_block"] <= rc["block"] <= b["last_block"])
    msg = message(b)
    pols = {}
    for s in rc["signatures"]:
        k = keys.get(s["kid"])
        ok = bool(k) and k["alg"] == s["alg"] and hashlib.sha256(base64.b64decode(k["public_key_b64"])).hexdigest() == k["public_key_sha256"] \
            and sig_ok(s["alg"], base64.b64decode(k["public_key_b64"]), base64.b64decode(s["sig_b64"]), msg)
        check("%s signature, policy %s" % (s["alg"], s["policy"]), ok)
        pols.setdefault(s["policy"], {})[FAMILY.get(s["alg"])] = ok
    check("one policy with BOTH post-quantum and classical signatures valid",
          any(v.get("post-quantum") and v.get("classical") for v in pols.values()))
    print("\nVALID" if all(res) else "\nINVALID")
    return 0 if all(res) else 1

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print(__doc__); sys.exit(2)
    sys.exit(main(sys.argv[1], sys.argv[2]))