open-standards/wai/jwp-conformance/jwp_verify.py

#!/usr/bin/env python3
"""Second verifier for jwp-conformance/ — labelled leaves, labelled group receipts, the
group total and the WAI_RECEIPT property (extensions/jwp-receipts.md §2.2-§2.4.2, §3,
§3.3 and §5.7; moq-streaming-format.md §5).

Written from the extension text. It shares no code with wai-rs (`crate::jwp_delivery`),
which generates the corpus. It takes from ../claim-conformance/claim_verify.py — itself
independent of wai-rs — BLAKE3 written from its specification, jwp-receipts §3.1 step 3's
Ed25519 rule (key decoding, small order, S < L, the cofactorless equation through the
`cryptography` package), and the claim verifier an energy claim is checked with.

    python3 jwp_verify.py [vectors.json]   verify every vector against its declared outcome
    python3 jwp_verify.py verdicts <file>  print `name<TAB>verdict` for every vector,
                                           ignoring any declared outcome

Exit status 0 only if every vector reports its declared outcome.
"""
import json, os, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import claim_verify as CV  # noqa: E402

LEAF1 = b"moq-jwp:merkle-leaf\x01"
LEAF2 = b"moq-jwp:merkle-leaf\x02"
NODE = b"moq-jwp:merkle-node\x01"
SIGN1 = b"moq-jwp:group-sign\x01"
SIGN2 = b"moq-jwp:group-sign\x02"
LINK = b"moq-jwp:receipt-link\x01"
OBJECT = b"moq-jwp:object\x01"
FLOOR = 288676  # energy-measurement §2.3: joules_micro x relative_ppm, at least
U32, U64 = 2 ** 32 - 1, 2 ** 64 - 1
RANK = {"Estimator": 1, "ModelBased": 2, "HwShunt": 3, "OnChipCounter": 4, "CalibratedInstrument": 5}
TAG = {"HwShunt": 1, "ModelBased": 2, "Estimator": 3, "OnChipCounter": 4, "CalibratedInstrument": 5}
FILTERING = {"Unknown": 0, "Off": 1, "On": 2}
CLASS_KEYS = ["energy_provenance", "energy_uncertainty", "energy_counter_filtering", "energy_calibration_ref"]
GROUP = ["kind", "track_alias", "group_id", "root_hash", "leaf_count", "joules_total", "group_start", "group_end",
         "signer_pubkey", "signer_id", "sig", "parent_receipt_hash"]
GROUP_LABEL = ["sign_version"] + CLASS_KEYS + ["energy_coverage"]
OBJ = ["kind", "track_alias", "group_id", "object_id", "content_hash", "joules_micro", "origin", "merkle_proof"]
OBJ_LABEL = ["leaf_version"] + CLASS_KEYS


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

    def label(self):
        return f"reject reason={self.reason}" if self.step is None else f"reject step={self.step} reason={self.reason}"


class Unreadable(Exception):
    """A class whose JSON members do not make one class (energy-measurement §2.7)."""


def blake3(b):
    return CV.blake3(b)


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


def hex_bytes(v, n=None):
    if not CV.is_hex(v, n):
        return None
    return bytes.fromhex(v)


def present(m, k):
    return m.get(k) is not None


# ------------------------------------------------------------ classes (§2.7)

