#!/usr/bin/env python3
"""Independent verifier of pq-conformance/ (extensions/pq-signatures.md).

Written from the extension's text and the documents it cites. It shares no code with wai-rs,
which generates the corpus. ML-DSA-65 is verified by `_mldsa.py`, written here from FIPS 204;
SHA-512 is the standard library's. From the claim corpus's own second verifier
(`../claim-conformance/claim_verify.py`) it takes what pq-signatures reuses from jwp-receipts
§3.1: the JSON reading of step 1, the canonical form (JCS), the body checks of step 2, the class
rules of step 4, and through it WAI's strict Ed25519 rule (`_ed25519.py`, SPEC §7 *Signatures*)
and BLAKE3 (`_blake3.py`). From the key-discovery corpus's (`../signer-conformance/
signer_verify.py`) it takes §7's classical rule 1 checks, the `signer_id` grammar, the log
date-times and the authorities a claim needs. Everything pq-signatures adds is written below:

- the composite of `draft-ietf-lamps-pq-composite-sigs-19` §3.3 with WAI's order (§4), the
  hybrid members (§3) and the policy (§5);
- the hybrid group receipt, in its JSON and its binary form, and its objects by inclusion
  (§6.2, jwp-receipts §2.2, §2.3, §2.7, §3);
- the hybrid claim (§6.3) and the records of the `wai:<domain>\\x01 ‖ JCS` pattern (§3.4);
- the hybrid key log: entries, roots, the hybrid commitment, rules 1, 3, 5 and 6, the head walk
  under an anchor that requires hybrid, and confirmation step 4a (§6.1);
- the sizes §9 states, recomputed from the corpus's own items.

Before the corpus it checks `_mldsa.py` against FIPS 204's published vectors vendored in
`acvp/` (every ML-DSA-65 sigVer vector, and every deterministic sigGen signature as one that
must verify) and the suite against the draft's own test vector. After it, it mutates the corpus
(a flipped byte of either signature, a stripped post-quantum half) and requires every mutated
case to change its verdict, and it drops each of its own checks in turn and requires the corpus
to catch the weakened verifier.

    python3 -I pq_verify.py [pq-conformance]          verify everything above
    python3 -I pq_verify.py verdicts <vectors.json>   print `name<TAB>verdict` for every case

Exit status 0 only if every check holds. Python 3.8 or later. ML-DSA in pure Python takes a few
milliseconds a verification on a current machine (about 4 ms on an arm64 laptop), and the whole
run, with every mutation and every dropped check, about ten seconds.
"""
import copy
import hashlib
import json
import os
import struct
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.dont_write_bytecode = True  # write no cache into the corpus directory
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
sys.path.insert(0, os.path.join(HERE, "..", "signer-conformance"))
import _mldsa as ML  # noqa: E402
import claim_verify as CV  # noqa: E402
import signer_verify as SV  # noqa: E402

ED = CV.ED
blake3 = CV.blake3
EXTENSION = os.path.join(HERE, "..", "extensions", "pq-signatures.md")

# §2: the suite, draft-ietf-lamps-pq-composite-sigs-19 §2.2 and §6.
SUITE = "id-MLDSA65-Ed25519-SHA512"
LABEL = b"COMPSIG-MLDSA65-Ed25519-SHA512"
PREFIX = b"CompositeAlgorithmSignatures2025"
PQ_KEY, PQ_SIG = 1952, 3309  # FIPS 204 §4, Table 2
MEMBERS = ("pq_suite", "pq_key", "pq_sig")

D_GROUP = {b"moq-jwp:group-sign\x01": 1, b"moq-jwp:group-sign\x02": 2, b"moq-jwp:group-sign\x03": 3}
D_LINK = b"moq-jwp:receipt-link\x01"
D_LEAF2, D_NODE = b"moq-jwp:merkle-leaf\x02", b"moq-jwp:merkle-node\x01"
D_CLAIM = b"moq-jwp:claim\x01"
D_SET, D_SET_LINK = b"wai:key-set\x01", b"wai:key-set-link\x01"
D_NEXT1, D_NEXT2 = b"wai:key-set-next\x01", b"wai:key-set-next\x02"

# energy-measurement §2.7: the class tags (the corpus uses those with no evidence).
TAGS = {"HwShunt": 1, "ModelBased": 2, "Estimator": 3, "OnChipCounter": 4, "CalibratedInstrument": 5}
NAMES = {v: k for k, v in TAGS.items()}

# The checks `drop` can leave out, one at a time, to show the corpus catches the defect.
CHECKS = ("pq-verify", "ed-over-representative", "policy", "suite", "pq-lengths",
          "4a-downgrade", "4a-mismatch", "4a-anchor", "head-hybrid-only", "rule1-pq", "rule5-pq-root",
          "rule6-pq-key", "group-v3-label", "binary-v3")
DROP = set()


class Refuse(Exception):
    def __init__(self, reason, step=None):
        super().__init__(reason)
        self.reason, self.step = reason, step


def on(check):
    return check not in DROP


def hexb(v, n):
    return SV.hexb(v, n)


# ---------------------------------------------------------------- §4: the composite

def representative(m):
    """M' = Prefix ‖ Label ‖ len(ctx) ‖ ctx ‖ SHA-512(M), ctx empty (draft §2.2, §3.2)."""
    return PREFIX + LABEL + b"\x00" + hashlib.sha512(m).digest()


_ML_CACHE = {}


