#!/usr/bin/env python3
"""Independent verifier of the ad-presentation corpus (extensions/ad-presentation.md).

Written from the extension's text. It shares no code with wai-rs, which generates the corpus,
and never runs it. Each room, its policy, schedule and receipts are verified by the live-ads
corpus's own second implementation, live_ads_verify.py (live-ads §9), to which this verifier
gives the judge of the claims. Ed25519 signatures are checked with
../claim-conformance/_ed25519.py, WAI's strict rule (SPEC §7 *Signatures*). Everything else is
here, from the standard library:

- the session a session key names (§3) and the creative digest (§4);
- the claim, its spans and its rules (§5);
- the viewability rules, the qualifying run and the quartile (§6);
- the token, its challenge, and RSASSA-PSS-VERIFY with SHA-384 and a 48-byte salt over a DER
  SubjectPublicKeyInfo of id-RSASSA-PSS (§7, RFC 9577, RFC 9578 §6.4, RFC 8017 §8.1.2), checked
  first against the first test vector of RFC 9578 Appendix A.2;
- each claim's verdict, each impression's count and the presented spend (§9).

For every room it requires the outcome, reason and place, everything live-ads §9.4 and §9.4
here state, each claim's verdict and facts, and every HLS tag file byte for byte; of the corpus,
that every reason of §11 appears in it and that no other file is there. 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 ad_presentation_verify.py <live-conformance/presentation>

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

import copy
import hashlib
import json
import os
import struct
import sys

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

EXTENSION = os.path.join(HERE, "..", "extensions", "ad-presentation.md")

CLAIM_MAGIC = b"WPC1"
CLAIM_DOMAIN = b"wai:ad-presentation-claim\x01"                                   # §5.3
SESSION_DOMAIN = b"wai:ad-presentation-session\x01"                               # §3
ATTEST_DOMAIN = b"wai:ad-presentation-attest\x01"                                 # §7.2
CREATIVE_DOMAIN = b"wai:ad-creative\x01"                                          # §4
AUDIBLE, RESUMED = 1, 2                                                           # §5.2
CLAIM, ATTEST, AUDIBLE_RUN = 1, 2, 4                                              # §8
VIDEO, CROSS_MEDIA = 1, 2                                                         # §6.1
NK = {1: 48, 2: 256}                                                              # §7.1
PSS_OID = bytes.fromhex("2a864886f70d01010a")                                     # 1.2.840.113549.1.1.10
SALT = 48

# The checks a weakened verifier may drop (main(), the last part).
SKIP = set()


def sha(*parts):
    h = hashlib.sha256()
    for p in parts:
        h.update(p)
    return h.digest()


def reasons():
    """§11's codes, from the text."""
    with open(EXTENSION, encoding="utf-8") as f:
        text = f.read()
    at = text.index("\n## 11. Reasons")
    end = text.index("\n## ", at + 1)
    return [line.split("`")[1] for line in text[at:end].splitlines() if line.startswith("| `")]


REASONS = reasons()


class Bad(Exception):
    """A claim's verdict other than presented."""


# ---------------------------------------------------------------- §3, §4

def session_of(room, key):
    return sha(SESSION_DOMAIN, room, key)[:16]


def creative_digest(objects):
    out = [CREATIVE_DOMAIN, struct.pack("<I", len(objects))]
    for o in objects:
        out += [o["content_hash"], struct.pack("<I", o["duration_ms"])]
    return sha(*out)


def touched(objects, spans):
    """§9.2 step 6: the objects the spans touch, in order."""
    out, s = [], 0
    for o in objects:
        e = s + o["duration_ms"]
        if any(x["start"] < e and x["start"] + x["length"] > s for x in spans):
            out.append(o["content_hash"])
        s = e
    return out


# ---------------------------------------------------------------- §5 the claim