def read_class(m):
    """The class a receipt's JSON members name, as energy-measurement §2.7 reads it, or
    None for no class. Evidence without a class, an unknown class, or evidence a class
    does not carry make no class: Unreadable."""
    label = m.get("energy_provenance")
    u, f, r = m.get("energy_uncertainty"), m.get("energy_counter_filtering"), m.get("energy_calibration_ref")
    if label is None:
        if u is not None or f is not None or r is not None:
            raise Unreadable()
        return None
    if not isinstance(label, str) or label not in TAG:
        raise Unreadable()
    if label in ("HwShunt", "ModelBased", "Estimator"):
        if u is not None or f is not None or r is not None:
            raise Unreadable()
        return {"label": label}
    if not isinstance(u, dict) or not u64(u.get("relative_ppm")) or not u64(u.get("window_us")):
        raise Unreadable()
    ppm, win = u["relative_ppm"], u["window_us"]
    if ppm == 0 or ppm > U32 or win == 0:
        raise Unreadable()
    if label == "OnChipCounter":
        if r is not None:
            raise Unreadable()
        if f is None:
            f = "Unknown"  # an unstated filtering state is Unknown, never Off (§2.4)
        if f not in FILTERING:
            raise Unreadable()
        return {"label": label, "ppm": ppm, "window": win, "filtering": f}
    if f is not None or not isinstance(r, str) or r.strip(CV.WHITE_SPACE) == "":
        raise Unreadable()
    return {"label": label, "ppm": ppm, "window": win, "ref": r}


def encode(c):
    """The signed encoding (§2.7); None is tag 0, the unmetered marker."""
    if c is None:
        return b"\x00"
    b = bytes([TAG[c["label"]]])
    if c["label"] == "OnChipCounter":
        b += c["ppm"].to_bytes(4, "big") + c["window"].to_bytes(8, "big") + bytes([FILTERING[c["filtering"]]])
    elif c["label"] == "CalibratedInstrument":
        ref = c["ref"].encode()
        b += c["ppm"].to_bytes(4, "big") + c["window"].to_bytes(8, "big") + len(ref).to_bytes(4, "big") + ref
    return b


def label_ok(figure, c):
    """A class as an optional label (energy-measurement §2.8): never HwShunt, never
    beside a zero figure, at the resolution floor where it declares an uncertainty."""
    if c["label"] == "HwShunt" or figure == 0:
        return False
    return "ppm" not in c or figure * c["ppm"] >= FLOOR


# ------------------------------------------------------------ the tree (§2.2)

def node(a, b):
    return blake3(NODE + a + b)


def leaf(o):
    if not o["labelled"]:
        return blake3(LEAF1 + o["content_hash"])
    origin = o["origin"].encode()
    return blake3(LEAF2 + o["content_hash"] + o["joules_micro"].to_bytes(8, "big") + encode(o["class"])
                  + len(origin).to_bytes(4, "big") + origin)


def depth(n):
    return 0 if n <= 1 else (n - 1).bit_length()


# ------------------------------------------------------------ group receipts (§2.4.1, §3 steps 1-3)

def read_group(g):
    """§3 step 1's reading of a group receipt."""
    if not isinstance(g, dict) or "kind" not in g:
        raise Refuse(1, "malformed")
    if g["kind"] != "group":
        raise Refuse(1, "wrong_kind")
    if any(k not in g for k in GROUP) or any(k not in GROUP and k not in GROUP_LABEL for k in g):
        raise Refuse(1, "malformed")
    ok = all(u64(g[k]) for k in ("track_alias", "group_id", "leaf_count", "joules_total")) \
        and all(isinstance(g[k], str) for k in ("group_start", "group_end", "signer_id")) \
        and CV.is_hex(g["root_hash"], 32) and CV.is_hex(g["signer_pubkey"], 32) and CV.is_hex(g["sig"], 64) \
        and (g["parent_receipt_hash"] is None or CV.is_hex(g["parent_receipt_hash"], 32))
    if not ok:
        raise Refuse(1, "malformed")
    out = dict(g, labelled=False)
    sv = g.get("sign_version")
    if sv is None:
        if any(present(g, k) for k in GROUP_LABEL):
            raise Refuse(1, "group_label_invalid")
        return out
    # sign_version is a labelled member (§2.4.1): any value but 2 is not the labelled
    # form, and is group_label_invalid (§3.3), as leaf_version is object_label_invalid.
    if not (isinstance(sv, int) and not isinstance(sv, bool) and sv == 2):
        raise Refuse(1, "group_label_invalid")
    try:
        c = read_class(g)
    except Unreadable:
        raise Refuse(1, "group_label_invalid")
    cov = g.get("energy_coverage")
    if c is None:
        if cov is not None:
            raise Refuse(1, "group_label_invalid")
        cov = "Whole"
    elif cov not in ("Whole", "Partial"):
        raise Refuse(1, "group_label_invalid")
    out.update(labelled=True, cls=c, coverage=cov)
    return out