def ml_verify(pk, mp, sig):
    key = hashlib.sha256(pk + b"|" + mp + b"|" + sig).digest()
    if key not in _ML_CACHE:
        _ML_CACHE[key] = ML.verify(pk, mp, sig, LABEL)
    return _ML_CACHE[key]


def hybrid_check(pk, msg, sig, pq, require):
    """§4's order: the Ed25519 key (`bad_key`); under the policy, an item with no post-quantum
    half (`hybrid_required`); the Ed25519 half, over M' for a hybrid item and over M for a
    classical one (`bad_signature`); the post-quantum half's lengths (`pq_malformed`) and
    ML-DSA.Verify over M' with the label as context (`pq_signature_invalid`). None when it
    holds."""
    if not ED.strict_key(pk):
        return "bad_key"
    if pq is None:
        if require and on("policy"):
            return "hybrid_required"
        return ED.verify(pk, msg, sig)
    mp = representative(msg) if on("ed-over-representative") else msg
    why = ED.verify(pk, mp, sig)
    if why is not None:
        return why
    key, psig = pq
    if on("pq-lengths") and (len(key) != PQ_KEY or len(psig) != PQ_SIG):
        return "pq_malformed"
    if on("pq-verify") and not ml_verify(key, representative(msg), psig):
        return "pq_signature_invalid"
    return None


def read_members(obj, key_member="pq_key", sig_member="pq_sig"):
    """§3: none of the hybrid members, or all of them; `pq_suite` names §2's suite; the key and
    the signature are lowercase hex of the suite's lengths. Returns None or (key, sig)."""
    names = ["pq_suite", key_member] + ([sig_member] if sig_member else [])
    present = [n for n in names if n in obj]
    if not present:
        return None
    if len(present) != len(names):
        raise Refuse("pq_malformed")
    s = obj.get("pq_suite")
    if not isinstance(s, str):
        raise Refuse("pq_malformed")
    if s != SUITE and on("suite"):
        raise Refuse("pq_suite_unknown")
    key = hexb(obj.get(key_member), PQ_KEY)
    if key is None:
        if on("pq-lengths") or not isinstance(obj.get(key_member), str):
            raise Refuse("pq_malformed")
        key = bytes.fromhex(obj[key_member])
    sig = None
    if sig_member:
        sig = hexb(obj.get(sig_member), PQ_SIG)
        if sig is None:
            if on("pq-lengths") or not isinstance(obj.get(sig_member), str):
                raise Refuse("pq_malformed")
            sig = bytes.fromhex(obj[sig_member])
    return (key, sig)


def strip(obj, names):
    return {k: v for k, v in obj.items() if k not in names}


# ---------------------------------------------------------------- jwp-receipts §2.3, §2.7: groups

GROUP = ["kind", "track_alias", "group_id", "root_hash", "leaf_count", "joules_total", "group_start", "group_end",
         "signer_pubkey", "signer_id", "sig", "parent_receipt_hash"]
LABELLED = ["sign_version", "energy_provenance", "energy_uncertainty", "energy_counter_filtering",
            "energy_calibration_ref", "energy_coverage"]
HYBRID = ["pq_key", "pq_sig"]


def u64(v):
    return isinstance(v, int) and not isinstance(v, bool) and 0 <= v < 1 << 64


def read_class(obj, where):
    """energy-measurement §2.7's JSON fields, as (tag, evidence bytes). The corpus's classes
    carry no evidence; a refined class's is read here too."""
    p = obj.get("energy_provenance")
    if p is None:
        if any(obj.get(k) is not None for k in LABELLED[2:5]):
            raise Refuse(where)
        return 0, b""
    if p not in TAGS:
        raise Refuse(where)
    tag = TAGS[p]
    if tag < 4:
        if any(obj.get(k) is not None for k in LABELLED[2:5]):
            raise Refuse(where)
        return tag, b""
    u = obj.get("energy_uncertainty")
    if not isinstance(u, dict) or set(u) != {"relative_ppm", "window_us"}:
        raise Refuse(where)
    ev = struct.pack(">IQ", u["relative_ppm"], u["window_us"])
    if tag == 4:
        f = {"Unknown": 0, "Off": 1, "On": 2}.get(obj.get("energy_counter_filtering"))
        if f is None or obj.get("energy_calibration_ref") is not None:
            raise Refuse(where)
        return tag, ev + bytes([f])
    ref = obj.get("energy_calibration_ref")
    if not isinstance(ref, str) or obj.get("energy_counter_filtering") is not None:
        raise Refuse(where)
    return tag, ev + struct.pack(">I", len(ref.encode())) + ref.encode()


