#!/usr/bin/env python3
"""Independent verifier of the end-to-end encrypted rooms corpus (extensions/confidentiality.md).

Written from the extension's text and the documents it reads. It shares no code with the
reference (joule-moq's secure_objects and joule-moq-mls), which generates the corpus, and
never runs it. Standard library only:

- AES (FIPS 197), AES-GCM (NIST SP 800-38D) and RFC 9605's AES-CTR with HMAC-SHA-256 are
  written here from their specifications, FOR TEST VECTORS ONLY: no constant time, no
  speed, no use outside this file;
- HKDF (RFC 5869) and HMAC come from Python's hmac and hashlib;
- BLAKE3 from ../claim-conformance/_blake3.py, and Ed25519 under WAI's strict rule (SPEC §7
  *Signatures*) from ../claim-conformance/_ed25519.py.

What it checks:

- every test vector of RFC 9605 Appendix C (rfc9605.json) and of
  draft-ietf-moq-secure-objects-01 Appendix B (secure-objects-01.json), the latter exact in
  the GCM suites and, in the counter-mode suites, the difference confidentiality §11.2
  states;
- every vector of vectors.json: the confidentiality records (§3.3); the key schedule
  (§4.2, §4.5, §5.7); every sealed object, sealed again from its inputs with every derived
  value required, then opened through the receiving checks (§6); every room, replayed step
  by step through the model of §5.4 to §5.6 (lifecycle, the clock's proposals, the
  committer, what a commit holds, each epoch's roster) and every delivery of every sealed
  object opened or refused (§6);
- that the corpus raises every reason of §13, and that §13 lists exactly the codes the
  corpus lists.

The boundary: it does not implement MLS. Each epoch's application secret is taken from the
corpus as given (§14), and everything after it is checked.

Then it mutates copies of the corpus, one defect each, and requires every copy to fail.

    python3 -I e2ee_verify.py <live-conformance/e2ee>

Exit status 0 only if every vector verifies and every mutated copy fails. Python 3.8 or later.
"""

import copy
import hashlib
import hmac as _hmac
import json
import os
import sys

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


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


# ================================================================ AES, FOR VECTORS ONLY

def _xtime(a):
    a <<= 1
    return (a ^ 0x11B) & 0xFF if a & 0x100 else a


def _gmul(a, b):
    r = 0
    while b:
        if b & 1:
            r ^= a
        a = _xtime(a)
        b >>= 1
    return r


def _make_sbox():
    """FIPS 197 §5.1.1: the inverse in GF(2^8), then the affine transformation."""
    inv = [0] * 256
    for x in range(1, 256):
        for y in range(1, 256):
            if _gmul(x, y) == 1:
                inv[x] = y
                break
    box = []
    for x in range(256):
        b = inv[x]
        s = b
        for i in range(1, 5):
            s ^= ((b << i) | (b >> (8 - i))) & 0xFF
        box.append(s ^ 0x63)
    return box


SBOX = _make_sbox()


def key_expansion(key):
    """FIPS 197 §5.2: the round keys."""
    nk = len(key) // 4
    if nk not in (4, 6, 8) or len(key) % 4:
        raise ValueError("an AES key is 16, 24 or 32 bytes")
    nr = nk + 6
    w = [list(key[4 * i:4 * i + 4]) for i in range(nk)]
    rcon = 1
    for i in range(nk, 4 * (nr + 1)):
        t = list(w[i - 1])
        if i % nk == 0:
            t = [SBOX[b] for b in t[1:] + t[:1]]
            t[0] ^= rcon
            rcon = _xtime(rcon)
        elif nk > 6 and i % nk == 4:
            t = [SBOX[b] for b in t]
        w.append([w[i - nk][j] ^ t[j] for j in range(4)])
    return [sum(w[4 * r:4 * r + 4], []) for r in range(nr + 1)]


def encrypt_block(rk, block):
    """FIPS 197 §5.1, the state column by column."""
    s = [b ^ k for b, k in zip(block, rk[0])]
    nr = len(rk) - 1
    for r in range(1, nr + 1):
        s = [SBOX[b] for b in s]
        s = [s[4 * ((c + i) % 4) + i] for c in range(4) for i in range(4)]
        if r != nr:
            m = []
            for c in range(4):
                a = s[4 * c:4 * c + 4]
                m += [_gmul(a[0], 2) ^ _gmul(a[1], 3) ^ a[2] ^ a[3],
                      a[0] ^ _gmul(a[1], 2) ^ _gmul(a[2], 3) ^ a[3],
                      a[0] ^ a[1] ^ _gmul(a[2], 2) ^ _gmul(a[3], 3),
                      _gmul(a[0], 3) ^ a[1] ^ a[2] ^ _gmul(a[3], 2)]
            s = m
        s = [b ^ k for b, k in zip(s, rk[r])]
    return bytes(s)


# ---------------------------------------------------------------- GCM (SP 800-38D)

_R = 0xE1 << 120


def _gf_mult(x, y):
    """SP 800-38D §6.3, Algorithm 1."""
    z, v = 0, y
    for i in range(127, -1, -1):
        if (x >> i) & 1:
            z ^= v
        v = (v >> 1) ^ _R if v & 1 else v >> 1
    return z


def _ghash(h, data):
    y = 0
    for i in range(0, len(data), 16):
        y = _gf_mult(y ^ int.from_bytes(data[i:i + 16], "big"), h)
    return y


def _pad16(b):
    return b + bytes(-len(b) % 16)


def _inc32(block):
    return block[:12] + ((int.from_bytes(block[12:], "big") + 1) % (1 << 32)).to_bytes(4, "big")


