open-standards/wai/live-conformance/prior_carriage_verify.py

#!/usr/bin/env python3
"""Independent verifier of the prior-carriage corpus (extensions/prior-carriage.md §15 to §19;
extensions/http-delivery.md §8.5, §8.7, §8.8, §12.5).

Written from the texts. It shares no code with wai-rs, which generates the corpus, and never
runs it. Everything those sections define is written here: a set root (an absolute `https` URI,
RFC 3986, with a host, no userinfo, query or fragment, whose path ends in `/`, at most 2048
bytes), the location hints of a `wai_priors` entry and of a WAI multivariant playlist's one
`EXT-X-SESSION-DATA` line (RFC 8216 §4.2's attribute lists, its roots separated by spaces), the
order in which a sink asks roots, and §8.7's steps over the responses each case states, with
SPEC §3.1's listing rules for a set of several files and the usage statement a listing binds as
its file `wai-usage.json`; and §8.8's: the `Available-Dictionary` value (RFC 9651's Byte Sequence),
a dictionary-coded response's checks in order, and its decoding, through `_rfc8878_zstd.py`, a
Zstandard decoder written here from RFC 8878. Where the Zstandard tool is on the PATH it also
runs the stock client's check the corpus states. And prior-carriage §17's rule over each worked
plan: savings, costs charged per sink, the generations kept within the cap, and each title's
rendition. And §18: the usage statement (SPDX 2.3 Annex D's licence expressions, with the
ai-usage statement read by the ai-usage corpus's second implementation, `ai_usage_verify.py`),
sealed sets opened with AES-128 in counter mode written here from FIPS 197 and NIST SP 800-38A
(`_fips197_aes.py`), and http-delivery §12.5's held sets per site. And §19: tree heads, their
pieces' BLAKE3 chaining values and the root they merge to, through the BLAKE3 every Python
verifier of WAI's corpora shares (`../claim-conformance/_blake3.py`, written from the BLAKE3
specification), and fetches by range from several sources, verified per piece.

For every case it requires the outcome the corpus states, field for field; that every reason
http-delivery §13 lists for §8.7 and §8.8, and every reason prior-carriage §20 lists, appears in
the corpus; then it mutates copies of the corpus,
one defect each, and requires every copy to fail; and it drops each of its own checks in turn
and requires the corpus to catch the weakened verifier.

    python3 prior_carriage_verify.py <live-conformance/prior-carriage>

Exit status 0 only if every case verifies, every mutated copy fails and every weakened verifier
is caught. Python 3.8 or later, the standard library only.
"""

import base64
import copy
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import _rfc8878_zstd as ZSTD  # noqa: E402
import _fips197_aes as AES  # noqa: E402
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import _blake3 as B3  # noqa: E402
import ai_usage_verify as AIU  # noqa: E402

HTTP_DELIVERY = os.path.join(HERE, "..", "extensions", "http-delivery.md")
PRIOR_CARRIAGE = os.path.join(HERE, "..", "extensions", "prior-carriage.md")
WELL_KNOWN = "/.well-known/wai-sets/"
DATA_ID = "science.transaction.wai.sets"
MAX_ROOT = 2048
MAX_ROOTS = 8
DCZ_MAGIC = bytes([0x5E, 0x2A, 0x4D, 0x18, 0x20, 0x00, 0x00, 0x00])
DCB_MAGIC = bytes([0xFF, 0x44, 0x43, 0x42])
MB = 1 << 20

DROPPED = None


def drops(name):
    return DROPPED == name


def sha(b):
    return hashlib.sha256(b).hexdigest()


# --- set roots (http-delivery §8.7; RFC 3986) -------------------------------------------------

UNRESERVED = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~")
SUB_DELIMS = set(b"!$&'()*+,;=")
HEXDIG = set(b"0123456789ABCDEFabcdef")


def chars_ok(s, extra):
    """RFC 3986 characters: unreserved, sub-delims, `extra`, or a percent-encoded octet."""
    b = s.encode("utf-8")
    i = 0
    while i < len(b):
        if b[i] == 0x25:
            if not (i + 2 < len(b) and b[i + 1] in HEXDIG and b[i + 2] in HEXDIG):
                return False
            i += 3
            continue
        if not (b[i] in UNRESERVED or b[i] in SUB_DELIMS or b[i] in extra):
            return False
        i += 1
    return True


def is_set_root(s):
    if not isinstance(s, str) or len(s.encode("utf-8")) > MAX_ROOT:
        return False
    if not s.startswith("https://") and not drops("scheme"):
        return False
    rest = s.split("://", 1)[1] if "://" in s else s
    if ("?" in rest or "#" in rest) and not drops("query"):
        return False
    rest = rest.split("?", 1)[0].split("#", 1)[0]
    if "/" not in rest:
        return False
    authority, path = rest[:rest.index("/")], rest[rest.index("/"):]
    if "@" in authority:
        if not drops("userinfo"):
            return False
        authority = authority.split("@", 1)[1]
    if authority.startswith("["):
        close = authority.find("]")
        if close < 0:
            return False
        literal, after = authority[1:close], authority[close + 1:]
        if not literal or "%" in literal or not chars_ok(literal, b":"):
            return False
        host, port = "[]", (after[1:] if after.startswith(":") else (None if after == "" else False))
        if port is False:
            return False
    else:
        host, port = (authority.rsplit(":", 1) if ":" in authority else (authority, None))
        if not chars_ok(host, b""):
            return False
    if host == "":
        return False
    if port is not None and not (port != "" and all(c in "0123456789" for c in port)):
        return False
    if not path.endswith("/") and not drops("trailing_slash"):
        return False
    return chars_ok(path, b":@/")