def read_claim(b):
    """§5.4: the claim, or Bad(claim_malformed)."""
    p = 0

    def take(n):
        nonlocal p
        if p + n > len(b):
            raise Bad("claim_malformed")
        p += n
        return b[p - n:p]

    if take(4) != CLAIM_MAGIC or take(1) != b"\x01":
        raise Bad("claim_malformed")
    c = {"room": take(32), "cue": take(32), "session": take(16), "slot": take(1)[0], "creative_hash": take(32),
         "player": take(32), "geometry": take(1)[0], "frames": struct.unpack("<I", take(4))[0]}
    n = struct.unpack("<H", take(2))[0]
    c["objects"] = [take(32) for _ in range(n)]
    m = struct.unpack("<H", take(2))[0]
    c["spans"] = []
    for _ in range(m):
        start, length, vis, flags = struct.unpack("<IIHB", take(11))
        c["spans"].append({"start": start, "length": length, "visible": vis, "flags": flags})
    c["attestation"] = take(struct.unpack("<H", take(2))[0])
    c["body"] = b[:p]
    c["sig"] = take(64)
    if p != len(b):
        raise Bad("claim_malformed")
    if "rules" in SKIP:
        return c
    if c["slot"] > 15 or c["geometry"] not in (1, 2, 3) or (n == 0) != (m == 0):
        raise Bad("claim_malformed")
    prev = None
    for s in c["spans"]:
        if s["length"] == 0 or s["visible"] > 10000 or s["flags"] > 3:
            raise Bad("claim_malformed")
        if prev is None and s["flags"] & RESUMED:
            raise Bad("claim_malformed")
        if prev is not None and (s["start"] < prev or (s["start"] > prev and not s["flags"] & RESUMED)):
            raise Bad("claim_malformed")
        prev = s["start"] + s["length"]
    return c


# ---------------------------------------------------------------- §6

def quartile(spans, d):
    """§6.3."""
    if not spans:
        return 0
    cov = 0
    if spans[0]["start"] == 0:
        for s in spans:
            if s["start"] != cov:
                break
            cov = s["start"] + s["length"]
    if cov >= d:
        return 5
    if 4 * cov >= 3 * d:
        return 4
    if 4 * cov >= 2 * d:
        return 3
    if 4 * cov >= d:
        return 2
    return 1


def runs(c, q):
    """§6.2: (the run, the viewable time)."""
    rule = q["rule"]
    admitted = (rule == VIDEO and c["geometry"] in (1, 3)) or rule == CROSS_MEDIA
    need = 5000 if rule == VIDEO else 10000
    best = cur = total = 0
    prev = None
    for s in c["spans"]:
        ok = admitted and s["visible"] >= need and (not q["require"] & AUDIBLE_RUN or s["flags"] & AUDIBLE)
        cont = prev == s["start"] and not s["flags"] & RESUMED
        if ok:
            total += s["length"]
            cur = cur + s["length"] if cont else s["length"]
            best = max(best, cur)
        else:
            cur = 0
        prev = s["start"] + s["length"]
    return best, total


# ---------------------------------------------------------------- §7 attestation

def redemption_context(c):
    return sha(ATTEST_DOMAIN, c["room"], c["cue"], c["session"], bytes([c["slot"]]), c["creative_hash"], c["player"])


def token_challenge(tt, issuer, ctx):
    """RFC 9577 §2.1.1, with no origin_info (§7.2)."""
    return struct.pack(">HH", tt, len(issuer)) + issuer + bytes([32]) + ctx + b"\x00\x00"


def der(b, i):
    """One DER element at i: (tag, contents, the index after it)."""
    if i + 2 > len(b):
        raise ValueError("short")
    tag, first = b[i], b[i + 1]
    if first < 0x80:
        n, h = first, 2
    else:
        k = first & 0x7F
        if not 1 <= k <= 2 or i + 2 + k > len(b):
            raise ValueError("length")
        n, h = int.from_bytes(b[i + 2:i + 2 + k], "big"), 2 + k
    if i + h + n > len(b):
        raise ValueError("short")
    return tag, b[i + h:i + h + n], i + h + n