def _gctr(rk, icb, data):
    out, cb = bytearray(), icb
    for i in range(0, len(data), 16):
        ks = encrypt_block(rk, list(cb))
        out += bytes(a ^ b for a, b in zip(data[i:i + 16], ks))
        cb = _inc32(cb)
    return bytes(out)


def _gcm_tag(rk, h, j0, aad, c):
    s = _ghash(h, _pad16(aad) + _pad16(c) + (8 * len(aad)).to_bytes(8, "big") + (8 * len(c)).to_bytes(8, "big"))
    return _gctr(rk, j0, s.to_bytes(16, "big"))


def gcm_seal(key, nonce, aad, pt):
    """SP 800-38D §7.1 with a 96-bit IV and a 128-bit tag."""
    rk = key_expansion(key)
    h = int.from_bytes(encrypt_block(rk, [0] * 16), "big")
    j0 = nonce + b"\x00\x00\x00\x01"
    c = _gctr(rk, _inc32(j0), pt)
    return c + _gcm_tag(rk, h, j0, aad, c)


def gcm_open(key, nonce, aad, ct):
    """SP 800-38D §7.2: the plaintext, or None."""
    if len(ct) < 16:
        return None
    c, t = ct[:-16], ct[-16:]
    rk = key_expansion(key)
    h = int.from_bytes(encrypt_block(rk, [0] * 16), "big")
    j0 = nonce + b"\x00\x00\x00\x01"
    if not _hmac.compare_digest(_gcm_tag(rk, h, j0, aad, c), t):
        return None
    return _gctr(rk, _inc32(j0), c)


# --------------------------------------------- RFC 9605 §4.5.1: AES-CTR with HMAC-SHA-256

def _ctr(enc_key, nonce, data):
    rk = key_expansion(enc_key)
    out, n = bytearray(), int.from_bytes(nonce + bytes(4), "big")
    for i in range(0, len(data), 16):
        ks = encrypt_block(rk, list(n.to_bytes(16, "big")))
        out += bytes(a ^ b for a, b in zip(data[i:i + 16], ks))
        n = (n + 1) % (1 << 128)
    return bytes(out)


def rfc9605_tag(auth_key, nonce, aad, ct, nt):
    data = len(aad).to_bytes(8, "big") + len(ct).to_bytes(8, "big") + nt.to_bytes(8, "big") + nonce + aad + ct
    return _hmac.new(auth_key, data, hashlib.sha256).digest()[:nt]


def ctr_seal(key, nonce, aad, pt, nt):
    ct = _ctr(key[:16], nonce, pt)
    return ct + rfc9605_tag(key[16:], nonce, aad, ct, nt)


def ctr_open(key, nonce, aad, ct, nt):
    if len(ct) < nt:
        return None
    inner, tag = ct[:-nt], ct[-nt:]
    if not _hmac.compare_digest(rfc9605_tag(key[16:], nonce, aad, inner, nt), tag):
        return None
    return _ctr(key[:16], nonce, inner)


# The SFrame suites (RFC 9605 §4.5, §8.1), which secure objects §7.2 uses:
# id: (hash, Nk, Nt, counter mode)
SUITES = {1: ("sha256", 48, 10, True), 2: ("sha256", 48, 8, True), 3: ("sha256", 48, 4, True),
          4: ("sha256", 16, 16, False), 5: ("sha512", 32, 16, False)}


def seal(suite, key, nonce, aad, pt):
    h, nk, nt, ctr = SUITES[suite]
    return ctr_seal(key, nonce, aad, pt, nt) if ctr else gcm_seal(key, nonce, aad, pt)


def open_(suite, key, nonce, aad, ct):
    h, nk, nt, ctr = SUITES[suite]
    return ctr_open(key, nonce, aad, ct, nt) if ctr else gcm_open(key, nonce, aad, ct)


# --------------------------------------------------------------------- HKDF (RFC 5869)

def hkdf_extract(hname, salt, ikm):
    h = getattr(hashlib, hname)
    return _hmac.new(salt or bytes(h().digest_size), ikm, h).digest()


def hkdf_expand(hname, prk, info, n):
    h = getattr(hashlib, hname)
    out, t, i = b"", b"", 1
    while len(out) < n:
        t = _hmac.new(prk, t + info + bytes([i]), h).digest()
        out += t
        i += 1
    return out[:n]


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


# ============================================================ MOQT -22 integers and pairs

def vi(v):
    """MOQT §8.1, the shortest encoding."""
    if v < 0 or v >= 1 << 64:
        raise ValueError("vi64 range")
    for n in range(1, 9):
        bits = 7 * n if n < 8 else 56
        if v < (1 << bits):
            prefix = (0xFF << (9 - n)) & 0xFF if n > 1 else 0
            return (v | (prefix << (8 * (n - 1)))).to_bytes(n, "big")
    return b"\xff" + v.to_bytes(8, "big")


def vi_read(b, i):
    """Any encoding §8.1 allows; ValueError when the bytes end inside it."""
    if i >= len(b):
        raise ValueError("truncated")
    first = b[i]
    n = 9
    for k in range(8):
        if not (first >> (7 - k)) & 1:
            n = k + 1
            break
    if i + n > len(b):
        raise ValueError("truncated")
    if n == 9:
        return int.from_bytes(b[i + 1:i + 9], "big"), i + 9
    v = first & (0xFF >> n) if n > 1 else first & 0x7F
    for x in b[i + 1:i + n]:
        v = (v << 8) | x
    return v, i + n