# --- hints ------------------------------------------------------------------------------------

def take(hint, i, root):
    """The element at i; False once the hint holds its eight roots."""
    roots, malformed = hint
    if len(roots) == MAX_ROOTS and not drops("eight"):
        return False
    if isinstance(root, str) and is_set_root(root):
        if root not in roots or drops("dedupe"):
            roots.append(root)
    else:
        malformed.append(i)
    return True


def catalog_hint(entry):
    hint = ([], [])
    v = entry.get("locations") if isinstance(entry, dict) else None
    if isinstance(v, list):
        for i, e in enumerate(v):
            if not take(hint, i, e):
                break
    return hint


def attributes(s):
    """An attribute list (RFC 8216 §4.2) as name -> (quoted, value), or None."""
    out = {}
    while True:
        m = re.match(r"([A-Z0-9-]+)=", s)
        if not m:
            return None
        name = m.group(1)
        if name in out:
            return None
        s = s[m.end():]
        if s.startswith('"'):
            end = s.find('"', 1)
            if end < 0 or "\r" in s[1:end] or "\n" in s[1:end]:
                return None
            out[name] = (True, s[1:end])
            s = s[end + 1:]
        else:
            m = re.match(r'[^",\s]+', s)
            if not m:
                return None
            out[name] = (False, m.group(0))
            s = s[m.end():]
        if s == "":
            return out
        if not s.startswith(","):
            return None
        s = s[1:]


def playlist_hint(text):
    """The first line of the DATA-ID with no LANGUAGE; its VALUE split at every space."""
    hint = ([], [])
    for line in text.split("\n"):
        if line.endswith("\r"):
            line = line[:-1]
        if not line.startswith("#EXT-X-SESSION-DATA:"):
            continue
        a = attributes(line[len("#EXT-X-SESSION-DATA:"):])
        if a is None or a.get("DATA-ID") != (True, DATA_ID):
            continue
        if "LANGUAGE" in a and not drops("language"):
            continue
        v = a.get("VALUE")
        if v is not None and v[0] and ("URI" not in a or drops("uri_line")):
            for i, e in enumerate([v[1]] if drops("split") else v[1].split(" ")):
                if not take(hint, i, e):
                    break
        if not drops("first_line"):
            break
    return hint


def well_known(origin):
    if not origin.startswith("https://") or "/" in origin[len("https://"):]:
        return None
    r = origin + WELL_KNOWN
    return r if is_set_root(r) else None


def order(own, hint_roots, origin_root):
    out = []
    for r in list(own) + list(hint_roots) + ([origin_root] if origin_root else []):
        if r not in out:
            out.append(r)
    return out


# --- a fetch (http-delivery §8.7 steps 1-4; SPEC §3.1 listings) ------------------------------

SEGMENT = re.compile(r"[A-Za-z0-9_-][A-Za-z0-9_.-]*")


def listing(b):
    """SPEC §3.1's listing: [(path, sha256 hex)], or None where it breaks a rule."""
    if not b or not b.endswith(b"\n") or b"\r" in b:
        return None
    lines = []
    for raw in b[:-1].split(b"\n"):
        try:
            t = raw.decode("utf-8")
        except UnicodeDecodeError:
            return None
        m = re.fullmatch(r"([0-9a-f]{64})  (.+)", t)
        if not m:
            return None
        path = m.group(2)
        if not all(SEGMENT.fullmatch(seg) and not seg.endswith(".") for seg in path.split("/")):
            return None
        lines.append((path, m.group(1)))
    paths = [p for p, _ in lines]
    if any(paths[i].encode() >= paths[i + 1].encode() for i in range(len(paths) - 1)):
        return None
    if len({p.lower() for p in paths}) != len(paths):
        return None
    if any(q.startswith(p + "/") for p in paths for q in paths):
        return None
    return lines


def available_dictionary(digest_hex):
    return ":" + base64.b64encode(bytes.fromhex(digest_hex)).decode() + ":"


def parse_available_dictionary(v):
    """RFC 9651 §3.3.5's Byte Sequence of exactly 32 bytes, padded and canonical, with spaces
    around it discarded: the digest in hexadecimal, or None."""
    v = v.strip(" ")
    if len(v) < 2 or v[0] != ":" or v[-1] != ":":
        return None
    inner = v[1:-1]
    if not re.fullmatch(r"[A-Za-z0-9+/]*={0,2}", inner) or len(inner) % 4:
        return None
    b = base64.b64decode(inner)
    if len(b) != 32 or base64.b64encode(b).decode() != inner:
        return None
    return b.hex()


def window_admitted(window, dict_len):
    if window > 128 * MB and not drops("window"):
        return False
    return window <= 8 * MB or 4 * window <= 5 * dict_len or drops("window")