def rsa_key(spki):
    """A DER SubjectPublicKeyInfo of id-RSASSA-PSS with a 2048-bit modulus (§7.3): (n, e)."""
    t, outer, end = der(spki, 0)
    if t != 0x30 or end != len(spki):
        raise ValueError("outer")
    ta, alg, i = der(outer, 0)
    tb, bits, j = der(outer, i)
    if ta != 0x30 or tb != 0x03 or j != len(outer):
        raise ValueError("spki")
    to, oid, _ = der(alg, 0)
    if to != 0x06 or oid != PSS_OID or bits[:1] != b"\x00":
        raise ValueError("algorithm")
    ts, seq, k = der(bits, 1)
    if ts != 0x30 or k != len(bits):
        raise ValueError("key")
    tn, n, a = der(seq, 0)
    te, e, z = der(seq, a)
    if tn != 0x02 or te != 0x02 or z != len(seq):
        raise ValueError("integers")
    if n[:1] == b"\x00":
        n = n[1:]
    if len(n) != 256 or n[0] < 0x80 or not e or e[0] >= 0x80:
        raise ValueError("size")
    return int.from_bytes(n, "big"), int.from_bytes(e, "big")


def mgf1(seed, length):
    out, c = b"", 0
    while len(out) < length:
        out += hashlib.sha384(seed + struct.pack(">I", c)).digest()
        c += 1
    return out[:length]


def pss_verify(spki, msg, sig):
    """RSASSA-PSS-VERIFY, RFC 8017 §8.1.2, SHA-384, MGF1-SHA-384, sLen 48 (RFC 9578 §6.4)."""
    try:
        n, e = rsa_key(spki)
    except ValueError:
        return False
    if len(sig) != 256:
        return False
    s = int.from_bytes(sig, "big")
    if s >= n:
        return False
    em = pow(s, e, n).to_bytes(256, "big")
    h_len = 48
    if em[-1] != 0xBC:
        return False
    db_len = 256 - h_len - 1
    masked, h = em[:db_len], em[db_len:-1]
    if masked[0] & 0x80:
        return False
    db = bytes(x ^ y for x, y in zip(masked, mgf1(h, db_len)))
    db = bytes([db[0] & 0x7F]) + db[1:]
    ps = db_len - SALT - 1
    if db[:ps] != bytes(ps) or db[ps] != 1:
        return False
    return hashlib.sha384(bytes(8) + hashlib.sha384(msg).digest() + db[ps + 1:]).digest() == h


def read_token(b):
    if len(b) < 2:
        return None
    tt = struct.unpack(">H", b[:2])[0]
    if tt not in NK or len(b) != 98 + NK[tt]:
        return None
    return {"type": tt, "digest": b[34:66], "key_id": b[66:98], "input": b[:98], "auth": b[98:]}


def attestation(c, q, issuers):
    """§9.2 step 9: None, or the reason."""
    if not c["attestation"]:
        return None if "attest_absent" in SKIP else "attestation_absent"
    t = read_token(c["attestation"])
    if t is None or t["type"] != q["token_type"]:
        return "attestation_invalid"
    if "attest_challenge" not in SKIP:
        if t["digest"] != sha(token_challenge(q["token_type"], q["issuer"], redemption_context(c))):
            return "attestation_invalid"
    if "attest_authenticator" in SKIP:
        return None
    for i in issuers:
        if i["name"] == q["issuer"] and i["token_type"] == t["type"]:
            for k in i["keys"]:
                if sha(k) == t["key_id"] and t["type"] == 2 and pss_verify(k, t["input"], t["auth"]):
                    return None
    return "attestation_invalid"


# ---------------------------------------------------------------- §9 the judge