def pairs_read(b):
    """MOQT §8.3: delta-coded types; an even type's value one integer, an odd type's a
    length of at most 2^16 - 1 and that many bytes."""
    out, i, prev = [], 0, 0
    while i < len(b):
        d, i = vi_read(b, i)
        ty = prev + d
        if ty >= 1 << 64:
            raise ValueError("type overflow")
        prev = ty
        if ty % 2 == 0:
            v, i = vi_read(b, i)
        else:
            n, i = vi_read(b, i)
            if n > 0xFFFF or i + n > len(b):
                raise ValueError("value")
            v, i = bytes(b[i:i + n]), i + n
        out.append((ty, v))
    return out


def pairs_write(ps):
    """In non-decreasing type order (pairs of one type keep their order), shortest
    integers."""
    out, prev = b"", 0
    for ty, v in sorted(ps, key=lambda p: p[0]):
        out += vi(ty - prev)
        prev = ty
        out += vi(v) if ty % 2 == 0 else vi(len(v)) + v
    return out


# ===================================================================== secure objects

KEY_ID = 0x3A4A
IMMUTABLE = 0x0B
ENCRYPTED_LIST = 0x000A
PADDING = 0x0032
INNER_DIGEST = 0x3A4B


def ser_namespace(fields):
    out = vi(len(fields))
    for f in fields:
        out += vi(len(f)) + f
    return out


def ser_ftn(fields, name):
    return ser_namespace(fields) + vi(len(name)) + name


def derive_key_salt(suite, kid, base, ftn):
    """Secure objects §3.7 (confidentiality §4.5)."""
    h, nk, nt, _ = SUITES[suite]
    secret = hkdf_extract(h, b"", base)
    tail = ftn + suite.to_bytes(2, "big") + kid.to_bytes(8, "big")
    return (hkdf_expand(h, secret, b"MOQ 1.0 Secure Objects Secret key " + tail, nk),
            hkdf_expand(h, secret, b"MOQ 1.0 Secret salt " + tail, 12))


def ctr_of(group, obj):
    if obj > 0xFFFFFFFF:
        raise Refuse("object_id_too_large")
    return group.to_bytes(8, "big") + obj.to_bytes(4, "big")


def aad_of(group, obj, priority, sip):
    return ctr_of(group, obj) + bytes([priority]) + sip


def xor(a, b):
    return bytes(x ^ y for x, y in zip(a, b))


def track_base_key(secret, fields):
    """Confidentiality §5.7: HKDF-Expand with SHA-256, info mocha_track and the namespace."""
    return hkdf_expand("sha256", secret, b"mocha_track" + ser_namespace(fields), 32)


def kid_of(epoch, leaf):
    if epoch >= 1 << 40 or leaf >= 1 << 24:
        raise ValueError("KID range")
    return (epoch << 24) | leaf


def plaintext_write(payload, digest, entries, padding):
    pt = vi(len(payload)) + payload
    es = b""
    if digest:
        es += INNER_DIGEST.to_bytes(2, "big") + vi(32) + blake3(payload)
    for ty, v in entries:
        es += ty.to_bytes(2, "big") + vi(len(v)) + v
    if padding is not None:
        es += PADDING.to_bytes(2, "big") + padding.to_bytes(2, "big") + bytes(padding)
    if es:
        pt += ENCRYPTED_LIST.to_bytes(2, "big") + vi(len(es)) + es
    return pt


def plaintext_read(b):
    """Confidentiality §4.4: (payload, inner digest or None)."""
    try:
        n, i = vi_read(b, 0)
        if i + n > len(b):
            raise Refuse("plaintext_malformed")
        payload, i = b[i:i + n], i + n
        if i == len(b):
            return payload, None
        if i + 2 > len(b) or int.from_bytes(b[i:i + 2], "big") != ENCRYPTED_LIST:
            raise Refuse("plaintext_malformed")
        m, i = vi_read(b, i + 2)
        if m != len(b) - i:
            raise Refuse("plaintext_malformed")
        digest = None
        while i < len(b):
            if i + 2 > len(b):
                raise Refuse("plaintext_malformed")
            ty = int.from_bytes(b[i:i + 2], "big")
            i += 2
            if ty == PADDING:
                if i + 2 > len(b):
                    raise Refuse("plaintext_malformed")
                n = int.from_bytes(b[i:i + 2], "big")
                if i + 2 + n != len(b):
                    raise Refuse("plaintext_malformed")
                break
            n, i = vi_read(b, i)
            if i + n > len(b):
                raise Refuse("plaintext_malformed")
            v, i = b[i:i + n], i + n
            if ty == INNER_DIGEST:
                if digest is not None or len(v) != 32:
                    raise Refuse("plaintext_malformed")
                digest = v
        return payload, digest
    except ValueError:
        raise Refuse("plaintext_malformed")


def key_id_of(props_bytes):
    """Confidentiality §6 step 1: (KID, Serialized Immutable Properties)."""
    try:
        outer = pairs_read(props_bytes)
    except ValueError:
        raise Refuse("properties_malformed")
    imm = [v for t, v in outer if t == IMMUTABLE]
    if any(t == KEY_ID for t, _ in outer) or len(imm) > 1:
        raise Refuse("properties_malformed")
    if not imm:
        raise Refuse("key_id_missing")
    sip = imm[0]
    try:
        inner = pairs_read(sip)
    except ValueError:
        raise Refuse("properties_malformed")
    if any(t == IMMUTABLE for t, _ in inner):
        raise Refuse("properties_malformed")
    kids = [v for t, v in inner if t == KEY_ID]
    if not kids:
        raise Refuse("key_id_missing")
    if len(kids) > 1:
        raise Refuse("properties_malformed")
    return kids[0], sip


# ===================================================================== a receiver (§6)

def owner_of(fields):
    """Live-rooms §6.1: the namespace's second field, when it is a principal's 64 hex digits."""
    if len(fields) < 2:
        return None
    f = fields[1]
    if len(f) == 64 and all(c in b"0123456789abcdef" for c in f):
        return bytes.fromhex(f.decode())
    return None


