#!/usr/bin/env python3
"""Mechanically re-derive every claim in CANDIDATE-SYSDRV-E04B68-COUNT-UNDERFLOW.md.

The chain was assembled by hand over many passes, and in that stretch four separate readings
turned out to be wrong (a truncated listing, immediates mistaken for a stored value, a
misattributed base register, a register reassignment missed). Prose is therefore not adequate
evidence. This script asserts each structural claim directly against the bytes of
`psp_out/sys_drv_code.bin`, so the candidate stands or falls on reproducible checks.

Every check prints PASS/FAIL and the instruction it matched. A FAIL means the write-up is
wrong, not that the firmware changed.

Usage:  ./verify_e04b68_chain.py
"""
from __future__ import annotations

import re
import struct
import sys
from pathlib import Path

import capstone as cs

REPO = Path(__file__).resolve().parent.parent
BASE = 0xE00100
THUMB = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)

B = (REPO / "psp_out" / "sys_drv_code.bin").read_bytes()


def dis(lo: int, n: int):
    return list(THUMB.disasm(B[lo - BASE:lo - BASE + n], lo))


def at(lo: int, n: int, addr: int):
    for i in dis(lo, n):
        if i.address == addr:
            return i
    return None


def literal(i) -> int | None:
    """Resolve `ldr rX,[pc,#imm]` to the loaded word."""
    if not (i.mnemonic.startswith("ldr") and "pc" in (i.op_str or "")):
        return None
    m = re.search(r"#(0x[0-9a-f]+|\d+)", i.op_str)
    if not m:
        return None
    lit = ((i.address + 4) & ~3) + int(m.group(1), 0)
    off = lit - BASE
    if not (0 <= off < len(B) - 4):
        return None
    return struct.unpack_from("<I", B, off)[0]


RESULTS: list[tuple[bool, str, str]] = []


def check(name: str, ok: bool, detail: str) -> None:
    RESULTS.append((ok, name, detail))