def read_group(g):
    """§3 step 1 of a group receipt's JSON members, with pq-signatures §6.2. Returns a dict."""
    if not isinstance(g, dict) or "kind" not in g:
        raise Refuse("malformed", 1)
    if g["kind"] != "group":
        raise Refuse("wrong_kind", 1)
    allowed = GROUP + LABELLED + HYBRID
    if any(k not in g for k in GROUP) or any(k not in allowed for k in g):
        raise Refuse("malformed", 1)
    ok = (u64(g["track_alias"]) and u64(g["group_id"]) and hexb(g["root_hash"], 32) is not None
          and u64(g["leaf_count"]) and u64(g["joules_total"]) and isinstance(g["group_start"], str)
          and isinstance(g["group_end"], str) and hexb(g["signer_pubkey"], 32) is not None
          and isinstance(g["signer_id"], str) and hexb(g["sig"], 64) is not None
          and (g["parent_receipt_hash"] is None or hexb(g["parent_receipt_hash"], 32) is not None))
    if not ok:
        raise Refuse("malformed", 1)
    out = dict(g)
    sv = g.get("sign_version")
    hybrid_named = any(k in g for k in HYBRID)
    version = sv if isinstance(sv, int) and not isinstance(sv, bool) else None
    if sv is None:
        if any(g.get(k) is not None for k in LABELLED):
            raise Refuse("group_label_invalid", 1)
        out["version"] = 1
    elif version == 2 or (version == 3 and (hybrid_named or not on("group-v3-label"))):
        tag, ev = read_class(g, "group_label_invalid")
        cov = g.get("energy_coverage")
        if tag == 0:
            if cov is not None:
                raise Refuse("group_label_invalid", 1)
            cov_b = 0
        else:
            if cov not in ("Whole", "Partial"):
                raise Refuse("group_label_invalid", 1)
            cov_b = 0 if cov == "Whole" else 1
        out.update(version=version, tag=tag, evidence=ev, coverage=cov_b)
    else:
        raise Refuse("group_label_invalid", 1)
    if out["version"] == 3:
        key, sig = hexb(g.get("pq_key"), PQ_KEY), hexb(g.get("pq_sig"), PQ_SIG)
        if key is None or sig is None:
            raise Refuse("pq_malformed", 1)
        out["pq"] = (key, sig)
    elif hybrid_named:
        raise Refuse("pq_malformed", 1)
    return out


def string(s):
    b = s.encode()
    return struct.pack(">I", len(b)) + b


def group_payload(g):
    """§2.3's payload in the group's form: legacy, labelled, or hybrid (pq-signatures §6.2)."""
    v = g["version"]
    dom = {1: b"moq-jwp:group-sign\x01", 2: b"moq-jwp:group-sign\x02", 3: b"moq-jwp:group-sign\x03"}[v]
    b = dom + struct.pack(">QQ", g["track_alias"], g["group_id"]) + bytes.fromhex(g["root_hash"])
    b += struct.pack(">QQ", g["leaf_count"], g["joules_total"]) + string(g["group_start"]) + string(g["group_end"])
    b += bytes.fromhex(g["signer_pubkey"]) + string(g["signer_id"])
    if v >= 2:
        b += bytes([g["tag"]]) + g["evidence"] + bytes([g["coverage"]])
        p = g["parent_receipt_hash"]
        b += (b"\x01" + bytes.fromhex(p)) if p is not None else b"\x00" + bytes(32)
    if v == 3:
        b += g["pq"][0]
    return b


def canonical(g):
    b = group_payload(g) + bytes.fromhex(g["sig"])
    return b + g["pq"][1] if g["version"] == 3 else b


def verify_group(g, require):
    """§3 steps 1-2."""
    why = hybrid_check(bytes.fromhex(g["signer_pubkey"]), group_payload(g), bytes.fromhex(g["sig"]), g.get("pq"), require)
    if why == "bad_key":
        raise Refuse("bad_signer_key", 1)
    if why == "bad_signature":
        raise Refuse("signature_verify_failed", 2)
    if why is not None:
        raise Refuse(why, 2)
    if g["version"] >= 2:
        # §2.4.1 with the group alone: the unmetered marker exactly beside a zero total, never
        # HwShunt (the corpus's classes carry no evidence to check against the floor).
        if (g["tag"] == 0) != (g["joules_total"] == 0) or g["tag"] == 1:
            raise Refuse("group_label_invalid", 2)


class Bytes:
    def __init__(self, b, step):
        self.b, self.at, self.step = b, 0, step

    def take(self, n):
        if self.at + n > len(self.b):
            raise Refuse("malformed", self.step)
        s = self.b[self.at:self.at + n]
        self.at += n
        return s

    def u8(self):
        return self.take(1)[0]

    def u32(self):
        return struct.unpack(">I", self.take(4))[0]

    def u64(self):
        return struct.unpack(">Q", self.take(8))[0]

    def string(self):
        try:
            return self.take(self.u32()).decode("utf-8")
        except UnicodeDecodeError:
            raise Refuse("malformed", self.step)

    def cls(self, m):
        tag = self.u8()
        if tag == 0:
            return
        if tag not in NAMES:
            raise Refuse("malformed", self.step)
        m["energy_provenance"] = NAMES[tag]
        if tag >= 4:
            m["energy_uncertainty"] = {"relative_ppm": self.u32(), "window_us": self.u64()}
            if tag == 4:
                f = self.u8()
                if f > 2:
                    raise Refuse("malformed", self.step)
                m["energy_counter_filtering"] = ["Unknown", "Off", "On"][f]
            else:
                m["energy_calibration_ref"] = self.string()