def check_coded(coding, body, named, limit):
    """http-delivery §8.8's checks 1 to 5: (decoded bytes, None) or (None, reason)."""
    if (named is None or coding != "dcz") and not drops("unrequested"):
        return None, "dictionary_unrequested"
    if named is None:
        return None, "dictionary_coding_malformed"
    d_hex, dictionary = named
    if (len(body) < 40 or body[:8] != DCZ_MAGIC) and not drops("magic"):
        return None, "dictionary_coding_malformed"
    if body[8:40] != bytes.fromhex(d_hex) and not drops("dictionary_hash"):
        return None, "dictionary_mismatch"
    frame = body[40:]
    try:
        window = ZSTD.frame_window(frame)[0]
    except (ZSTD.ZstdError, IndexError, KeyError):
        return None, "dictionary_coding_malformed"
    if not window_admitted(window, len(dictionary)):
        return None, "dictionary_window_too_large"
    try:
        return ZSTD.decode(frame, dictionary, limit), None
    except ZSTD.TooLarge:
        return None, "set_too_large"
    except (ZSTD.ZstdError, IndexError, KeyError):
        return None, "dictionary_coding_malformed"


def outcome_of(digest, r, named, limit):
    """(bytes or None, reason or None, ask the root again without the dictionary)."""
    if r is None or (r["status"] != 200 and not drops("status")):
        return None, "set_location_unavailable", False
    enc = r["content_encoding"]
    if enc is None:
        if len(r["body"]) > limit and not drops("bound"):
            return None, "set_too_large", False
        body = r["body"]
    elif enc in ("dcz", "dcb"):
        body, why = check_coded(enc, r["body"], named, limit)
        if why is not None:
            return None, why, why != "set_too_large" and named is not None and not drops("retry")
    elif not drops("coding"):
        return None, "content_coding_present", False
    else:
        body = r["body"]
    if sha(body) != digest and not drops("digest"):
        return None, "set_location_mismatch", False
    return body, None, False


def fetch_one(digest, roots, limit, get, attempts, named=None):
    for i, root in enumerate(roots):
        url = root + digest
        asked = named
        while True:
            r = get((url, available_dictionary(asked[0]) if asked else None))
            body, why, again = outcome_of(digest, r, asked, limit)
            attempts.append([url, why or "held"])
            if why is None:
                return i, body
            if not again:
                break
            asked = None
    return None


def fetch(digest, form, roots, limit, get, named=None):
    """(attempts, result, the usage statement a held listing binds or None)."""
    attempts = []
    listing_form = form == "wai.prior.set"
    got = fetch_one(digest, roots, limit, get, attempts, None if listing_form else named)
    if got is None:
        return attempts, "prior_not_held", None
    at, body = got
    if not listing_form:
        return attempts, "held", None
    lines = listing(body)
    if lines is None:
        return attempts, "set_malformed", None
    file_roots = roots if drops("listing_root_first") else [roots[at]] + [r for i, r in enumerate(roots) if i != at]
    held = {}
    for path, d in lines:
        f = fetch_one(d, file_roots, limit, get, attempts)
        if f is None:
            return attempts, "part_mismatch", None
        held[path] = f[1]
    # prior-carriage §18.1: the statement at the top-level path wai-usage.json, read once held.
    bound = [b for p, b in held.items() if p == "wai-usage.json" or (drops("usage_top_level") and p.endswith("wai-usage.json"))]
    if not bound or drops("listing_usage"):
        return attempts, "held", None
    usage = read_usage(bound[0])
    if usage is None:
        return attempts, "set_usage_malformed", None
    return attempts, "held", usage


def pin_target(pin):
    """The two members of a pin a fetch reads: `sha256` and `form`."""
    d = pin.get("sha256")
    if not (isinstance(d, str) and re.fullmatch(r"[0-9a-f]{64}", d)):
        return None
    f = pin.get("form")
    if f is not None and not isinstance(f, str):
        return None
    return d, f


# --- amortisation (prior-carriage §17) ---------------------------------------------------------

INSTANT = re.compile(r"[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z")


def is_uint(v):
    return isinstance(v, int) and not isinstance(v, bool) and 0 <= v <= (1 << 53) - 1


def read_plan(p):
    """The plan's parts, or None where it does not read (§20's plan_malformed)."""
    try:
        f, t = p["span"]["from"], p["span"]["to"]
        if not (isinstance(f, str) and isinstance(t, str) and INSTANT.fullmatch(f) and INSTANT.fullmatch(t) and f < t):
            return None
        cap = p["cap"]
        if not is_uint(cap) or cap < 1:
            return None
        gens = []
        for g in p["generations"]:
            gid, epoch, size, delta = g["id"], g["epoch"], g["set_bytes"], g["delta"]
            if not (isinstance(gid, str) and gid and is_uint(epoch) and is_uint(size) and size >= 1):
                return None
            if gid in [x[0] for x in gens] or (gens and gens[-1][1] >= epoch):
                return None
            if delta is not None:
                if not (isinstance(delta["from"], str) and delta["from"] in [x[0] for x in gens]
                        and is_uint(delta["bytes"]) and delta["bytes"] >= 1):
                    return None
                delta = (delta["from"], delta["bytes"])
            gens.append((gid, epoch, size, delta))
        if not gens:
            return None
        titles = {}
        for x in p["titles"]:
            tid, gen = x["id"], x["generation"]
            if not (isinstance(tid, str) and tid and tid not in titles and gen in [g[0] for g in gens]
                    and is_uint(x["pinned_bytes"]) and is_uint(x["classical_bytes"])):
                return None
            titles[tid] = (gen, x["pinned_bytes"], x["classical_bytes"])
        sinks = p["sinks"]
        if not sinks:
            return None
        for sk in sinks:
            holds = sk["holds"]
            if not (isinstance(holds, list) and all(isinstance(h, str) for h in holds) and len(set(holds)) == len(holds)
                    and all(h in [g[0] for g in gens] for h in holds)):
                return None
            if not isinstance(sk["demand"], dict) or not all(k in titles and is_uint(v) for k, v in sk["demand"].items()):
                return None
        return cap, gens, titles, sinks
    except (KeyError, TypeError):
        return None