class Judge:
    """The claims of one room: read once (steps 1, 2), judged as each receipt is reached."""

    def __init__(self, doc, v, policy):
        self.q = policy["presentation"]
        self.room_creatives = [{"digest": creative_digest(c), "objects": c} for c in doc["_creatives"]]
        self.issuers = doc["_issuers"]
        self.read, self.facts = [], []
        for h in v["claims"]:
            b = bytes.fromhex(h)
            fact = {"digest": None, "verdict": None, "receipt": None, "slot": None, "quartile": None, "run_ms": None,
                    "viewable_ms": None}
            try:
                c = read_claim(b)
                fact["digest"] = sha(c["body"]).hex()
                bad = None if "signature" in SKIP else ED.verify(c["player"], CLAIM_DOMAIN + c["body"], c["sig"])
                if bad:
                    raise Bad(bad)
                self.read.append(c)
                self.facts.append(None)
            except Bad as e:
                fact["verdict"] = str(e)
                self.read.append(None)
                self.facts.append(fact)

    def receipt(self, k, r):
        """Steps 3 to 9 for receipt k's impressions: (its presented spend, what is stated of it)."""
        counts, spend = [], 0
        for imp in r["impressions"]:
            own, verdict = None, "claim_absent"
            for j, c in enumerate(self.read):
                if c is None or self.facts[j] is not None:
                    continue
                if (c["room"], c["cue"], c["session"], c["slot"]) != (r["room"], r["cue"], r["session"], imp["slot"]):
                    continue
                f = {"digest": sha(c["body"]).hex(), "verdict": "presented", "receipt": k, "slot": c["slot"],
                     "quartile": None, "run_ms": None, "viewable_ms": None}
                if session_of(r["room"], c["player"]) != r["session"] and "session" not in SKIP:
                    f["verdict"] = "claim_key_unbound"
                elif own is not None and "duplicate" not in SKIP:
                    f["verdict"] = "claim_duplicate"
                else:
                    if own is None:
                        own = j
                    qt = quartile(c["spans"], imp["duration_ms"])
                    run, viewable = runs(c, self.q)
                    f.update(quartile=qt, run_ms=run, viewable_ms=viewable)
                    f["verdict"] = self.judge(c, imp, qt, run) or "presented"
                    if own == j:
                        verdict = f["verdict"]
                self.facts[j] = f
            if verdict == "presented":
                spend += imp["price_micros"]
            counts.append({"slot": imp["slot"], "claim": own, "verdict": verdict})
        return spend, {"presented_spend_micros": spend, "impressions": counts}

    def judge(self, c, imp, qt, run):
        """Steps 6 to 9."""
        if "asset" not in SKIP:
            cr = next((x for x in self.room_creatives if x["digest"] == imp["creative_hash"]), None)
            if (c["creative_hash"] != imp["creative_hash"] or cr is None
                    or sum(o["duration_ms"] for o in cr["objects"]) != imp["duration_ms"]
                    or any(s["start"] + s["length"] > imp["duration_ms"] for s in c["spans"])
                    or touched(cr["objects"], c["spans"]) != c["objects"]):
                return "claim_asset_mismatch"
        if qt != imp["quartiles_reached"] and "quartile" not in SKIP:
            return "claim_quartile_mismatch"
        if run < 2000 and "viewability" not in SKIP:
            return "claim_not_viewable"
        if self.q["require"] & ATTEST:
            return attestation(c, self.q, self.issuers)
        return None

    def finish(self):
        out = []
        for c, f in zip(self.read, self.facts):
            if f is None:
                f = {"digest": sha(c["body"]).hex(), "verdict": "claim_no_decision", "receipt": None, "slot": None,
                     "quartile": None, "run_ms": None, "viewable_ms": None}
            out.append(f)
        return out


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

def held(doc):
    """The creatives and issuers the verifier holds."""
    doc["_creatives"] = [[{"content_hash": bytes.fromhex(o["content_hash"]), "duration_ms": o["duration_ms"]}
                          for o in c["objects"]] for c in doc["creatives"]]
    doc["_issuers"] = [{"name": i["name"].encode(), "token_type": i["token_type"],
                        "keys": [bytes.fromhex(k) for k in i["keys"]]} for i in doc["issuers"]]