def read_receipts_bin(b):
    """jwp-receipts §2.7.3 and §2.7.4, with the hybrid group of pq-signatures §6.2: the
    group's JSON members, its canonical length, and each object's members."""
    c = Bytes(b, 1)
    if c.take(4) != b"WJR1":
        raise Refuse("malformed", 1)
    start = c.at
    dom = c.take(19)
    v = D_GROUP.get(dom)
    if v is None or (v == 3 and not on("binary-v3")):
        raise Refuse("malformed", 1)
    g = {"kind": "group", "track_alias": c.u64(), "group_id": c.u64(), "root_hash": c.take(32).hex(),
         "leaf_count": c.u64(), "joules_total": c.u64(), "group_start": c.string(), "group_end": c.string(),
         "signer_pubkey": c.take(32).hex(), "signer_id": c.string()}

    def parent():
        flag, h = c.u8(), c.take(32)
        if flag == 1:
            return h.hex()
        if flag == 0 and h == bytes(32):
            return None
        raise Refuse("malformed", 1)

    label = {}
    if v >= 2:
        c.cls(label)
        cov = c.u8()
        if cov > 1:
            raise Refuse("malformed", 1)
        if label:
            label["energy_coverage"] = ["Whole", "Partial"][cov]
        g["parent_receipt_hash"] = parent()
        label["sign_version"] = v
    if v == 3:
        g["pq_key"] = c.take(PQ_KEY).hex()
    g["sig"] = c.take(64).hex()
    if v == 3:
        g["pq_sig"] = c.take(PQ_SIG).hex()
    length = c.at - start
    if v == 1:
        g["parent_receipt_hash"] = parent()
    g.update(label)
    count = c.u32()
    objs = []
    c.step = 4
    for _ in range(count):
        o = {"object_id": c.u64(), "content_hash": c.take(32), "joules_micro": c.u64(), "origin": c.string()}
        n = c.u8()
        o["siblings"] = [c.take(32) for _ in range(n)]
        o["direction_bits"], o["leaf_index"] = c.u64(), c.u64()
        lv = c.u8()
        if lv not in (1, 2):
            raise Refuse("malformed", 4)
        o["leaf_version"] = lv
        if lv == 2:
            m = {}
            c.cls(m)
            o["class"] = m
        objs.append(o)
    c.step = 1
    if c.at != len(b):
        raise Refuse("malformed", 1)
    return g, b[start:start + length], objs


def ceil_log2(n):
    return (n - 1).bit_length() if n > 1 else 0


def check_object(o, g, root):
    """§3 step 5 for a labelled object: its label (§2.4.2), its position, its leaf walked
    to the verified root."""
    if o["leaf_version"] != 2:
        raise Refuse("unsupported: the corpus's objects are labelled")
    tag, ev = read_class(o["class"], "object_label_invalid")
    if (tag == 0) != (o["joules_micro"] == 0) or tag == 1:
        raise Refuse("object_label_invalid", 5)
    i = o["leaf_index"]
    if not (i < g["leaf_count"] and len(o["siblings"]) == ceil_log2(g["leaf_count"]) and o["direction_bits"] == i
            and o["object_id"] == i):
        raise Refuse("proof_position_mismatch", 5)
    node = blake3(D_LEAF2 + o["content_hash"] + struct.pack(">Q", o["joules_micro"]) + bytes([tag]) + ev + string(o["origin"]))
    for level, s in enumerate(o["siblings"]):
        node = blake3(D_NODE + (s + node if (o["direction_bits"] >> level) & 1 else node + s))
    if node != root:
        raise Refuse("inclusion_proof_failed", 5)
    return tag


# ---------------------------------------------------------------- §6.3: claims; §3.4: records

ENVELOPE = ["kind", "class", "subject_type", "subject", "seq", "issued_at", "body", "signer_pubkey", "signer_id",
            "parent_receipt_hash", "sig"]


def read_claim(text, require):
    """jwp-receipts §3.1 steps 1-4 with the hybrid members (pq-signatures §6.3)."""
    try:
        c = CV.parse_strict(text)
    except ValueError:
        raise Refuse("malformed", 1)
    try:
        if not isinstance(c, dict):
            raise CV.Reject(1, "malformed")
        if "kind" not in c:
            raise CV.Reject(1, "missing_required_field")
        CV.need(c["kind"] == "claim", 1, "wrong_kind")
        CV.members(strip(c, MEMBERS), ENVELOPE, 1)
        CV.need(isinstance(c["class"], str), 1)
        CV.need(c["class"] != "delivery", 1, "class_reserved")
        CV.need(c["class"] in CV.ADMITS, 1, "unknown_class")
        CV.need(c["subject_type"] in ("object", "receipt"), 1)
        CV.need(CV.is_hex(c["subject"], 32), 1)
        CV.need(CV.is_int(c["seq"]), 1)
        CV.need(isinstance(c["issued_at"], str) and CV.is_rfc3339(c["issued_at"]), 1)
        CV.need(CV.is_hex(c["signer_pubkey"], 32), 1)
        CV.need(isinstance(c["signer_id"], str), 1)
        CV.need(c["parent_receipt_hash"] is None or CV.is_hex(c["parent_receipt_hash"], 32), 1)
        CV.need(CV.is_hex(c["sig"], 64), 1)
        CV.need(c["subject_type"] in CV.ADMITS[c["class"]], 1, "subject_type_not_admitted")
    except CV.Reject as r:
        raise Refuse(r.reason, r.step)
    try:
        pq = read_members(c)
    except Refuse as r:
        raise Refuse(r.reason, 1)
    try:
        CV.check_body(c["class"], c["body"])
    except CV.Reject as r:
        raise Refuse(r.reason, r.step)
    msg = D_CLAIM + CV.jcs(strip(c, ("sig", "pq_sig"))).encode()
    why = hybrid_check(bytes.fromhex(c["signer_pubkey"]), msg, bytes.fromhex(c["sig"]), pq, require)
    if why is not None:
        raise Refuse({"bad_key": "bad_signer_key", "bad_signature": "signature_verify_failed"}.get(why, why), 3)
    try:
        CV.class_rules(c)
    except CV.Reject as r:
        raise Refuse(r.reason, r.step)
    return c, pq