def plan(p):
    got = read_plan(p)
    if got is None:
        return "plan_malformed"
    cap, gens, titles, sinks = got
    if drops("per_sink"):
        # A verifier that pools the sinks' fetches, as if the set were fetched once for all.
        pooled = {}
        for sk in sinks:
            for k, v in sk["demand"].items():
                pooled[k] = pooled.get(k, 0) + v
        sinks = [{"holds": sinks[0]["holds"], "demand": pooled}] * len(sinks)
    out = []
    for sk in sinks:
        held, demand = sk["holds"], {k: v for k, v in sk["demand"].items() if v > 0}
        rows = []
        for i, (gid, _, size, delta) in enumerate(gens):
            mine = [k for k in demand if titles[k][0] == gid]
            if not mine and gid not in held:
                continue
            saving = sum(demand[k] * (titles[k][2] - titles[k][1]) for k in mine
                         if titles[k][2] > titles[k][1] or drops("negative_titles"))
            if gid in held:
                cost, via = 0, "held"
            elif delta is not None and delta[0] in held and not drops("delta"):
                cost, via = delta[1], "delta"
            else:
                cost, via = size, "set"
            rows.append((i, saving, cost, via, saving - cost))
        worth = [r for r in rows if r[4] > 0 or (drops("strict") and r[4] >= 0 and r[3] != "held")]
        worth.sort(key=lambda r: (-r[4], r[0] if drops("tie") else -r[0]))
        kept = [r[0] for r in worth][:len(worth) if drops("cap") else cap]
        if not drops("filler"):
            for r in sorted((r for r in rows if r[3] == "held" and r[0] not in kept), key=lambda r: -r[0]):
                if len(kept) < cap:
                    kept.append(r[0])
        fetch = [[gens[r[0]][0], r[3], r[2]] for r in rows if r[0] in kept and r[3] != "held"]
        total = sum(f[2] for f in fetch)
        classical = 0
        per_title = {}
        for k in sorted(demand):
            gen, pinned_b, classical_b = titles[k]
            gi = [g[0] for g in gens].index(gen)
            pin = gi in kept and classical_b > pinned_b
            per_title[k] = "pinned" if pin else "classical"
            total += demand[k] * (pinned_b if pin else classical_b)
            classical += demand[k] * classical_b
        out.append({
            "generations": [{"id": gens[r[0]][0], "saving": r[1], "cost": r[2], "via": r[3], "net": r[4]} for r in rows],
            "hold": [gens[i][0] for i in kept],
            "fetch": fetch,
            "evict": [h for h in held if h not in [gens[i][0] for i in kept]],
            "titles": per_title,
            "bytes": total,
            "classical_bytes": classical,
        })
    return out


# --- usage, sealed sets and held sets (prior-carriage §18; http-delivery §12.5) ----------------

IDSTRING = re.compile(r"[A-Za-z0-9.-]+")


def licence_expression(s):
    """SPDX 2.3 Annex D, operators upper case, single spaces, identifiers by their characters."""
    if not isinstance(s, str) or not 1 <= len(s.encode()) <= 1024:
        return False
    toks = []
    for word in s.split(" "):
        if word == "":
            return False
        while word.startswith("("):
            toks.append("(")
            word = word[1:]
        n = 0
        while word.endswith(")"):
            n += 1
            word = word[:-1]
        if word:
            toks.append(word)
        toks += [")"] * n

    def simple(t):
        if t in ("AND", "OR", "WITH", "(", ")"):
            return False
        if ":" in t:
            doc, ref = t.split(":", 1)
            return (doc.startswith("DocumentRef-") and bool(IDSTRING.fullmatch(doc[12:]))
                    and ref.startswith("LicenseRef-") and bool(IDSTRING.fullmatch(ref[11:])))
        if t.startswith("LicenseRef-"):
            return bool(IDSTRING.fullmatch(t[11:]))
        return bool(IDSTRING.fullmatch(t[:-1] if t.endswith("+") else t))

    pos = [0]

    def expr(depth):
        if depth > 32:
            return False
        while True:
            t = toks[pos[0]] if pos[0] < len(toks) else None
            if t == "(":
                pos[0] += 1
                if not expr(depth + 1) or pos[0] >= len(toks) or toks[pos[0]] != ")":
                    return False
                pos[0] += 1
            elif t is not None and simple(t):
                pos[0] += 1
                if pos[0] < len(toks) and toks[pos[0]] == "WITH":
                    pos[0] += 1
                    e = toks[pos[0]] if pos[0] < len(toks) else None
                    if e is None or e in ("AND", "OR", "WITH") or not IDSTRING.fullmatch(e):
                        return False
                    pos[0] += 1
            else:
                return False
            if pos[0] < len(toks) and toks[pos[0]] in ("AND", "OR"):
                pos[0] += 1
            else:
                return True

    if drops("operators"):
        toks = [t.upper() if t.lower() in ("and", "or", "with") else t for t in toks]
    return expr(0) and pos[0] == len(toks)