def main() -> int:
    # --- Step 3: sub_E04B68 reads buf[0x88], clamps to 0x400, subtracts UNSIGNED ----------
    i = at(0xE04B68, 0x80, 0xE04B70)
    check("E04B68 reads the count at buf+0x88 before any validation",
          i is not None and i.mnemonic.startswith("ldr") and "#0x88" in (i.op_str or ""),
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    i = at(0xE04B68, 0x80, 0xE04B8E)
    clamp = i is not None and i.mnemonic.startswith("mov") and "0x400" in (i.op_str or "")
    check("E04B68 clamps the total to 0x400", clamp,
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    i = at(0xE04B68, 0x80, 0xE04BBA)
    check("E04B68 computes copy2 length as an UNSIGNED subtraction (subs r2, r7, r0)",
          i is not None and i.mnemonic == "subs" and re.match(r"r2,\s*r7,\s*r0", i.op_str or ""),
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    i = at(0xE04B68, 0x80, 0xE04BAE)
    check("E04B68 copy1 uses the count unclamped (call follows ldr r2,[r4,#0x88])",
          i is not None and i.mnemonic in ("blx", "bl"),
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    # --- Step 2: sub_E1D6BC arm copies payload into slot->buf, bounded only by capacity ----
    i = at(0xE1D6BC, 0x260, 0xE1D8A4)
    check("E1D6BC arm loads the slot buffer (ldr r5,[r4,#0x4c])",
          i is not None and "#0x4c" in (i.op_str or ""),
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    i = at(0xE1D6BC, 0x260, 0xE1D8AA)
    check("E1D6BC arm performs the payload copy into that buffer",
          i is not None and i.mnemonic in ("blx", "bl"),
          f"{i.address:#x} {i.mnemonic} {i.op_str}" if i else "not found")

    i = at(0xE1D6BC, 0x260, 0xE1D6F8)
    j = at(0xE1D6BC, 0x260, 0xE1D6FC)
    check("E1D6BC bounds the message length only by slot capacity [slot+0x48]",
          i is not None and "#0x48" in (i.op_str or "") and j is not None and j.mnemonic == "bhi",
          f"{i.address:#x} {i.mnemonic} {i.op_str} ; {j.address:#x} {j.mnemonic}" if i and j else "not found")

    # --- arm selection: slot[0] == 0x80000032 -------------------------------------------
    sel = None
    for i in dis(0xE1D6BC, 0x260):
        if literal(i) == 0x7FFFFFCE:
            sel = i
            break
    check("E1D6BC selects that arm via the 0x7fffffce normalisation (slot[0] == 0x80000032)",
          sel is not None, f"{sel.address:#x} {sel.mnemonic} {sel.op_str} ;=0x7fffffce" if sel else "not found")

    # --- Step 1: the allocator accepts format 0x80000032 the same way ---------------------
    acc = None
    for i in dis(0xE055D0, 0x3AA):
        if literal(i) == 0x7FFFFFCE:
            acc = i
            break
    ok_acc = False
    detail = "not found"
    if acc is not None:
        nxt = [x for x in dis(0xE055D0, 0x3AA) if acc.address < x.address <= acc.address + 6]
        ok_acc = any(x.mnemonic.startswith("adds") for x in nxt) and any(x.mnemonic == "beq" for x in nxt)
        detail = f"{acc.address:#x} ldr ;=0x7fffffce then " + ", ".join(
            f"{x.mnemonic} {x.op_str}" for x in nxt)
    check("allocator accepts format 0x80000032 via the same adds/beq normalisation",
          ok_acc, detail)

    # --- capacity is caller-scaled and stored at slot+0x48 --------------------------------
    i = at(0xE055D0, 0x3AA, 0xE05900)
    j = at(0xE055D0, 0x3AA, 0xE05902)
    k = at(0xE055D0, 0x3AA, 0xE05908)
    ok_cap = (i is not None and i.mnemonic == "lsrs" and "#4" in (i.op_str or "")
              and j is not None and "lsl #3" in (j.op_str or "")
              and k is not None and "#0x48" in (k.op_str or ""))
    check("capacity := 9*(n>>4)+0x10 stored to slot+0x48",
          ok_cap,
          "; ".join(f"{x.address:#x} {x.mnemonic} {x.op_str}" for x in (i, j, k) if x))

    # the operand of that shift must be reloaded from the request, not the format code
    reload_site = at(0xE055D0, 0x3AA, 0xE058DA)
    check("that capacity operand is a request field (ldr r7,[r2,#0xc]), NOT the format code",
          reload_site is not None and re.match(r"r7,\s*\[r2,\s*#0xc\]", reload_site.op_str or ""),
          f"{reload_site.address:#x} {reload_site.mnemonic} {reload_site.op_str}" if reload_site else "not found")

    # --- slot[0] is the caller-supplied format code --------------------------------------
    i = at(0xE055D0, 0x3AA, 0xE0590C)
    j = at(0xE055D0, 0x3AA, 0xE0590E)
    check("slot[0] := the caller's format code (ldr r3,[r2,#4] ; strd r3,r6,[r1])",
          i is not None and re.match(r"r3,\s*\[r2,\s*#4\]", i.op_str or "")
          and j is not None and j.mnemonic.startswith("strd"),
          "; ".join(f"{x.address:#x} {x.mnemonic} {x.op_str}" for x in (i, j) if x))

    # --- the count field is inside the message-writable region ---------------------------
    caps = [9 * (n >> 4) + 0x10 for n in (0xE0, 0x100, 0x400)]
    check("caller-selectable capacities exceed the count offset 0x8c",
          all(c >= 0x8C for c in caps), f"n=0xe0/0x100/0x400 -> {[hex(c) for c in caps]}")

    # --- the underflow itself --------------------------------------------------------
    count, clampv = 0x500, 0x400
    check("count > clamp underflows copy2's length",
          ((clampv - count) & 0xFFFFFFFF) > 0xFFFF0000,
          f"0x400 - 0x500 = {((clampv - count) & 0xFFFFFFFF):#x}")

    width = max(len(n) for _, n, _ in RESULTS)
    bad = 0
    for ok, name, detail in RESULTS:
        print(f"[{'PASS' if ok else 'FAIL'}] {name.ljust(width)}  {detail}")
        bad += 0 if ok else 1
    print(f"\n{len(RESULTS) - bad}/{len(RESULTS)} structural claims verified against the binary")
    if bad:
        print("FAIL means the write-up is wrong and must be corrected.")
    return 1 if bad else 0


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