def read_record(domain, text, require):
    """§3.4: the hybrid members, `signer_pubkey` and `sig`, then the composite over
    `domain ‖ JCS(record without sig and pq_sig)`."""
    try:
        r = CV.parse_strict(text)
    except ValueError:
        raise Refuse("malformed")
    if not isinstance(r, dict):
        raise Refuse("malformed")
    pq = read_members(r)
    pk, sig = hexb(r.get("signer_pubkey"), 32), hexb(r.get("sig"), 64)
    if pk is None or sig is None:
        raise Refuse("malformed")
    why = hybrid_check(pk, domain + CV.jcs(strip(r, ("sig", "pq_sig"))).encode(), sig, pq, require)
    if why is not None:
        raise Refuse(why)
    return pq


# ---------------------------------------------------------------- §6.1: key logs

def commitment1(k):
    return blake3(D_NEXT1 + k)


def commitment2(k, p):
    return blake3(D_NEXT2 + k + p)


def rev_hash(r):
    return blake3(D_SET_LINK + CV.jcs(r).encode()).hex()


def root_pq(r):
    return bytes.fromhex(r["root_pq_key"]) if "root_pq_key" in r else None


def root_commitment(r):
    k = bytes.fromhex(r["root"])
    p = root_pq(r)
    return commitment2(k, p) if p is not None and on("rule5-pq-root") else commitment1(k)


def check_revision(r):
    """Rule 1 with §6.1: the hybrid members of the revision and of its entries, read first,
    then §7.2's rule 1 on the rest (signer_verify.py's check), then what §6.1 adds."""
    bad = Refuse("key_log_malformed")
    if not isinstance(r, dict) or not isinstance(r.get("keys"), list):
        raise bad
    try:
        rpq = read_members(r, "root_pq_key") if on("rule1-pq") else None
        epq = [read_members(e, "pq_key", None) if isinstance(e, dict) and on("rule1-pq") else None for e in r["keys"]]
    except Refuse:
        raise bad
    plain = strip(r, ("pq_suite", "root_pq_key", "pq_sig"))
    plain["keys"] = [strip(e, ("pq_suite", "pq_key")) if isinstance(e, dict) else e for e in r["keys"]]
    try:
        SV.check_revision(plain)
    except SV.Stop:
        raise bad
    root, nxt = bytes.fromhex(r["root"]), bytes.fromhex(r["next_root"])
    if rpq is not None and on("rule1-pq") and nxt == commitment2(root, rpq[0]):
        raise bad
    keys = [e[0] for e in epq if e is not None]
    if on("rule1-pq"):
        if len(set(keys)) != len(keys) or (rpq is not None and rpq[0] in keys):
            raise bad
        for e, p in zip(r["keys"], epq):
            if p is not None and commitment2(bytes.fromhex(e["key"]), p[0]) == nxt:
                raise bad
    return r


def check_held(r, origin):
    """Rules 2-3: the origin, then the revision's signature, both halves of a hybrid one."""
    if r["origin"] != origin:
        raise Refuse("origin_mismatch")
    pq = read_members(r, "root_pq_key")
    msg = D_SET + CV.jcs(strip(r, ("sig", "pq_sig"))).encode()
    if hybrid_check(bytes.fromhex(r["root"]), msg, bytes.fromhex(r["sig"]), pq, False) is not None:
        raise Refuse("key_log_signature_invalid")


def same_root(c, p):
    return c["root"] == p["root"] and (root_pq(c) == root_pq(p) or not on("rule5-pq-root"))


def is_rotation(c, p):
    return not same_root(c, p) and root_commitment(c).hex() == p["next_root"]


def authorised(c, p):
    return same_root(c, p) or root_commitment(c).hex() == p["next_root"]


def check_link(p, c):
    """Rules 4-6, a hybrid root's ML-DSA key part of the root and an entry's never changing."""
    if c["prev"] != rev_hash(p) or c["seq"] != p["seq"] + 1 or c["issued_at"] < p["issued_at"]:
        raise Refuse("key_log_chain_break")
    if not ((same_root(c, p) and c["next_root"] == p["next_root"]) or is_rotation(c, p)):
        raise Refuse("key_log_root_not_committed")
    if len(c["keys"]) < len(p["keys"]) or c["revoked_modules"][:len(p["revoked_modules"])] != p["revoked_modules"]:
        raise Refuse("key_log_rewrites_history")
    for a, b in zip(p["keys"], c["keys"]):
        if (a["id"], a["key"], a["may_sign"]) != (b["id"], b["key"], b["may_sign"]) or SV.RANK[b["status"]] < SV.RANK[a["status"]]:
            raise Refuse("key_log_rewrites_history")
        if on("rule6-pq-key") and (a.get("pq_key"), a.get("pq_suite")) != (b.get("pq_key"), b.get("pq_suite")):
            raise Refuse("key_log_rewrites_history")
        for t in ("until", "compromised_since"):
            if a[t] is not None and a[t] != b[t]:
                raise Refuse("key_log_rewrites_history")


def read_log(origin, text):
    """§7.3's checks before anchoring, in their order."""
    try:
        d = CV.parse_strict(text)
    except ValueError:
        raise Refuse("key_log_malformed")
    if not (isinstance(d, dict) and set(d) == {"kind", "revisions"} and d["kind"] == "wai-key-log"
            and isinstance(d["revisions"], list) and d["revisions"]):
        raise Refuse("key_log_malformed")
    revs = []
    for r in d["revisions"]:
        check_revision(r)
        check_held(r, origin)
        revs.append(r)
    if revs[0]["seq"] != 0:
        raise Refuse("key_log_chain_break")
    for a, b in zip(revs, revs[1:]):
        check_link(a, b)
    return revs