def read_usage(b):
    """prior-carriage §18.1: the statement's fields, or None (set_usage_malformed)."""
    if len(b) > 65536:
        return None
    try:
        v = AIU.strict_json(b.decode("utf-8"))
    except (ValueError, UnicodeDecodeError):
        return None
    names = {"kind", "licence", "terms", "content_derived", "redistribute", "ai_usage"}
    if not isinstance(v, dict) or (set(v) != names and not drops("members")):
        return None
    if v.get("kind") != "wai-set-usage":
        return None
    lic, terms = v.get("licence"), v.get("terms")
    if lic is not None and not licence_expression(lic):
        return None
    if terms is not None and not (isinstance(terms, str) and 8 <= len(terms) <= 2048 and terms.startswith("https://")
                                  and all(0x21 <= ord(c) <= 0x7E for c in terms)):
        return None
    if not isinstance(v.get("content_derived"), bool) or not isinstance(v.get("redistribute"), bool):
        return None
    ai = v.get("ai_usage")
    digest = None
    if ai is not None:
        try:
            digest = AIU.read_statement(ai)["digest"].hex()
        except AIU.Refuse:
            return None
    return {"licence": lic, "terms": terms, "content_derived": v["content_derived"], "redistribute": v["redistribute"],
            "ai_usage": digest}


ADMITTED = ("", "wai.prior.safetensors", "wai.prior.gguf")


def open_sealed(b, keys):
    """prior-carriage §18.2's steps: (header fields, plaintext) or the reason."""
    if len(b) < 5 or b[:4] != b"WSK1":
        return None, "sealed_malformed"
    k = b[4]
    if k > 64 or len(b) < 73 + k:
        return None, "sealed_malformed"
    form = b[5:5 + k]
    if not all(0x21 <= c <= 0x7E for c in form):
        return None, "sealed_malformed"
    kid, iv, inner = b[5 + k:21 + k], b[21 + k:37 + k], b[37 + k:69 + k]
    u = int.from_bytes(b[69 + k:73 + k], "little")
    if u > 65536 or len(b) < 81 + k + u:
        return None, "sealed_malformed"
    n = int.from_bytes(b[73 + k + u:81 + k + u], "little")
    if not 1 <= n <= 1 << 40 or len(b) - (81 + k + u) != n:
        return None, "sealed_malformed"
    if form.decode() not in ADMITTED and not drops("inner_form"):
        return None, "sealed_inner_form"
    usage = None
    if u:
        usage = read_usage(b[73 + k:73 + k + u])
        if usage is None:
            return None, "set_usage_malformed"
    key = keys.get(kid.hex())
    if key is None:
        return None, "set_key_absent"
    plain = AES.ctr(bytes.fromhex(key), iv, b[81 + k + u:])
    if sha(plain) != inner.hex() and not drops("inner_digest"):
        return None, "set_key_mismatch"
    return {"result": "held", "inner_form": form.decode(), "inner": inner.hex(), "usage": usage}, None


def held_sets_for(store, site, wanted, enabled):
    if not enabled and not drops("off"):
        return None
    out = {}
    for cap, d in wanted:
        if any((h["site"] == site or drops("site")) and h["capability"] == cap and h["sha256"] == d
               and (not h["revoked"] or drops("revoked")) for h in store):
            lst = out.setdefault(cap, [])
            if d not in lst:
                lst.append(d)
    return out


def may_name(store, site, d):
    return any((h["site"] == site or drops("site")) and h["sha256"] == d and not h["revoked"] for h in store)


# --- the chunk-tree form (prior-carriage §19) -----------------------------------------------------

TREE_INNER = ADMITTED + ("wai.prior.sealed",)


def merge(values, top):
    if len(values) == 1:
        return values[0]
    left = 1 << ((len(values) - 1).bit_length() - 1)
    if drops("left_split"):
        left = len(values) // 2
    return B3.parent_cv(merge(values[:left], False), merge(values[left:], False), root=top)