def rfc_vector(doc):
    """RFC 9578 Appendix A.2, its first test vector: its token verifies under its key."""
    r = doc["rfc9578_a2_vector1"]
    pki, ch, tok = bytes.fromhex(r["pkI"]), bytes.fromhex(r["token_challenge"]), bytes.fromhex(r["token"])
    t = read_token(tok)
    return (t is not None and t["type"] == 2 and t["digest"] == sha(ch) and t["key_id"] == sha(pki)
            and pss_verify(pki, t["input"], t["auth"]))


def check(doc, files, quiet=False):
    """Every room against its statement, and every tag file: the failures."""
    fails = []
    try:
        held(doc)
    except Exception as e:  # a mutated copy that does not read
        return ["the creatives or issuers do not read: %s" % e], {}
    if not rfc_vector(doc):
        fails.append("RFC 9578 A.2's first test vector does not verify")
    written = {}
    for v in doc.get("vectors", []):
        try:
            got, tags = LAV.ads(v, judge=lambda vv, d, p: Judge(doc, vv, p))
        except Exception as e:  # a crash on a mutated copy is a failure to read it
            got, tags = None, None
            bad = ["the vector does not verify: %s: %s" % (type(e).__name__, e)]
        else:
            bad = [k for k in got if v.get(k) != got[k]]
            bad += [k for k in ("outcome", "reason", "at", "receipt_facts", "claim_facts") if k in v and k not in got]
        if tags is not None:
            name = v["name"] + ".m3u8"
            written[name] = tags
            if files.get(name) != tags:
                bad.append("its tag file %s" % name)
        if not quiet:
            what = "?" if got is None else got["outcome"] + (" " + got["reason"] if got["outcome"] == "refuse" else "")
            verdicts = "" if got is None else " (%s)" % ", ".join(f["verdict"] for f in got.get("claim_facts", []))
            print(("PASS" if not bad else "FAIL") + " room %s: %s%s" % (v.get("name"), what, verdicts))
        fails += ["room %s: %s differs from the corpus" % (v.get("name"), k) for k in bad]
    return fails, written


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


def mutations(doc, files):
    """Copies of the corpus, each with one defect a conforming verifier must catch."""
    vn = {v["name"]: k for k, v in enumerate(doc["vectors"])}
    video, cross, under = vn["video_rule"], vn["cross_media_attested"], vn["ceiling_counts_presented"]
    over = vn["ceiling_exceeded_by_presented"]

    def mutate(label, f):
        d, fs = copy.deepcopy(doc), dict(files)
        f(d, fs)
        return label, d, fs

    def vs(d):
        return d["vectors"]

    def swap(d, k, a, b):
        c = vs(d)[k]["claims"]
        c[a], c[b] = c[b], c[a]

    return [
        mutate("a byte of a presented claim's signature",
               lambda d, fs: vs(d)[video]["claims"].__setitem__(0, flip(vs(d)[video]["claims"][0], -3))),
        mutate("a byte of a presented claim's span", lambda d, fs: vs(d)[video]["claims"].__setitem__(1, flip(vs(d)[video]["claims"][1], 300))),
        mutate("a receipt's stated presented spend",
               lambda d, fs: vs(d)[video]["receipt_facts"][0].update(presented_spend_micros=69999)),
        mutate("the room's stated spend", lambda d, fs: vs(d)[cross].update(spend_micros=vs(d)[cross]["spend_micros"] + 1)),
        mutate("the duplicate claim withheld", lambda d, fs: vs(d)[video]["claims"].pop(5)),
        mutate("the unbound claim and the mismatched one given in the other order", lambda d, fs: swap(d, video, 6, 7)),
        mutate("a claim's stated quartile", lambda d, fs: vs(d)[video]["claim_facts"][1].update(quartile=5)),
        mutate("a claim's stated run", lambda d, fs: vs(d)[video]["claim_facts"][3].update(run_ms=3000)),
        mutate("an impression stated as counting", lambda d, fs: vs(d)[video]["receipt_facts"][3]["impressions"][1].update(verdict="presented")),
        mutate("a creative's object one millisecond longer", lambda d, fs: d["creatives"][1]["objects"][0].update(duration_ms=5001)),
        mutate("a creative's object hash", lambda d, fs: d["creatives"][0]["objects"][2].update(content_hash="00" * 32)),
        mutate("the issuer's key withheld", lambda d, fs: d["issuers"][0].update(keys=[])),
        mutate("the issuer named otherwise", lambda d, fs: d["issuers"][0].update(name="issuer.example.")),
        mutate("a byte of the RFC's token", lambda d, fs: d["rfc9578_a2_vector1"].update(token=flip(d["rfc9578_a2_vector1"]["token"], 200))),
        mutate("the second presented claim given to the room within its ceiling",
               lambda d, fs: vs(d)[under]["claims"].append(vs(d)[over]["claims"][1])),
        mutate("the attested claim given without its token",
               lambda d, fs: vs(d)[cross]["claims"].__setitem__(0, vs(d)[cross]["claims"][1])),
        mutate("a tag file removed", lambda d, fs: fs.pop("video_rule.m3u8")),
    ]