def head(held, anchor, require_hybrid):
    """§7.4 step 3. Under an anchor that requires hybrid (§6.1), a revision that is not
    hybrid-signed is the child of nothing."""
    by = {rev_hash(r): r for r in held}

    def children(h):
        p = by[h]
        return [(x, c) for x, c in sorted(by.items()) if c["prev"] == h and authorised(c, p)
                and (not require_hybrid or "pq_sig" in c or not on("head-hybrid-only"))]

    line = [anchor]
    while by[line[-1]]["prev"] in by and by[line[-1]]["prev"] not in line:
        line.append(by[line[-1]]["prev"])
    line.reverse()
    start, moved = anchor, None
    for t, nxt in zip(line, line[1:]):
        if is_rotation(by[nxt], by[t]):
            continue
        rots = [(x, c) for x, c in children(t) if x != nxt and is_rotation(c, by[t])]
        if len(rots) == 1:
            check_link(by[t], rots[0][1])
            start = moved = rots[0][0]
            break
        if len(rots) > 1:
            raise Refuse("key_log_fork")
    cur = start
    for _ in range(len(by) + 1):
        kids = children(cur)
        if not kids:
            return cur, by[cur], moved
        if len(kids) > 1:
            rots = [k for k in kids if is_rotation(k[1], by[cur])]
            if len(rots) != 1:
                raise Refuse("key_log_fork")
            kids = rots
        check_link(by[cur], kids[0][1])
        cur = kids[0][0]
    raise Refuse("key_log_chain_break")


def confirm(item, pin_text, require_hybrid, held_texts, origin):
    """§7.4 with step 4a of §6.1, against a log pinned from `pin_text` (its latest revision
    the anchor, §7.3) and the revisions of `held_texts` the verifier also holds."""
    form = SV.classify(item["signer_id"])
    if form[0] != "origin":
        return "unconfirmed step=1 reason=signer_id_not_origin_bound"
    _, o, kid = form
    try:
        pinned = read_log(origin, pin_text)
    except Refuse as r:
        return "error the pinned log does not anchor: " + r.reason
    held = {rev_hash(r): r for r in pinned}
    for t in held_texts:
        for r in CV.parse_strict(t)["revisions"]:
            try:
                check_revision(r)
                check_held(r, o)
            except Refuse:
                continue
            held.setdefault(rev_hash(r), r)
    anchor = rev_hash(pinned[-1])
    if o != origin:
        return "unconfirmed step=2 reason=origin_not_anchored"
    try:
        hh, hr, moved = head(list(held.values()), anchor, require_hybrid)
    except Refuse as r:
        return "unconfirmed step=3 reason=" + r.reason
    e = next((x for x in hr["keys"] if x["id"] == kid), None)
    if e is None:
        return "refused step=4 reason=key_not_listed"
    if bytes.fromhex(e["key"]) != item["pubkey"]:
        return "refused step=4 reason=key_mismatch"
    epq = bytes.fromhex(e["pq_key"]) if "pq_key" in e else None
    if epq is not None and item["pq_key"] is None and on("4a-downgrade"):
        return "refused step=4a reason=hybrid_downgrade"
    if item["pq_key"] is not None and epq != item["pq_key"] and on("4a-mismatch"):
        return "refused step=4a reason=pq_key_mismatch"
    if require_hybrid and item["pq_key"] is None and on("4a-anchor"):
        return "unconfirmed step=4a reason=hybrid_required"
    if e["status"] == "revoked":
        return "refused step=5 reason=key_revoked"
    if e["status"] == "retired" and SV.later_than(item["time"], e["until"]):
        return "refused step=5 reason=key_retired"
    if not all(a in e["may_sign"] for a in item["authorities"]):
        return "refused step=6 reason=authority_not_granted"
    if any(a in SV.PINNED_ONLY for a in item["authorities"]) and moved is not None:
        return "unconfirmed step=6a reason=anchor_not_pinned"
    return f"confirmed key_id={kid} head={hh}"


# ---------------------------------------------------------------- the cases

def verdict(doc, c):
    """The case's verdict, as `expect` writes it; an error that is not a verdict is written
    `error ...`, which no case states."""
    try:
        return _verdict(doc, c)
    except (KeyError, ValueError, TypeError, IndexError) as e:
        return f"error {type(e).__name__}: {e}"


def _verdict(doc, c):
    require = {"accept-classical": False, "require-hybrid": True}[c["policy"]]
    fam = c["family"]
    item = None
    try:
        if fam == "group":
            if "receipts_bin" in c:
                gv, canon, objs = read_receipts_bin(bytes.fromhex(c["receipts_bin"]))
                g = read_group(gv)
                if canonical(g) != canon:
                    raise Refuse("malformed", 1)
            else:
                g, objs = read_group(c["group"]), None
            verify_group(g, require)
            head_v = "accept link=" + blake3(D_LINK + canonical(g)).hex()
            classes = [NAMES[g["tag"]]] if g["version"] >= 2 and g["tag"] else []
            if objs is not None:
                root = bytes.fromhex(g["root_hash"])
                for o in objs:
                    tag = check_object(o, g, root)
                    if tag and NAMES[tag] not in classes:
                        classes.append(NAMES[tag])
                head_v += f" objects={len(objs)}"
            item = {"pubkey": bytes.fromhex(g["signer_pubkey"]), "signer_id": g["signer_id"], "time": g["group_end"],
                    "authorities": SV.item_authorities("group", classes),
                    "pq_key": g["pq"][0] if g["version"] == 3 else None}
        elif fam == "claim":
            cl, pq = read_claim(c["claim"], require)
            head_v = "accept link=" + CV.link_hash(cl)
            item = {"pubkey": bytes.fromhex(cl["signer_pubkey"]), "signer_id": cl["signer_id"], "time": cl["issued_at"],
                    "authorities": SV.claim_authorities(cl), "pq_key": pq[0] if pq else None}
        elif fam == "record":
            pq = read_record(c["record"]["domain"].encode(), c["record"]["text"], require)
            head_v = "accept form=" + ("hybrid" if pq else "classical")
        elif fam == "key_log":
            revs = read_log(doc["origin"], doc["logs"][c["log"]])
            return "anchored head=" + rev_hash(revs[-1])
        else:
            return "error unknown family"
    except Refuse as r:
        if fam == "key_log":
            return "refused reason=" + r.reason
        return f"reject step={r.step} reason={r.reason}" if r.step is not None else f"reject reason={r.reason}"
    k = c.get("confirm")
    if k is None:
        return head_v
    v = confirm(item, doc["logs"][k["pin"]], k["require_hybrid"], [doc["logs"][h] for h in k["held"]], doc["origin"])
    return head_v + " | " + v