def payload(g):
    """§2.3: the legacy payload, or the labelled one — the same fields under the 0x02
    domain, then the class, the coverage byte and the parent link."""
    b = (SIGN2 if g["labelled"] else SIGN1) + g["track_alias"].to_bytes(8, "big") + g["group_id"].to_bytes(8, "big")
    b += bytes.fromhex(g["root_hash"]) + g["leaf_count"].to_bytes(8, "big") + g["joules_total"].to_bytes(8, "big")
    for s in (g["group_start"], g["group_end"]):
        e = s.encode()
        b += len(e).to_bytes(4, "big") + e
    sid = g["signer_id"].encode()
    b += bytes.fromhex(g["signer_pubkey"]) + len(sid).to_bytes(4, "big") + sid
    if g["labelled"]:
        b += encode(g["cls"]) + (b"\x01" if g["cls"] is not None and g["coverage"] == "Partial" else b"\x00")
        parent = g["parent_receipt_hash"]
        b += (b"\x01" + bytes.fromhex(parent)) if parent is not None else b"\x00" + bytes(32)
    return b


def link(g):
    return blake3(LINK + payload(g) + bytes.fromhex(g["sig"]))


def verify_group(g):
    """§3 steps 1-2: the key, the signature under §3.1 step 3's rule, then a labelled
    group's own checks."""
    try:
        CV.check_signature(bytes.fromhex(g["signer_pubkey"]), bytes.fromhex(g["sig"]), payload(g))
    except CV.Reject as r:
        raise Refuse(1 if r.reason == "bad_signer_key" else 2, r.reason)
    if g["labelled"]:
        c, total = g["cls"], g["joules_total"]
        if c is None:
            ok = total == 0 and g["coverage"] == "Whole"
        else:
            ok = label_ok(total, c)
        if not ok:
            raise Refuse(2, "group_label_invalid")


def check_parent(g, previous):
    """§3 step 3. A previous group that does not verify is not held."""
    if previous is not None:
        try:
            verify_group(previous)
        except Refuse:
            previous = None
    parent = g["parent_receipt_hash"]
    if previous is None:
        if parent is not None:
            raise Refuse(3, "parent_not_held")
        return
    if parent is None or previous["signer_pubkey"] != g["signer_pubkey"] or bytes.fromhex(parent) != link(previous):
        raise Refuse(3, "chain_break")


# ------------------------------------------------------------ object receipts (§2.4.2, §3 steps 4-5)

def read_object(o):
    """The object receipt's members. Shape failures are `malformed` at step 4; a label
    that cannot be read is kept, and refused at step 5."""
    if not isinstance(o, dict) or "kind" not in o:
        raise Refuse(4, "malformed")
    if o["kind"] != "object":
        raise Refuse(4, "wrong_kind")
    if any(k not in o for k in OBJ) or any(k not in OBJ and k not in OBJ_LABEL for k in o):
        raise Refuse(4, "malformed")
    p = o["merkle_proof"]
    ok = all(u64(o[k]) for k in ("track_alias", "group_id", "object_id", "joules_micro")) \
        and CV.is_hex(o["content_hash"], 32) and isinstance(o["origin"], str) \
        and isinstance(p, dict) and sorted(p) == ["direction_bits", "leaf_index", "siblings"] \
        and u64(p["direction_bits"]) and u64(p["leaf_index"]) and isinstance(p["siblings"], list) \
        and all(CV.is_hex(s, 32) for s in p["siblings"])
    if not ok:
        raise Refuse(4, "malformed")
    r = {"track_alias": o["track_alias"], "group_id": o["group_id"], "object_id": o["object_id"],
         "content_hash": bytes.fromhex(o["content_hash"]), "joules_micro": o["joules_micro"], "origin": o["origin"],
         "siblings": [bytes.fromhex(s) for s in p["siblings"]], "bits": p["direction_bits"], "index": p["leaf_index"],
         "labelled": False, "class": None, "unreadable": False}
    lv = o.get("leaf_version")
    if lv is None:
        r["unreadable"] = any(present(o, k) for k in OBJ_LABEL)
    elif isinstance(lv, int) and not isinstance(lv, bool) and lv == 2:
        try:
            r.update(labelled=True, **{"class": read_class(o)})
        except Unreadable:
            r["unreadable"] = True
    else:
        r["unreadable"] = True
    return r