class Ring:
    """What a member holds: the epochs it entered, their secrets and rosters, the
    locations it opened, and its removal."""

    def __init__(self, suite, me):
        self.suite, self.me = suite, me
        self.joined = self.removed = None
        self.secrets, self.rosters, self.opened = {}, {}, set()

    def enter(self, epoch, secret, roster):
        if self.joined is None:
            self.joined = epoch
        self.secrets[epoch] = secret
        self.rosters[epoch] = dict(roster)

    def open(self, o):
        fields = [f.encode() for f in o["namespace"]]
        name = o["track"].encode()
        group, obj, prio = o["group"], o["object"], o["priority"]
        kid, sip = key_id_of(bytes.fromhex(o["properties"]))
        if obj > 0xFFFFFFFF:
            raise Refuse("object_id_too_large")
        epoch, leaf = kid >> 24, kid & 0xFFFFFF
        if self.removed is not None and epoch >= self.removed:
            raise Refuse("removed_before_epoch")
        if not self.rosters or self.joined is None or epoch < self.joined:
            raise Refuse("epoch_not_joined")
        latest = max(self.rosters)
        if epoch > latest:
            return {"outcome": "pending", "epoch": epoch}
        if epoch not in self.secrets or epoch not in self.rosters:
            raise Refuse("epoch_not_joined")
        sender = self.rosters[epoch].get(leaf)
        if sender is None:
            raise Refuse("key_id_unknown")
        if owner_of(fields) != sender:
            raise Refuse("sender_not_owner")
        for e in range(epoch, latest + 1):
            if e in self.rosters and self.rosters[e].get(leaf) != sender:
                raise Refuse("sender_removed")
        ftn = ser_ftn(fields, name)
        place = (ftn, group, obj)
        if place in self.opened:
            raise Refuse("replayed")
        key, salt = derive_key_salt(self.suite, kid, track_base_key(self.secrets[epoch], fields), ftn)
        pt = open_(self.suite, key, xor(salt, ctr_of(group, obj)), aad_of(group, obj, prio, sip), bytes.fromhex(o["payload"]))
        if pt is None:
            raise Refuse("bad_tag")
        payload, digest = plaintext_read(pt)
        if digest is not None and digest != blake3(payload):
            raise Refuse("inner_digest_mismatch")
        self.opened.add(place)
        return {"outcome": "opened", "payload_blake3": blake3(payload).hex()}


def try_open(ring, o):
    try:
        return ring.open(o)
    except Refuse as e:
        return {"outcome": e.code}


# ===================================================================== the record (§3)

def check_record(b, room, authority, mode):
    if len(b) < 4:
        raise Refuse("truncated")
    if b[:4] != b"WRE1":
        raise Refuse("malformed")
    if len(b) < 5:
        raise Refuse("truncated")
    if b[4] != 1:
        raise Refuse("unknown_version")
    if len(b) < 105:
        raise Refuse("truncated")
    if len(b) > 105:
        raise Refuse("malformed")
    if b[5:37] != room:
        raise Refuse("wrong_room")
    suite = int.from_bytes(b[37:39], "little")
    if suite not in (4, 5):
        raise Refuse("suite_unsupported")
    if int.from_bytes(b[39:41], "little") != 1:
        raise Refuse("mls_suite_unsupported")
    if mode != 0:
        raise Refuse("recording_refused")
    r = ED.verify(authority, b"wai:live-room-e2ee\x01" + b[:41], b[41:105])
    if r is not None:
        raise Refuse(r)
    return suite


# ============================================================ the group follows the room