def read_head(b):
    """§19.1's checks: the head's fields, or the reason."""
    if len(b) < 6 or b[:4] != b"WTR1":
        return None, "tree_head_malformed"
    g, k = b[4], b[5]
    if g > 16 or k > 64 or len(b) < 82 + k:
        return None, "tree_head_malformed"
    form = b[6:6 + k]
    if not all(0x21 <= c <= 0x7E for c in form):
        return None, "tree_head_malformed"
    n_bytes = int.from_bytes(b[6 + k:14 + k], "little")
    inner, root = b[14 + k:46 + k], b[46 + k:78 + k]
    n = int.from_bytes(b[78 + k:82 + k], "little")
    piece = 1024 << g
    if n_bytes == 0 or n != -(-n_bytes // piece) or not (2 <= n <= 65536 or drops("one_piece")) or len(b) != 82 + k + 32 * n:
        return None, "tree_head_malformed"
    if form.decode() not in TREE_INNER and not drops("tree_inner"):
        return None, "tree_inner_form"
    values = [b[82 + k + 32 * i:114 + k + 32 * i] for i in range(n)]
    if merge(values, True) != root and not drops("tree_root"):
        return None, "tree_root_mismatch"
    return {"g": g, "inner_form": form.decode(), "length": n_bytes, "inner": inner.hex(), "root": root.hex(),
            "values": values, "piece": piece}, None


def fetch_pieces(head, sources):
    attempts, held = [], []
    for i in range(len(head["values"])):
        s, e = i * head["piece"], min((i + 1) * head["piece"], head["length"])
        got = None
        for j in range(len(sources)):
            src = sources[(i + j) % len(sources)] if not drops("first_source") else sources[j]
            body = src["body"] if src["answers"] else None
            if body is None:
                outcome = "set_location_unavailable"
            elif src["content_encoding"] is not None:
                outcome = "content_coding_present"
            elif len(body[s:e]) != e - s:
                outcome = "content_range_mismatch"
            elif B3.subtree_cv(body[s:e], s) != head["values"][i] and not drops("piece"):
                outcome = "tree_piece_mismatch"
            else:
                got = body[s:e]
                outcome = "held"
            attempts.append([src["name"], i, outcome])
            if got is not None:
                break
        if got is None:
            return attempts, "prior_not_held", None
        held.append(got)
    whole = b"".join(held)
    if sha(whole) != head["inner"] and not drops("whole_digest"):
        return attempts, "tree_digest_mismatch", None
    return attempts, "held", sha(whole)


# --- the corpus ---------------------------------------------------------------------------------

def reasons_of_8_7_and_8_8():
    text = open(HTTP_DELIVERY, encoding="utf-8").read()
    at = text.index("\n## 13. Reasons")
    end = text.index("\n## ", at + 1)
    out = []
    for line in text[at:end].splitlines():
        if line.startswith("| `") and ("| §8.7" in line or "| §8.8" in line):
            out.extend(re.findall(r"`([a-z_]+)`", line.split("|")[1]))
    return out


def body_of(spec, sets, files):
    if "hex" in spec:
        return bytes.fromhex(spec["hex"])
    b = bytearray(sets[spec["set"]] if "set" in spec else files[spec["file"]])
    if "flip" in spec:
        b[spec["flip"]] ^= 0x01
    if "truncate" in spec:
        del b[spec["truncate"]:]
    return bytes(b)


def reasons_of_20():
    text = open(PRIOR_CARRIAGE, encoding="utf-8").read()
    at = text.index("\n## 20. Refusals")
    end = text.find("\n## ", at + 1)
    return re.findall(r"^\| `([a-z_]+)` \|", text[at:end if end > 0 else len(text)], re.M)


def strings(v, out):
    if isinstance(v, str):
        out.append(v)
    elif isinstance(v, list):
        for x in v:
            strings(x, out)
    elif isinstance(v, dict):
        for x in v.values():
            strings(x, out)


def check(doc, corpus, quiet=False):
    fails = []
    sets = {}
    files = {}
    for dirpath, _, names in os.walk(os.path.join(corpus, "files")):
        for n in names:
            p = os.path.join(dirpath, n)
            files[os.path.relpath(p, corpus).replace(os.sep, "/")] = open(p, "rb").read()
    for s in doc["sets"]:
        b = open(os.path.join(corpus, s["path"]), "rb").read()
        sets[s["name"]] = b
        if sha(b) != s["sha256"] or len(b) != s["length"]:
            fails.append("set %s: its digest or length is not the file's" % s["name"])
    for r in doc["roots"]:
        if is_set_root(r["root"]) != r["is_set_root"]:
            fails.append("root %s" % r["name"])
    for h in doc["hints"]:
        roots, malformed = catalog_hint(h["entry"]) if "entry" in h else playlist_hint(h["playlist"])
        if {"roots": roots, "malformed": malformed} != h["outcome"]:
            fails.append("hint %s: %s %s" % (h["name"], roots, malformed))
    for f in doc["fetches"]:
        roots, _ = catalog_hint(f["entry"])
        if drops("pin_member") and isinstance(f["pin"].get("locations"), list):
            roots = [r for r in f["pin"]["locations"] if isinstance(r, str)] + roots
        o = order(f["own"], roots, well_known(f["origin"]))
        target = pin_target(f["pin"])
        if target is None:
            fails.append("fetch %s: the pin does not read" % f["name"])
            continue
        responses = {}
        for r in f["responses"]:
            responses[(r["url"], r["available_dictionary"])] = {
                "status": r["status"], "content_encoding": r["content_encoding"], "body": body_of(r["body"], sets, files)}
        named = None
        if f["dictionary"] is not None:
            named = (sha(sets[f["dictionary"]]), sets[f["dictionary"]])
        attempts, result, usage = fetch(target[0], target[1], o, f["limit"], responses.get, named)
        got = {"order": o, "attempts": attempts, "result": result, "usage": usage}
        if got != f["outcome"]:
            fails.append("fetch %s: %s" % (f["name"], json.dumps(got)))
    for v in doc["fields"]:
        if "available_dictionary" in v:
            if parse_available_dictionary(v["available_dictionary"]) != v["sha256"]:
                fails.append("field %s" % v["name"])
            elif v["sha256"] is not None and available_dictionary(v["sha256"]) != v["available_dictionary"].strip(" "):
                fails.append("field %s: not written back the same" % v["name"])
        elif 'match="%s*"' % v["root_path"] != v["use_as_dictionary"]:
            fails.append("field %s" % v["name"])
    st = doc["stock"]
    body, dictionary = files.get(st["body"], b""), sets[st["dictionary"]]
    got, why = check_coded("dcz", body, (sha(dictionary), dictionary), 1 << 30)
    if why is not None or sha(got) != st["sha256"] or st["header"] != 40:
        fails.append("stock: %s" % why)
    for c in doc["plans"]:
        got = plan(c["plan"])
        if got != c["outcome"]:
            fails.append("plan %s: %s" % (c["name"], json.dumps(got)))
    for c in doc["usage"]:
        got = read_usage(c["statement"].encode("utf-8"))
        if (got if got is not None else "set_usage_malformed") != c["outcome"]:
            fails.append("usage %s: %s" % (c["name"], got))
    for c in doc["sealed"]:
        got, why = open_sealed(body_of(c["body"], sets, files), c["keys"])
        if (got if why is None else why) != c["outcome"]:
            fails.append("sealed %s: %s" % (c["name"], why or got))
    for c in doc["held_sets"]:
        got = held_sets_for(c["store"], c["site"], [tuple(w) for w in c["wanted"]], c["enabled"])
        if got != c["outcome"]:
            fails.append("held_sets %s: %s" % (c["name"], got))
    for c in doc["dictionary_naming"]:
        if may_name(c["store"], c["site"], c["sha256"]) != c["may_name"]:
            fails.append("dictionary_naming %s" % c["name"])
    for c in doc["tree_heads"]:
        b = body_of(c["head"], sets, files)
        h, why = read_head(b)
        got = why if why else {k: h[k] for k in ("inner_form", "length", "g", "inner", "root")}
        if h is not None:
            got["pieces"] = len(h["values"])
            # The root is the BLAKE3 hash of the set the head names, where the corpus holds it.
            held = [x for x in list(sets.values()) + list(files.values()) if sha(x) == h["inner"]]
            if held and B3.blake3(held[0]).hex() != h["root"]:
                fails.append("tree head %s: its root is not the BLAKE3 hash of its set" % c["name"])
        if got != c["outcome"]:
            fails.append("tree head %s: %s" % (c["name"], got))
    for c in doc["tree_fetches"]:
        h, why = read_head(body_of(c["head"], sets, files))
        if why:
            fails.append("tree fetch %s: %s" % (c["name"], why))
            continue
        srcs = [{"name": x["name"], "body": body_of(x["body"], sets, files), "answers": x["answers"],
                 "content_encoding": x["content_encoding"]} for x in c["sources"]]
        attempts, result, digest = fetch_pieces(h, srcs)
        got = {"attempts": attempts, "result": result, "sha256": digest}
        if got != c["outcome"]:
            fails.append("tree fetch %s: %s" % (c["name"], json.dumps(got)))
    stated = []
    for k in ("hints", "fetches", "plans", "usage", "sealed", "tree_heads", "tree_fetches"):
        for c in doc[k]:
            strings(c["outcome"], stated)
    for code in reasons_of_8_7_and_8_8() + reasons_of_20():
        if code == "set_location_malformed":
            given = any(h["outcome"]["malformed"] for h in doc["hints"])
        else:
            given = code in stated
        if not given:
            fails.append("no case gives " + code)
    if not quiet:
        for x in fails:
            print("FAIL", x)
    return fails


def mutations(doc):
    """Copies of the corpus, one defect each, every one of which must fail."""
    def named(d, kind, name):
        return next(c for c in d[kind] if c["name"] == name)

    def m_root(d):
        named(d, "roots", "query")["is_set_root"] = True

    def m_hint_order(d):
        o = named(d, "hints", "catalog_two_roots")["outcome"]
        o["roots"] = list(reversed(o["roots"]))

    def m_playlist(d):
        named(d, "hints", "playlist_elements_passed_over")["outcome"]["malformed"] = [1, 2, 4]

    def m_playlist_later_line(d):
        named(d, "hints", "playlist_a_later_line_ignored")["outcome"]["roots"].append(
            named(d, "hints", "playlist_one_line")["outcome"]["roots"][1])

    def m_second_root(d):
        named(d, "fetches", "second_root_after_tampered_mirror")["responses"][1]["body"]["flip"] = 7

    def m_order(d):
        named(d, "fetches", "own_root_first")["own"] = []

    def m_bound(d):
        named(d, "fetches", "past_the_bound")["limit"] = 1 << 20

    def m_listing_result(d):
        named(d, "fetches", "a_listed_file_no_root_holds")["outcome"]["result"] = "prior_not_held"

    def m_listing_usage(d):
        named(d, "fetches", "a_listing_binds_its_statement")["outcome"]["usage"]["redistribute"] = True

    def m_set_digest(d):
        d["sets"][0]["sha256"] = "00" * 32

    def m_reason_gone(d):
        d["fetches"].remove(named(d, "fetches", "a_window_past_the_limit"))

    def m_no_dictionary(d):
        named(d, "fetches", "successor_against_its_predecessor")["dictionary"] = None

    def m_wrong_dictionary(d):
        named(d, "fetches", "successor_against_its_predecessor")["dictionary"] = "title"

    def m_field(d):
        named(d, "fields", "pad_bits_set")["sha256"] = named(d, "fields", "rfc_9842_example")["sha256"]

    def m_stock(d):
        d["stock"]["sha256"] = d["sets"][0]["sha256"]

    def m_frame(d):
        named(d, "fetches", "successor_against_its_predecessor")["responses"][0]["body"]["flip"] = 300

    def m_plan_bytes(d):
        named(d, "plans", "worth_holding")["outcome"][0]["bytes"] = 10500

    def m_plan_cap(d):
        named(d, "plans", "two_generations_held_cap_two")["plan"]["cap"] = 3

    def m_plan_hold_order(d):
        h = named(d, "plans", "two_generations_held_cap_three")["outcome"][0]["hold"]
        h.reverse()

    def m_plan_delta(d):
        named(d, "plans", "successor_through_its_predecessor")["plan"]["sinks"][0]["holds"] = []

    def m_plan_malformed(d):
        named(d, "plans", "cap_zero")["plan"]["cap"] = 2

    def m_usage(d):
        named(d, "usage", "content_derived")["outcome"]["redistribute"] = True

    def m_key(d):
        c = named(d, "sealed", "opened_with_its_key")
        c["keys"] = {k: "6c" * 16 for k in c["keys"]}

    def m_sealed_inner(d):
        named(d, "sealed", "a_capability_s_one_file")["outcome"]["inner"] = d["sets"][1]["sha256"]

    def m_held_site(d):
        named(d, "held_sets", "not_what_another_site_fetched")["site"] = "a.example"

    def m_held_off(d):
        named(d, "held_sets", "off_unless_turned_on")["enabled"] = True

    def m_tree_piece(d):
        named(d, "tree_fetches", "two_peers")["sources"][1]["body"]["flip"] = 4000

    def m_tree_order(d):
        o = named(d, "tree_fetches", "a_flipped_piece_at_one_peer")["outcome"]["attempts"]
        o[3], o[4] = o[4], o[3]

    def m_tree_root(d):
        named(d, "tree_heads", "pieces_of_one_chunk")["outcome"]["root"] = "00" * 32

    return [m_root, m_hint_order, m_playlist, m_playlist_later_line, m_second_root, m_order, m_bound, m_listing_result,
            m_listing_usage, m_set_digest,
            m_reason_gone, m_no_dictionary, m_wrong_dictionary, m_field, m_stock, m_frame, m_plan_bytes, m_plan_cap,
            m_plan_hold_order, m_plan_delta, m_plan_malformed, m_usage, m_key, m_sealed_inner, m_held_site, m_held_off,
            m_tree_piece, m_tree_order, m_tree_root]


WEAKENINGS = ["scheme", "query", "userinfo", "trailing_slash", "eight", "dedupe", "uri_line", "language", "split",
              "first_line", "listing_usage", "usage_top_level", "status", "coding",
              "bound", "digest", "listing_root_first", "pin_member", "unrequested", "magic", "dictionary_hash",
              "window", "retry", "per_sink", "negative_titles", "delta", "strict", "tie", "cap", "filler", "operators",
              "members", "inner_form", "inner_digest", "off", "site", "revoked", "left_split", "one_piece", "tree_inner",
              "tree_root", "first_source", "piece", "whole_digest"]


def stock_tool(doc, corpus):
    """The stock client's check the corpus states, run with the Zstandard tool where it is on the
    PATH: the body after its header, decoded against the dictionary, has the stated SHA-256."""
    st = doc["stock"]
    tool = shutil.which("zstd")
    if tool is None:
        print("stock client: the Zstandard tool is not on the PATH; its check was not run")
        return []
    path = next(s["path"] for s in doc["sets"] if s["name"] == st["dictionary"])
    body = open(os.path.join(corpus, st["body"]), "rb").read()[st["header"]:]
    r = subprocess.run([tool, "-q", "-d", "-c", "-D", os.path.join(corpus, path)], input=body, capture_output=True)
    ok = r.returncode == 0 and sha(r.stdout) == st["sha256"]
    print("stock client: %s" % ("the Zstandard tool decodes the body to the stated SHA-256" if ok else "FAILED"))
    return [] if ok else ["stock client: the Zstandard tool's decoding differs"]


def main(argv):
    global DROPPED
    if len(argv) != 2:
        print(__doc__)
        return 2
    corpus = argv[1]
    doc = json.load(open(os.path.join(corpus, "vectors.json"), encoding="utf-8"))
    present = sorted(os.listdir(corpus))
    fails = check(doc, corpus)
    extra = [p for p in present if p not in ("vectors.json", "README.md", "files")]
    if extra:
        fails.append("files the corpus does not hold: %s" % extra)
    kinds = ("roots", "hints", "fetches", "fields", "plans", "usage", "sealed", "held_sets", "dictionary_naming",
             "tree_heads", "tree_fetches")
    if not AES.self_test():
        fails.append("AES-128 fails FIPS 197 C.1 or SP 800-38A F.5.1")
    print("%d cases (%s): %d failure(s)" % (sum(len(doc[k]) for k in kinds),
                                           ", ".join("%d %s" % (len(doc[k]), k) for k in kinds), len(fails)))
    fails += stock_tool(doc, corpus)
    caught = 0
    for m in mutations(doc):
        d = copy.deepcopy(doc)
        m(d)
        if check(d, corpus, quiet=True):
            caught += 1
        else:
            fails.append("mutation %s was not caught" % m.__name__)
    print("%d of %d mutated copies fail" % (caught, len(mutations(doc))))
    weak = 0
    for w in WEAKENINGS:
        DROPPED = w
        if check(doc, corpus, quiet=True):
            weak += 1
        else:
            fails.append("a verifier without its %s check passes the corpus" % w)
    DROPPED = None
    print("%d of %d weakened verifiers are caught" % (weak, len(WEAKENINGS)))
    for x in fails:
        print("FAIL", x)
    return 1 if fails else 0


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