def inclusion(o, g):
    """§3 step 5: the label, the position, the walk."""
    if o["unreadable"]:
        raise Refuse(5, "object_label_invalid")
    if o["labelled"]:
        c = o["class"]
        if (c is None and o["joules_micro"] != 0) or (c is not None and not label_ok(o["joules_micro"], c)):
            raise Refuse(5, "object_label_invalid")
    n = g["leaf_count"]
    if (o["track_alias"] != g["track_alias"] or o["group_id"] != g["group_id"] or o["object_id"] != o["index"]
            or not o["index"] < n or len(o["siblings"]) != depth(n) or o["bits"] != o["index"]):
        raise Refuse(5, "proof_position_mismatch")
    cur = leaf(o)
    for i, s in enumerate(o["siblings"]):
        cur = node(s, cur) if (o["bits"] >> i) & 1 else node(cur, s)
    if cur != bytes.fromhex(g["root_hash"]):
        raise Refuse(5, "inclusion_proof_failed")


def accept(o, g, envelope):
    """§3 steps 1-5."""
    verify_group(g)
    if blake3(OBJECT + envelope) != o["content_hash"]:
        raise Refuse(4, "content_hash_mismatch")
    inclusion(o, g)


# ------------------------------------------------------------ the group total (§3)

def compose(parts):
    """The composition rule over (figure, labelled, class) parts. None: no class may be
    carried; otherwise (class or None for the unmetered marker, coverage)."""
    any_zero = any(f == 0 for f, _, _ in parts)
    live = [(f, c) for f, lab, c in parts if f != 0]
    if any(c is None for _, c in live):
        return None
    if not live:
        return (None, "Whole")
    cov = "Partial" if any_zero else "Whole"
    weakest = min(RANK[c["label"]] for _, c in live)
    if weakest <= 2:
        return ({"label": "ModelBased" if weakest == 2 else "Estimator"}, cov)
    if any("ppm" not in c for _, c in live):
        return None
    num = sum(f * c["ppm"] for f, c in live)
    den = sum(f for f, _ in live)
    ppm = -(-num // den)
    win = max(c["window"] for _, c in live)
    if weakest == 4:
        fs = {c["filtering"] for _, c in live if c["label"] == "OnChipCounter"}
        return ({"label": "OnChipCounter", "ppm": ppm, "window": win, "filtering": fs.pop() if len(fs) == 1 else "Unknown"}, cov)
    refs = {c["ref"] for _, c in live}
    if len(refs) != 1:
        return None
    return ({"label": "CalibratedInstrument", "ppm": ppm, "window": win, "ref": refs.pop()}, cov)


def same_but_uncertainty(a, b):
    if a["label"] != b["label"]:
        return False
    if a["label"] == "OnChipCounter":
        return a["filtering"] == b["filtering"]
    if a["label"] == "CalibratedInstrument":
        return a["ref"] == b["ref"]
    return True


def signed_label(c, cov, checked):
    u = f" relative_ppm={c['ppm']} window_us={c['window']}" if "ppm" in c else ""
    return f"read signed class={c['label']}{u} coverage={cov} parts={'checked' if checked else 'unchecked'}"


def total(inp):
    g = read_group(inp["group"])
    verify_group(g)
    held = {}
    for d in inp["delivered"]:
        try:
            o = read_object(d["object"])
            accept(o, g, bytes.fromhex(d["envelope"]))
        except Refuse:
            continue
        held.setdefault(o["index"], []).append(o)
    complete = g["leaf_count"] > 0 and len(held) == g["leaf_count"] and all(len(v) == 1 for v in held.values())

    def own(checked):
        if not g["labelled"]:
            return "read unlabelled"
        if g["cls"] is None:
            return "read unmetered"
        return signed_label(g["cls"], g["coverage"], checked)
    if not complete:
        return own(False)
    parts = [held[i][0] for i in sorted(held)]
    if sum(o["joules_micro"] for o in parts) != g["joules_total"]:
        raise Refuse(None, "group_total_mismatch")
    comp = compose([(o["joules_micro"], o["labelled"], o["class"]) for o in parts])
    if not g["labelled"]:
        if all(o["joules_micro"] == 0 for o in parts):
            return "read unmetered"
        if comp is None or comp[0] is None:
            return "read unlabelled"
        return f"read weakest class={comp[0]['label']} coverage={comp[1]}"
    if comp is None:
        raise Refuse(None, "group_class_disagrees")
    got, want = g["cls"], comp[0]
    if (got is None) != (want is None) or (got is not None and not same_but_uncertainty(got, want)):
        raise Refuse(None, "group_class_disagrees")
    if g["coverage"] != comp[1]:
        raise Refuse(None, "group_coverage_disagrees")
    if got is not None and "ppm" in got and (got["ppm"] < want["ppm"] or got["window"] != want["window"]):
        raise Refuse(None, "group_uncertainty_below_parts")
    return own(True)


# ------------------------------------------------------------ the WAI_RECEIPT property

def property_class(p):
    """moq-streaming-format §5's rules on the property alone (`property_invalid`); the
    class it names, or None for a zero figure."""
    bad = Refuse(None, "property_invalid")
    acq, ev = p["acquisition"], bytes.fromhex(p["acquisition_evidence"])
    if p["joules_micro"] == 0:
        if acq != "" or ev:
            raise bad
        return None
    if acq in ("HwShunt", "ModelBased", "Estimator"):
        if ev:
            raise bad
        return {"label": acq}
    if acq == "OnChipCounter":
        if len(ev) != 13 or ev[12] > 2:
            raise bad
        c = {"label": acq, "ppm": int.from_bytes(ev[0:4], "big"), "window": int.from_bytes(ev[4:12], "big"),
             "filtering": {0: "Unknown", 1: "Off", 2: "On"}[ev[12]]}
    elif acq == "CalibratedInstrument":
        if len(ev) < 16 or int.from_bytes(ev[12:16], "big") != len(ev) - 16:
            raise bad
        try:
            ref = ev[16:].decode()
        except UnicodeDecodeError:
            raise bad
        if ref.strip(CV.WHITE_SPACE) == "":
            raise bad
        c = {"label": acq, "ppm": int.from_bytes(ev[0:4], "big"), "window": int.from_bytes(ev[4:12], "big"), "ref": ref}
    else:
        raise bad
    if c["ppm"] == 0 or c["window"] == 0 or p["joules_micro"] * c["ppm"] < FLOOR:
        raise bad
    return c


def claim_class(b):
    """An energy claim's class in this file's form."""
    label = b["energy_provenance"]
    if label in ("OnChipCounter", "CalibratedInstrument"):
        u = b["energy_uncertainty"]
        c = {"label": label, "ppm": u["relative_ppm"], "window": u["window_us"]}
        if label == "OnChipCounter":
            c["filtering"] = b["energy_counter_filtering"]
        else:
            c["ref"] = b["energy_calibration_ref"]
        return c
    return {"label": label}


def prop(inp):
    p = inp["property"]
    c = property_class(p)
    if "claim" in inp:
        try:
            claim = CV.verify_claim(inp["claim"])
        except CV.Reject:
            return "read unsigned" if c is not None else "read unmetered"
        if claim["class"] != "energy" or claim["subject"] != p["content_hash"]:
            return "read unsigned" if c is not None else "read unmetered"
        b = claim["body"]
        if c is not None and claim_class(b) == c and b["joules_micro"] == p["joules_micro"]:
            # Signed by the claim's signer, whoever that is: never read as the publisher's.
            return f"read signed class={c['label']} by={claim['signer_pubkey']}"
        raise Refuse(None, "property_disagrees")
    g = read_group(inp["group"])
    o = read_object(inp["object"])
    if bytes.fromhex(p["content_hash"]) != o["content_hash"] or bytes.fromhex(p["group_receipt"]) != link(g):
        raise Refuse(None, "property_disagrees")
    accept(o, g, bytes.fromhex(inp["envelope"]))
    if not o["labelled"]:
        return "read unmetered" if c is None else "read unsigned"
    if o["class"] is None and c is None:
        return "read unmetered"
    if o["class"] is not None and c is not None and o["class"] == c and o["joules_micro"] == p["joules_micro"]:
        return f"read signed class={c['label']} by={g['signer_pubkey']}"
    raise Refuse(None, "property_disagrees")


# ------------------------------------------------------------ vectors

def verdict(v):
    inp = v["input"]
    try:
        if v["procedure"] == "leaf":
            return "leaf=" + leaf(read_object(inp["object"])).hex()
        if v["procedure"] == "group":
            g = read_group(inp["group"])
            verify_group(g)
            prev = None
            if "previous" in inp:
                try:
                    prev = read_group(inp["previous"])
                except Refuse:
                    prev = None
            check_parent(g, prev)
            return f"accept link_hash={link(g).hex()} signing_payload={payload(g).hex()}"
        if v["procedure"] == "object":
            g = read_group(inp["group"])
            o = read_object(inp["object"])
            accept(o, g, bytes.fromhex(inp["envelope"]))
            return "accept"
        if v["procedure"] == "total":
            return total(inp)
        if v["procedure"] == "property":
            return prop(inp)
    except Refuse as r:
        return r.label()
    raise ValueError("unknown procedure " + str(v["procedure"]))


def expected_of(e):
    if "leaf" in e:
        return "leaf=" + e["leaf"]
    if e["outcome"] == "accept":
        return f"accept link_hash={e['link_hash']} signing_payload={e['signing_payload']}" if "link_hash" in e else "accept"
    if e["outcome"] == "read":
        return "read " + e["reading"]
    return f"reject step={e['step']} reason={e['reason']}" if "step" in e else f"reject reason={e['reason']}"


def main(argv):
    if argv[1:2] == ["verdicts"] and len(argv) == 3:
        doc = json.load(open(argv[2], encoding="utf-8"))
        for v in doc["vectors"]:
            try:
                got = verdict(v)
            except (ValueError, KeyError, TypeError, AttributeError):
                got = "error"
            print(f"{v['name']}\t{got}")
        return 0
    path = argv[1] if len(argv) > 1 else os.path.join(HERE, "vectors.json")
    doc = json.load(open(path, encoding="utf-8"))
    fail = 0
    b3 = blake3(b"abc").hex() == "6437b3ac38465133ffb63b75273a8db548c558465d79db03fd359c6cd5bd9d85"
    print(("PASS" if b3 else "FAIL") + " BLAKE3 reproduces its published value")
    fail += not b3
    for v in doc["vectors"]:
        want, got = expected_of(v["expected"]), verdict(v)
        ok = want == got
        fail += not ok
        print(f"{'PASS' if ok else 'FAIL'} {v['name']}  expected={want}  got={got}")
    print(f"{len(doc['vectors'])} vectors, {fail} failure(s)")
    return 1 if fail else 0


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