def sizes(doc):
    """§9's sizes, from the corpus's own items."""
    by = {c["name"]: c for c in doc["cases"]}

    def group_bin(name):
        g, canon, objs = read_receipts_bin(bytes.fromhex(by[name]["receipts_bin"]))
        return len(canon) + (33 if g.get("sign_version") is None else 0)

    def rec(name):
        return len(by[name]["record"]["text"].encode())

    def first_rev(log):
        return len(CV.jcs(CV.parse_strict(doc["logs"][log])["revisions"][0]).encode())

    def obj_len(name):
        b = bytes.fromhex(by[name]["receipts_bin"])
        g, canon, objs = read_receipts_bin(b)
        rest = len(b) - 4 - len(canon) - 4
        return rest // len(objs)

    return {
        "group_receipt_binary": {"classical": group_bin("group_classical_labelled"), "hybrid": group_bin("group_hybrid_binary")},
        "receipts_record_3_objects": {"classical": len(by["group_classical_labelled"]["receipts_bin"]) // 2,
                                      "hybrid": len(by["group_hybrid_binary"]["receipts_bin"]) // 2},
        "object_receipt_binary": obj_len("group_hybrid_binary"),
        "claim_json": {"classical": len(by["claim_classical"]["claim"].encode()), "hybrid": len(by["claim_hybrid"]["claim"].encode())},
        "first_revision_json": {"classical": first_rev("classical-r0"), "hybrid": first_rev("hybrid-r0")},
        "synthesis_record_json": {"classical": rec("record_synthesis_classical"), "hybrid": rec("record_synthesis_hybrid")},
    }


def reasons_table():
    """§8's table: the codes, in its order."""
    text = open(EXTENSION, encoding="utf-8").read()
    sec = text.split("\n## 8. ", 1)[1].split("\n## ", 1)[0]
    return [l[3:].split("`", 1)[0] for l in sec.splitlines() if l.startswith("| `")]


# ---------------------------------------------------------------- the published vectors

def check_published():
    fails = []
    acvp = os.path.join(HERE, "acvp")
    sv = json.load(open(os.path.join(acvp, "ml-dsa-65-sigver.json"), encoding="utf-8"))
    n = 0
    for g in sv["testGroups"]:
        for t in g["tests"]:
            pk, sig = bytes.fromhex(t["pk"]), bytes.fromhex(t["signature"])
            if g["signatureInterface"] == "external":
                got = ML.verify(pk, bytes.fromhex(t["message"]), sig, bytes.fromhex(t["context"]))
            elif g["externalMu"]:
                got = ML.verify_mu(pk, bytes.fromhex(t["mu"]), sig)
            else:
                got = ML.verify_internal(pk, bytes.fromhex(t["message"]), sig)
            n += 1
            if got != t["testPassed"]:
                fails.append(f"ACVP sigVer tgId {g['tgId']} tcId {t['tcId']}: {got}")
    sg = json.load(open(os.path.join(acvp, "ml-dsa-65-siggen.json"), encoding="utf-8"))
    for g in sg["testGroups"]:
        for t in g["tests"]:
            pk, sig = bytes.fromhex(t["pk"]), bytes.fromhex(t["signature"])
            if g["signatureInterface"] == "external":
                got = ML.verify(pk, bytes.fromhex(t["message"]), sig, bytes.fromhex(t["context"]))
            elif g["externalMu"]:
                got = ML.verify_mu(pk, bytes.fromhex(t["mu"]), sig)
            else:
                got = ML.verify_internal(pk, bytes.fromhex(t["message"]), sig)
            n += 1
            if not got:
                fails.append(f"ACVP sigGen tgId {g['tgId']} tcId {t['tcId']}: the stated signature does not verify")
    v = json.load(open(os.path.join(HERE, "draft-composite-sigs-19.json"), encoding="utf-8"))
    pk, m, ctx = bytes.fromhex(v["pk"]), bytes.fromhex(v["m"]), bytes.fromhex(v["ctx"])
    for field, cx in (("s", b""), ("sWithContext", ctx)):
        s = bytes.fromhex(v[field])
        mp = PREFIX + LABEL + bytes([len(cx)]) + cx + hashlib.sha512(m).digest()
        if not ML.verify(pk[:PQ_KEY], mp, s[:PQ_SIG], LABEL) or ED.verify(pk[PQ_KEY:], mp, s[PQ_SIG:]) is not None:
            fails.append(f"the draft's {field} does not verify")
        if not cx and hybrid_check(pk[PQ_KEY:], m, s[PQ_SIG:], (pk[:PQ_KEY], s[:PQ_SIG]), True) is not None:
            fails.append("the draft's vector fails this verifier's hybrid rule")
    return n, fails


# ---------------------------------------------------------------- mutations

def mutants(c):
    """Copies of an accepted case, one defect each: a flipped byte of the ML-DSA signature, of
    the Ed25519 signature, and the post-quantum half stripped."""
    out = []
    fam = c["family"]
    if fam == "group" and "group" in c and "pq_sig" in c["group"]:
        for k, at in (("pq_sig", 77), ("sig", 5)):
            m = copy.deepcopy(c)
            b = bytearray.fromhex(m["group"][k])
            b[at] ^= 1
            m["group"][k] = b.hex()
            out.append((f"{k} flipped", m))
        m = copy.deepcopy(c)
        for k in HYBRID:
            del m["group"][k]
        m["group"]["sign_version"] = 2
        out.append(("stripped", m))
    if fam == "group" and "receipts_bin" in c:
        b = bytearray.fromhex(c["receipts_bin"])
        if b[4 + 18] == 3:
            for name, at in (("pq_sig flipped", len(b) - 3 * 141 - 4 - 100), ("pq_key flipped", 300)):
                m = copy.deepcopy(c)
                bb = bytearray(b)
                bb[at] ^= 1
                m["receipts_bin"] = bb.hex()
                out.append((name, m))
    if fam in ("claim", "record"):
        key = "claim" if fam == "claim" else None
        text = c["claim"] if key else c["record"]["text"]
        v = json.loads(text)
        if "pq_sig" in v:
            for k, at in (("pq_sig", 1000), ("sig", 3)):
                w = copy.deepcopy(v)
                b = bytearray.fromhex(w[k])
                b[at] ^= 1
                w[k] = b.hex()
                out.append((f"{k} flipped", w))
            out.append(("stripped", strip(v, MEMBERS)))
            res = []
            for name, w in out:
                m = copy.deepcopy(c)
                if key:
                    m["claim"] = CV.jcs(w)
                else:
                    m["record"]["text"] = CV.jcs(w)
                res.append((name, m))
            return res
    return out


def main(argv):
    if len(argv) > 2 and argv[1] == "verdicts":
        doc = json.load(open(argv[2], encoding="utf-8"))
        for c in doc["cases"]:
            print(f"{c['name']}\t{verdict(doc, c)}")
        return 0
    corpus = argv[1] if len(argv) > 1 else HERE
    doc = json.load(open(os.path.join(corpus, "vectors.json"), encoding="utf-8"))
    fail = 0
    t0 = time.time()

    n, pf = check_published()
    for f in pf:
        print("FAIL " + f)
    fail += bool(pf)
    print(f"{'PASS' if not pf else 'FAIL'} _mldsa.py gives all {n} ACVP ML-DSA-65 vectors their verdicts, and the draft's composite vector verifies ({time.time() - t0:.1f}s)")

    s = doc["suite"]
    ok = (s["pq_suite"], s["label"].encode(), s["prefix"].encode(), s["pq_key_bytes"], s["pq_sig_bytes"]) == (SUITE, LABEL, PREFIX, PQ_KEY, PQ_SIG)
    fail += not ok
    print(f"{'PASS' if ok else 'FAIL'} the corpus names §2's suite")

    table = reasons_table()
    ok = table == doc["codes"]
    fail += not ok
    print(f"{'PASS' if ok else 'FAIL'} the corpus's codes are §8's table, in order: {table}")

    t1 = time.time()
    got = {}
    for c in doc["cases"]:
        v = verdict(doc, c)
        got[c["name"]] = v
        ok = v == c["expect"]
        fail += not ok
        if not ok:
            print(f"FAIL {c['name']}\n   expected {c['expect']}\n   got      {v}")
    print(f"{'PASS' if fail == 0 else 'FAIL'} {len(doc['cases'])} cases, each its stated verdict ({time.time() - t1:.1f}s)")

    every = " ".join(got.values())
    missing = [r for r in table if f"reason={r}" not in every]
    fail += bool(missing)
    print(f"{'PASS' if not missing else 'FAIL'} the corpus gives every code of §8 {missing or ''}")

    ok = sizes(doc) == doc["sizes"]
    fail += not ok
    print(f"{'PASS' if ok else 'FAIL'} the sizes the corpus states are its items' own")

    nm = bad = 0
    for c in doc["cases"]:
        if not c["expect"].startswith("accept"):
            continue
        for name, m in mutants(c):
            nm += 1
            v = verdict(doc, m)
            if v.split(" | ")[0] == c["expect"].split(" | ")[0]:
                bad += 1
                print(f"FAIL mutant {c['name']} ({name}) still {v}")
    fail += bad > 0 or nm == 0
    print(f"{'PASS' if bad == 0 and nm else 'FAIL'} {nm} mutated cases, each refused")

    weak = []
    for chk in CHECKS:
        DROP.clear()
        DROP.add(chk)
        caught = any(verdict(doc, c) != c["expect"] for c in doc["cases"])
        if not caught:
            weak.append(chk)
    DROP.clear()
    fail += bool(weak)
    print(f"{'PASS' if not weak else 'FAIL'} the corpus catches the verifier with any one of its {len(CHECKS)} checks dropped {weak or ''}")

    print(f"{'all checks pass' if fail == 0 else f'{fail} failure(s)'} in {time.time() - t0:.1f}s")
    return 1 if fail else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))