class Model:
    """Confidentiality §5.4 to §5.6, as state."""

    def __init__(self, room, clock):
        self.room, self.clock = room, clock
        self.generation, self.group = 0, None  # group: [epoch, {leaf: principal}]
        self.seq, self.keyed, self.gone, self.pending = 0, {}, {}, []

    def group_id(self):
        return self.room + self.generation.to_bytes(4, "big")

    def event(self, ev):
        self.seq += 1
        s = bytes.fromhex(ev["subject"])
        kind = ev["event"]
        if kind == "join":
            now = True if ev["keyed"] else None
        elif kind == "leave":
            now = False
        elif kind == "grant_exchange":
            now = ev["keyed_after"] if ev["keyed_before"] != ev["keyed_after"] else None
        else:
            raise ValueError("an event this text does not define")
        if now is True:
            self.keyed[s] = self.seq
            self.gone.pop(s, None)
        elif now is False and self.keyed.pop(s, None) is not None:
            self.gone[s] = self.seq
        if self.group is not None and not any(p in self.keyed for p in self.group[1].values()):
            self.group, self.pending = None, []
            self.generation += 1
            self.gone = {}

    def reflected(self):
        if self.group is None:
            return set()
        r = set(self.group[1].values())
        for p in self.pending:
            if p["from"] != "clock":
                continue
            if "add" in p["change"]:
                r.add(bytes.fromhex(p["change"]["add"]["identity"]))
            else:
                x = self.group[1].get(p["change"]["remove"]["leaf"])
                if x is not None:
                    r.discard(x)
        return r

    def next_owed(self):
        """The earliest change, in the order of the events calling for it, that can be
        proposed now."""
        if self.group is None:
            if not self.keyed:
                return None
            p = min(self.keyed, key=lambda q: (self.keyed[q], q))
            return ("create", p)
        refl = self.reflected()
        owed = []
        for p in sorted(self.keyed):
            if p not in refl:
                owed.append((self.keyed[p], 1, ("add", p)))
        for p in sorted(refl):
            if p not in self.keyed:
                leaf = next((l for l, q in self.group[1].items() if q == p), None)
                if leaf is not None:
                    owed.append((self.gone.get(p, 0), 0, ("remove", p, leaf)))
        owed.sort(key=lambda o: (o[0], o[1]))
        return owed[0][2] if owed else None

    def create(self, by, gid, ext):
        if gid != self.group_id() or ext != [self.clock]:
            raise Refuse("group_mismatch")
        if self.next_owed() != ("create", by):
            raise Refuse("proposal_mismatch")
        self.group = [0, {0: by}]

    def propose(self, p):
        if self.group is None:
            raise Refuse("proposal_mismatch")
        if p["from"] == "clock":
            if p["epoch"] != self.group[0]:
                raise Refuse("proposal_mismatch")
            ch, owed = p["change"], self.next_owed()
            if "add" in ch:
                if ch["add"]["identity"] != ch["add"]["signature_key"]:
                    raise Refuse("credential_mismatch")
                ok = owed is not None and owed[0] == "add" and owed[1] == bytes.fromhex(ch["add"]["identity"])
            else:
                ok = owed is not None and owed[0] == "remove" and owed[2] == ch["remove"]["leaf"]
            if not ok:
                raise Refuse("proposal_mismatch")
        self.pending.append(p)

    def commit(self, sender, holds):
        """§5.5's checks, in its order, then §5.6: the removed and the added."""
        if self.group is None:
            raise Refuse("commit_mismatch")
        if any(not isinstance(i, int) or i < 0 or i >= len(self.pending) for i in holds):
            raise Refuse("commit_mismatch")
        held = [self.pending[i] for i in holds]
        removed = {p["change"]["remove"]["leaf"] for p in held if "remove" in p["change"]}
        epoch, roster = self.group
        committer = next((l for l in sorted(roster) if l not in removed), None)
        if committer != sender:
            raise Refuse("commit_not_committer")
        if any(p["from"] != "clock" for p in held):
            raise Refuse("proposal_not_clock")
        clock = [i for i, p in enumerate(self.pending) if p["from"] == "clock"]
        if sorted(holds) != clock[:len(holds)] or any(l not in roster for l in removed):
            raise Refuse("commit_mismatch")
        new = {l: p for l, p in roster.items() if l not in removed}
        added = []
        for p in held:
            if "add" in p["change"]:
                q = bytes.fromhex(p["change"]["add"]["identity"])
                leaf = 0
                while leaf in new:
                    leaf += 1
                new[leaf] = q
                added.append(q)
        gone = [roster[l] for l in sorted(removed)]
        for q in gone:
            self.gone.pop(q, None)
        self.group, self.pending = [epoch + 1, new], []
        return gone


def roster_json(r):
    return {str(l): p.hex() for l, p in sorted(r.items())}


def replay(v):
    suite = v["suite"]
    m = Model(bytes.fromhex(v["room"]), bytes.fromhex(v["clock"]))
    rings, rosters, deliveries = {}, [], []

    def enter_all(secret):
        epoch, roster = m.group
        rosters.append({"generation": m.generation, "epoch": epoch, "roster": roster_json(roster)})
        for p in roster.values():
            rings.setdefault(p, Ring(suite, p)).enter(epoch, secret, roster)

    outcome = "accept"
    for i, s in enumerate(v["steps"]):
        try:
            if "event" in s:
                m.event(s["event"])
            elif "create" in s:
                c = s["create"]
                m.create(bytes.fromhex(c["by"]), bytes.fromhex(c["group_id"]), [bytes.fromhex(x) for x in c["external_senders"]])
                enter_all(bytes.fromhex(c["application_secret"]))
            elif "proposal" in s:
                m.propose(s["proposal"])
            elif "commit" in s:
                c = s["commit"]
                gone = m.commit(c["sender"], c["holds"])
                for q in gone:
                    if q in rings:
                        rings[q].removed = m.group[0]
                enter_all(bytes.fromhex(c["application_secret"]))
            elif "object" in s:
                out = []
                for d in s["deliveries"]:
                    who = bytes.fromhex(d["receiver"])
                    r = try_open(rings.setdefault(who, Ring(suite, who)), s["object"])
                    r["receiver"] = who.hex()
                    out.append(r)
                deliveries.append({"step": i, "deliveries": out})
            else:
                raise ValueError("a step this corpus does not define")
        except Refuse as e:
            outcome = {"refuse": e.code, "at": i}
            break
    return {"outcome": outcome, "rosters": rosters, "deliveries": deliveries}


# =========================================================================== the checks

def check_object(v):
    suite = v["suite"]
    if "input" in v:
        i = v["input"]
        fields = [f.encode() for f in i["namespace"]]
        name = i["track"].encode()
        payload = bytes.fromhex(i["payload"])
        kid = kid_of(i["epoch"], i["leaf"])
        base = track_base_key(bytes.fromhex(i["application_secret"]), fields)
        ftn = ser_ftn(fields, name)
        key, salt = derive_key_salt(suite, kid, base, ftn)
        imm = [(p["type"], p["varint"] if "varint" in p else bytes.fromhex(p["bytes"])) for p in i["immutable"]]
        sip = pairs_write([p for p in imm if p[0] != KEY_ID] + [(KEY_ID, kid)])
        pt = plaintext_write(payload, i["inner_digest"], [(e["type"], bytes.fromhex(e["value"])) for e in i["entries"]], i["padding"])
        aad = aad_of(i["group"], i["object"], i["priority"], sip)
        nonce = xor(salt, ctr_of(i["group"], i["object"]))
        sealed = seal(suite, key, nonce, aad, pt)
        derived = {"track_base_key": base.hex(), "kid": kid, "serialized_full_track_name": ftn.hex(), "moq_key": key.hex(),
                   "moq_salt": salt.hex(), "immutable_properties": sip.hex(), "plaintext": pt.hex(), "aad": aad.hex(), "nonce": nonce.hex()}
        if derived != v["derived"]:
            bad = [k for k in derived if derived[k] != v["derived"].get(k)]
            raise AssertionError("derived values differ: %s" % bad)
        o = v["object"]
        if (o["properties"] != pairs_write([(IMMUTABLE, sip)]).hex() or o["payload"] != sealed.hex()
                or [f.encode() for f in o["namespace"]] != fields or o["track"] != i["track"]
                or (o["group"], o["object"], o["priority"]) != (i["group"], i["object"], i["priority"])):
            raise AssertionError("the sealed object is not what its inputs seal")
        if v["content_hash"] != blake3(sealed).hex():
            raise AssertionError("the content hash is not BLAKE3 of the sealed payload")
    ctx = v["context"]
    ring = Ring(suite, bytes.fromhex(ctx["receiver"]))
    for e in ctx["epochs"]:
        ring.enter(e["epoch"], bytes.fromhex(e["application_secret"]), {int(k): bytes.fromhex(p) for k, p in e["roster"].items()})
    return try_open(ring, v["object"])