def main(corpus_dir):
    fails = []
    st = ED._self_test()
    print(("PASS" if not st else "FAIL") + " Ed25519 strict rule self-test " + ",".join(st))
    fails += ["ed25519 self-test"] if st else []
    ok = len(REASONS) == 10
    print(("PASS" if ok else "FAIL") + " §11 read from the extension: %d reasons" % len(REASONS))
    fails += [] if ok else ["§11 does not read"]
    with open(os.path.join(corpus_dir, "vectors.json"), encoding="utf-8") as f:
        doc = json.load(f)
    files = {}
    for name in sorted(os.listdir(corpus_dir)):
        if name.endswith(".m3u8"):
            with open(os.path.join(corpus_dir, name), "rb") as f:
                files[name] = f.read()
    held(doc)
    ok = rfc_vector(doc)
    print(("PASS" if ok else "FAIL") + " RFC 9578 Appendix A.2, its first test vector: its token verifies under its key")
    got, written = check(copy.deepcopy(doc), files)
    fails += got
    extra = sorted(set(os.listdir(corpus_dir)) - set(written) - {"vectors.json", "README.md"})
    print(("PASS" if not extra else "FAIL") + " no file but vectors.json, README.md and the tag files the rooms write")
    fails += ["an extra file: %s" % n for n in extra]
    seen = set()
    for v in doc["vectors"]:
        seen |= {f["verdict"] for f in v.get("claim_facts", [])}
        seen |= {i["verdict"] for r in v.get("receipt_facts", []) for i in r.get("impressions", [])}
    missing = [r for r in REASONS if r not in seen]
    print(("PASS" if not missing else "FAIL") + " every reason of §11 appears in the corpus")
    fails += ["no claim or impression gives %s" % r for r in missing]
    for label, d, fs in mutations(doc, files):
        caught, _ = check(d, fs, quiet=True)
        print(("PASS" if caught else "FAIL") + " mutated copy fails: %s%s" % (label, " (%s)" % caught[0][:100] if caught else ""))
        if not caught:
            fails.append("a mutated copy verifies: %s" % label)
    for name in ("rules", "signature", "session", "duplicate", "asset", "quartile", "viewability", "attest_absent",
                 "attest_challenge", "attest_authenticator"):
        SKIP.clear()
        SKIP.add(name)
        caught, _ = check(copy.deepcopy(doc), files, quiet=True)
        SKIP.clear()
        print(("PASS" if caught else "FAIL") + " the corpus catches a verifier without its %s check%s"
              % (name, " (%s)" % caught[0][:80] if caught else ""))
        if not caught:
            fails.append("a verifier without its %s check passes the corpus" % name)
    for f in fails:
        print("FAIL " + f)
    print(("PASS" if not fails else "FAIL") + " ad-presentation corpus: %d rooms, %d tag files" % (len(doc["vectors"]), len(files)))
    return 1 if fails else 0


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("usage: ad_presentation_verify.py <live-conformance/presentation>", file=sys.stderr)
        sys.exit(2)
    sys.exit(main(sys.argv[1]))