def check_published(r, so):
    """RFC 9605 Appendix C (`r`) and secure objects' Appendix B (`so`)."""
    fails = []

    def header(kid, ctr):
        def compact(x):
            return x.to_bytes(max(1, (x.bit_length() + 7) // 8), "big")
        b, tail = 0, b""
        if kid < 8:
            b |= kid << 4
        else:
            k = compact(kid)
            b |= 0x80 | ((len(k) - 1) << 4)
            tail += k
        if ctr < 8:
            b |= ctr
        else:
            c = compact(ctr)
            b |= 0x08 | (len(c) - 1)
            tail += c
        return bytes([b]) + tail

    X = bytes.fromhex
    for c in r["header"]:
        if header(c["kid"], c["ctr"]).hex() != c["header"]:
            fails.append("RFC 9605 header %x/%x" % (c["kid"], c["ctr"]))
    for c in r["aead"]:
        s = c["cipher_suite"]
        if X(c["key"])[:16] != X(c["enc_key"]) or X(c["key"])[16:] != X(c["auth_key"]):
            fails.append("RFC 9605 C.2 subkeys, suite %d" % s)
        ct = seal(s, X(c["key"]), X(c["nonce"]), X(c["aad"]), X(c["pt"]))
        if ct.hex() != c["ct"] or open_(s, X(c["key"]), X(c["nonce"]), X(c["aad"]), ct) != X(c["pt"]):
            fails.append("RFC 9605 C.2 suite %d" % s)
    for c in r["sframe"]:
        s = c["cipher_suite"]
        h, nk, nt, _ = SUITES[s]
        secret = hkdf_extract(h, b"", X(c["base_key"]))
        kl = b"SFrame 1.0 Secret key " + c["kid"].to_bytes(8, "big") + s.to_bytes(2, "big")
        sl = b"SFrame 1.0 Secret salt " + c["kid"].to_bytes(8, "big") + s.to_bytes(2, "big")
        key, salt = hkdf_expand(h, secret, kl, nk), hkdf_expand(h, secret, sl, 12)
        nonce = xor(salt, c["ctr"].to_bytes(12, "big"))
        hdr = header(c["kid"], c["ctr"])
        aad = hdr + X(c["metadata"])
        ok = (secret.hex() == c["sframe_secret"] and kl.hex() == c["sframe_key_label"] and sl.hex() == c["sframe_salt_label"]
              and key.hex() == c["sframe_key"] and salt.hex() == c["sframe_salt"] and nonce.hex() == c["nonce"] and aad.hex() == c["aad"])
        ct = hdr + seal(s, key, nonce, aad, X(c["pt"]))
        if not ok or ct.hex() != c["ct"] or open_(s, key, nonce, aad, ct[len(hdr):]) != X(c["pt"]):
            fails.append("RFC 9605 C.3 suite %d" % s)
    if not (len(r["header"]) > 100 and len(r["aead"]) == 3 and len(r["sframe"]) == 5):
        fails.append("RFC 9605's vectors are not all there")

    for c in so["full_track_name"]:
        if ser_ftn([t.encode() for t in c["namespace_tuples"]], c["track_name"].encode()).hex() != c["serialized_full_track_name"]:
            fails.append("secure objects B.1 %s" % c["track_name"])
    for c in so["key_derivation"]:
        s = c["cipher_suite"]
        h, nk, nt, _ = SUITES[s]
        ftn = X(c["serialized_full_track_name"])
        tail = ftn + s.to_bytes(2, "big") + c["key_id"].to_bytes(8, "big")
        key, salt = derive_key_salt(s, c["key_id"], X(c["track_base_key"]), ftn)
        if (hkdf_extract(h, b"", X(c["track_base_key"])).hex() != c["moq_secret"]
                or (b"MOQ 1.0 Secure Objects Secret key " + tail).hex() != c["moq_key_label"]
                or (b"MOQ 1.0 Secret salt " + tail).hex() != c["moq_salt_label"]
                or key.hex() != c["moq_key"] or salt.hex() != c["moq_salt"]):
            fails.append("secure objects B.2 suite %d" % s)
    for c in so["nonce_aad"]:
        ctr = c["group_id"].to_bytes(8, "big") + c["object_id"].to_bytes(4, "big")
        # Appendix B's Key ID property: a 16-bit type 0x0002, a length, the value (§11.3).
        sip = b"\x00\x02" + vi(len(vi(c["key_id"]))) + vi(c["key_id"]) + X(c["serialized_immutable_properties"])
        if (ctr.hex() != c["ctr"] or xor(X(c["moq_salt"]), ctr).hex() != c["nonce"]
                or (ctr + bytes([c["publisher_priority"]]) + sip).hex() != c["aad"]):
            fails.append("secure objects B.3 group %d" % c["group_id"])
    gcm = ctr_suites = 0
    for c in so["encryption"]:
        s = c["cipher_suite"]
        h, nk, nt, is_ctr = SUITES[s]
        payload = X(c["original_payload"])
        pt = vi(len(payload)) + payload + X(c["encrypted_properties_list"])
        key, nonce, aad, want = X(c["moq_key"]), X(c["nonce"]), X(c["aad"]), X(c["ciphertext"])
        k2, s2 = derive_key_salt(s, c["key_id"], X(c["track_base_key"]), ser_ftn([t.encode() for t in c["namespace_tuples"]], c["track_name"].encode()))
        if pt.hex() != c["plaintext"] or k2 != key or xor(s2, ctr_of(c["group_id"], c["object_id"])) != nonce:
            fails.append("secure objects B.4 inputs, suite %d" % s)
            continue
        got = seal(s, key, nonce, aad, pt)
        if not is_ctr:
            if got != want or open_(s, key, nonce, aad, want) != pt:
                fails.append("secure objects B.4 suite %d" % s)
            gcm += 1
        else:
            # §11.2: the counter-mode encryption is RFC 9605's; the vector's tag is
            # HMAC(auth_key, aad ‖ nonce ‖ ct), and RFC 9605 §4.5.1's is not that.
            body = len(want) - nt
            alt = _hmac.new(key[16:], aad + nonce + want[:body], hashlib.sha256).digest()[:nt]
            if got[:body] != want[:body] or got[body:] == want[body:] or alt != want[body:]:
                fails.append("secure objects B.4 suite %d: the difference of §11.2 is not as stated" % s)
            ctr_suites += 1
    if (gcm, ctr_suites) != (2, 2):
        fails.append("secure objects' encryption vectors are not all there")
    return fails


def spec_codes():
    """The codes §13 of the extension lists, in order."""
    text = open(os.path.join(HERE, "..", "extensions", "confidentiality.md"), encoding="utf-8").read()
    start = text.index("## 13. Reasons")
    end = text.find("\n## ", start + 1)
    out = []
    for line in text[start:end].splitlines():
        if line.startswith("| `"):
            out += line.split("|")[1].split("`")[1::2]
    return out


def check(doc, quiet=False):
    fails, seen = [], set()

    def note(o):
        if isinstance(o, str):
            seen.add(o)
        elif isinstance(o, dict) and "refuse" in o:
            seen.add(o["refuse"])

    def report(kind, name, bad, what):
        if not quiet:
            print(("PASS" if not bad else "FAIL") + " %s %s: %s" % (kind, name, what))
        fails.extend("%s %s: %s" % (kind, name, b) for b in bad)

    for v in doc["records"]:
        try:
            got = {"outcome": "accept", "suite": check_record(bytes.fromhex(v["record"]), bytes.fromhex(v["room"]),
                                                              bytes.fromhex(v["authority"]), v["recording_mode"])}
        except Refuse as e:
            got = {"outcome": e.code}
        bad = [k for k in ("outcome", "suite") if got.get(k) != v.get(k)]
        report("record", v["name"], bad, got["outcome"])
        note(v["outcome"])
    for v in doc["key_schedule"]:
        fields = [f.encode() for f in v["namespace"]]
        base = track_base_key(bytes.fromhex(v["application_secret"]), fields)
        kid = kid_of(v["epoch"], v["leaf"])
        ftn = ser_ftn(fields, v["track"].encode())
        key, salt = derive_key_salt(v["suite"], kid, base, ftn)
        got = {"track_base_key": base.hex(), "kid": kid, "serialized_full_track_name": ftn.hex(), "moq_key": key.hex(), "moq_salt": salt.hex()}
        report("key schedule", "epoch %d leaf %d" % (v["epoch"], v["leaf"]), [k for k in got if got[k] != v.get(k)], "derived")
    for v in doc["objects"]:
        try:
            got = check_object(v)
            bad = [k for k in ("outcome", "payload_blake3") if got.get(k) != v.get(k)]
        except (AssertionError, KeyError, ValueError, TypeError) as e:
            got, bad = {"outcome": "?"}, ["%s: %s" % (type(e).__name__, e)]
        report("object", v["name"], bad, got["outcome"])
        note(v["outcome"])
    for v in doc["rooms"]:
        try:
            got = replay(v)
            bad = [k for k in ("outcome", "rosters", "deliveries") if got[k] != v.get(k)]
        except (KeyError, ValueError, TypeError, AttributeError) as e:
            got, bad = {"outcome": "?"}, ["%s: %s" % (type(e).__name__, e)]
        report("room", v["name"], bad, json.dumps(got["outcome"]))
        note(v["outcome"])
        for d in v.get("deliveries", []):
            for x in d["deliveries"]:
                note(x["outcome"])
    codes = spec_codes()
    if doc["codes"] != codes:
        fails.append("the corpus's codes are not the ones §13 lists")
    for c in codes:
        if c not in seen:
            fails.append("no vector raises %s" % c)
    return fails


def flip(h, at):
    b = bytearray(bytes.fromhex(h))
    b[at] ^= 0x01
    return b.hex()


def mutations(doc):
    """Copies of the corpus, each with one defect a conforming verifier must catch."""
    def mutate(label, f):
        d = copy.deepcopy(doc)
        f(d)
        return label, d

    room = doc["rooms"][0]
    first_commit = next(i for i, s in enumerate(room["steps"]) if "commit" in s)
    first_object = next(i for i, s in enumerate(room["steps"]) if "object" in s)
    return [
        mutate("a record's signature byte", lambda d: d["records"][0].update(record=flip(d["records"][0]["record"], 60))),
        mutate("a refused record stated accepted", lambda d: d["records"][-1].update(outcome="accept", suite=4)),
        mutate("a track base key", lambda d: d["key_schedule"][2].update(track_base_key=flip(d["key_schedule"][2]["track_base_key"], 0))),
        mutate("a KID", lambda d: d["key_schedule"][4].update(kid=d["key_schedule"][4]["kid"] + 1)),
        mutate("an object's nonce", lambda d: d["objects"][0]["derived"].update(nonce=flip(d["objects"][0]["derived"]["nonce"], 11))),
        mutate("an object's AAD", lambda d: d["objects"][3]["derived"].update(aad=flip(d["objects"][3]["derived"]["aad"], 12))),
        mutate("a sealed payload's byte", lambda d: d["objects"][1]["object"].update(payload=flip(d["objects"][1]["object"]["payload"], 5))),
        mutate("an object's content hash", lambda d: d["objects"][2].update(content_hash=flip(d["objects"][2]["content_hash"], 0))),
        mutate("a refused object stated opened", lambda d: d["objects"][-1].update(outcome="opened")),
        mutate("an application secret", lambda d: d["rooms"][0]["steps"][first_commit]["commit"].update(
            application_secret=flip(d["rooms"][0]["steps"][first_commit]["commit"]["application_secret"], 0))),
        mutate("a commit's sender", lambda d: d["rooms"][0]["steps"][first_commit]["commit"].update(sender=1)),
        mutate("a delivery's outcome", lambda d: d["rooms"][0]["deliveries"][0]["deliveries"][0].update(outcome="bad_tag")),
        mutate("a roster", lambda d: d["rooms"][0]["rosters"][2]["roster"].update({"2": d["rooms"][0]["rosters"][2]["roster"]["0"]})),
        mutate("a refused room's reason", lambda d: d["rooms"][1]["outcome"].update(refuse="commit_mismatch")),
        mutate("a refused room's step", lambda d: d["rooms"][1]["outcome"].update(at=d["rooms"][1]["outcome"]["at"] - 1)),
        mutate("an object's group in the room", lambda d: d["rooms"][0]["steps"][first_object]["object"].update(group=9)),
        mutate("a lifecycle event's subject", lambda d: d["rooms"][0]["steps"][0]["event"].update(subject="00" * 32)),
        mutate("a code dropped from the list", lambda d: d["codes"].pop()),
    ]


def published_mutations(r, so):
    """Copies of the published vectors, each with one defect."""
    def mutate(label, which, f):
        a, b = copy.deepcopy(r), copy.deepcopy(so)
        f(a if which == "rfc" else b)
        return label, a, b

    return [
        mutate("a byte of an RFC 9605 SFrame ciphertext", "rfc", lambda j: j["sframe"][4].update(ct=flip(j["sframe"][4]["ct"], -1))),
        mutate("an RFC 9605 header", "rfc", lambda j: j["header"][100].update(header="00")),
        mutate("a byte of an RFC 9605 counter-mode ciphertext", "rfc", lambda j: j["aead"][1].update(ct=flip(j["aead"][1]["ct"], 0))),
        mutate("a byte of a secure-objects GCM ciphertext", "so", lambda j: j["encryption"][2].update(ciphertext=flip(j["encryption"][2]["ciphertext"], 0))),
        mutate("a byte of a secure-objects counter-mode tag", "so", lambda j: j["encryption"][0].update(ciphertext=flip(j["encryption"][0]["ciphertext"], -1))),
        mutate("a secure-objects moq_key", "so", lambda j: j["key_derivation"][3].update(moq_key=flip(j["key_derivation"][3]["moq_key"], 0))),
        mutate("a secure-objects AAD", "so", lambda j: j["nonce_aad"][1].update(aad=flip(j["nonce_aad"][1]["aad"], 12))),
    ]


def main(argv):
    if len(argv) != 2:
        print("usage: e2ee_verify.py <live-conformance/e2ee>", file=sys.stderr)
        return 2
    d = argv[1]
    r = json.load(open(os.path.join(d, "rfc9605.json")))
    so = json.load(open(os.path.join(d, "secure-objects-01.json")))
    fails = check_published(r, so)
    print(("PASS" if not fails else "FAIL") + " the published vectors of RFC 9605 Appendix C and secure objects' Appendix B")
    pcaught = 0
    pmuts = published_mutations(r, so)
    for label, a, b in pmuts:
        try:
            ok = not check_published(a, b)
        except Exception:  # a crash on a mutated copy is a failure to read it
            ok = False
        if ok:
            fails.append("a mutated copy of the published vectors verifies: %s" % label)
        else:
            pcaught += 1
    print("%d of %d mutated copies of the published vectors fail, as they must" % (pcaught, len(pmuts)))
    doc = json.load(open(os.path.join(d, "vectors.json")))
    if not (doc["records"] and doc["objects"] and doc["rooms"] and doc["key_schedule"]):
        fails.append("an empty corpus")
    fails += check(doc)
    caught = 0
    muts = mutations(doc)
    for label, m in muts:
        try:
            ok = not check(m, quiet=True)
        except Exception:  # a crash on a mutated copy is a failure to read it
            ok = False
        if ok:
            fails.append("a mutated copy verifies: %s" % label)
        else:
            caught += 1
    print("%d of %d mutated copies of the corpus fail, as they must" % (caught, len(muts)))
    for f in fails:
        print("FAIL " + f)
    print("e2ee_verify: %s" % ("every vector verifies" if not fails else "%d failure(s)" % len(fails)))
    return 1 if fails else 0


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