#!/usr/bin/env python3
"""Fuzz TA_InvokeCommandEntryPoint over the REAL host path.

Every previous TA fuzzer in this tree called a dispatcher directly with
synthetic state and Python stubs standing in for kernel services. This one
drives the actual surface:

    host writes mailbox header + ring frame
      -> IRQ 80 doorbell -> bl_legacy_mailbox_loop -> tee_ring_drain
      -> tee_cmd_dispatch -> tee_ta_command -> TA_InvokeCommandEntryPoint

so sOS validates the buffers, the MMU is on, the TA runs in its own process, and
the command arrives the way the driver sends it.

CRASH ORACLE
    sOS's `sos_fault_handler` (VA 0xEE4) is entered on prefetch/data abort,
    undef and SError, and `sos_fault_kill_ta` kills just the faulting TA. If a
    fuzzed command makes the TA fault, the handler runs -- that is a real crash
    signal from the firmware itself, not an emulator artefact.

    The oracle is validated with a POSITIVE CONTROL: a deliberately corrupted
    TA entry point is executed and the handler must fire. Without that, "no
    crashes" would be unfalsifiable.
"""
import sys, os, struct, time, random, collections
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.dirname(__file__))

from unicorn import UC_HOOK_CODE, UC_HOOK_BLOCK

import importlib.util
_b = importlib.util.spec_from_file_location(
    "bm", os.path.join(os.path.dirname(__file__), "boot_and_mailbox.py"))
BM = importlib.util.module_from_spec(_b); _b.loader.exec_module(BM)
M = BM.M

SOS_FAULT_HANDLER = 0x0000EE4          # sos_fault_handler
HDCP_INVOKE = 0x004013B4               # TA_InvokeCommandEntryPoint
TA_LO, TA_HI = 0x400000, 0x500000
TA_IMAGE_PA = 0x04400000
INVOKE_BUDGET = 500_000
INVOKE_SLICE = 25_000       # poll granularity; a whole invoke measures ~2k blocks
SETTLE_SLICES = 4           # consecutive quiet slices before an invoke is done


def setup(verbose=True):
    c, ccp, irqs, exc = BM.build()
    run, _ = M.make_det_runner(c, irqs)
    run(6_000_000)
    M.install_debug_gpio(c)

    box = M.Mailbox(c, irqs)
    c.uc.mem_write(BM.RING_PA, b"\x00" * BM.RING_SZ)
    rp = bytearray(0x1C)
    rp[0x10:0x14] = struct.pack("<I", BM.RING_PA)
    rp[0x18:0x1C] = struct.pack("<I", BM.RING_SZ)
    st, _ = box.submit(2, bytes(rp), timeout=12, runner=run)
    assert st == 0, "ring init failed: %s" % st

    ring = M.Ring(c, irqs, BM.RING_PA, BM.RING_SZ, BM.CMDBUF_PA, run)

    from pspemu import armmmu
    faults = []
    ta_blocks = set()          # (VA, live PA) -- VA alone cannot identify a TA
    c.uc.hook_add(UC_HOOK_CODE, lambda uc, a, s, u: faults.append(a),
                  begin=SOS_FAULT_HANDLER, end=SOS_FAULT_HANDLER + 2)
    def ta_hook(uc, a, s, u):
        if TA_LO <= a < TA_HI:
            ta_blocks.add((a, armmmu.translate_live(uc, a)))
    c.uc.hook_add(UC_HOOK_BLOCK, ta_hook)
    disp = []
    c.uc.hook_add(UC_HOOK_CODE, lambda uc, a, s, u: disp.append(a),
                  begin=0x2027C4, end=0x2027C6)
    c.uc.ctl_flush_tb()

    img = open(os.path.join(BM.PSP_OUT, "hdcp_full.bin"), "rb").read()
    for base in (TA_IMAGE_PA, TA_IMAGE_PA + M.Ring.ALIAS):
        c.uc.mem_write(base, img)
    ld = bytearray(0x20)
    ld[0:4] = struct.pack("<I", TA_IMAGE_PA)
    ld[8:12] = struct.pack("<I", len(img))
    def live_slots():
        _k, pa = ring.pt.describe(0x206080)
        tbl = int.from_bytes(c.uc.mem_read(pa, 4), "little")
        out = []
        for i in range(0x30):
            _k2, p2 = ring.pt.describe(tbl + 32 * i + 76928)
            if p2 and c.uc.mem_read(p2, 1)[0] != 0xFF:
                out.append(i)
        return out

    ring.submit(1, bytes(ld), timeout=10, probe=disp)
    # was: for _ in range(10): run(5_000_000)  -- 50M blocks of setup on every
    # fuzzer start, regardless of when the TA actually landed in a session slot.
    slots = BM.drain_until(run, live_slots, slices=10, blocks=5_000_000)
    assert slots, "TA did not load through the ring"
    if verbose:
        print("setup: TA loaded via host ring, session slot %d" % slots[0])
    return dict(c=c, ccp=ccp, irqs=irqs, exc=exc, run=run, ring=ring,
                disp=disp, faults=faults, ta=ta_blocks, slot=slots[0])


def invoke(env, payload, frame_params=(0, 0, 0)):
    """One host-driven TA invoke with fuzzed parameters."""
    c, ring = env["c"], env["ring"]
    buf = bytearray(0x400)
    buf[0:4] = struct.pack("<I", 0x400)
    buf[4:8] = struct.pack("<I", 1)
    buf[8:12] = struct.pack("<I", 3)                 # command id 3 = TA command
    buf[0x1C:0x20] = struct.pack("<I", env["slot"])  # target session
    buf[0x20:0x20 + len(payload)] = payload
    ring._mirror(ring.cmdbuf, bytes(buf))

    frame = bytearray(64)
    frame[0:4] = struct.pack("<I", ring.cmdbuf)
    frame[12:16] = struct.pack("<I", frame_params[0] & 0xFFFFFFFF)
    frame[16:20] = struct.pack("<I", frame_params[1] & 0xFFFFFFFF)
    frame[20:24] = struct.pack("<I", frame_params[2] & 0xFFFFFFFF)
    c.uc.mem_write(ring.ring_phys + ring.widx * 4, bytes(frame))
    nxt = (ring.widx + 16) % ring.ndw
    hdr = bytearray(M.MBOX_LEN)
    hdr[0:4] = struct.pack("<I", M.DONE)
    hdr[0x0C:0x10] = struct.pack("<I", nxt)
    c.uc.mem_write(M.MBOX, bytes(hdr))
    ring.widx = nxt

    n0 = len(env["disp"])
    env["emu_err"] = False
    env["irqs"].raise_irq(M.DOORBELL_IRQ)

    # Phase 1: wait for the ring frame to reach tee_cmd_dispatch.
    # Phase 2: let the invoke run to completion, then stop once TA execution AND
    # fault state have both been unchanged for SETTLE_SLICES consecutive slices.
    #
    # Total patience is unchanged (4*INVOKE_BUDGET then INVOKE_BUDGET), so a slow
    # or wedged invoke gets exactly the same budget it always did. What changes
    # is the common case: measured, a whole invoke -- doorbell, dispatch, TA
    # entry, all 22 TA blocks -- completes within ~2k blocks, while the old code
    # unconditionally ran 1,000,000 (one full slice to see the dispatch, plus a
    # fixed 500k settle). Quiescence is used rather than a smaller constant so
    # that a genuinely slow invoke is still waited out instead of truncated.
    spent = 0
    while spent < 4 * INVOKE_BUDGET and len(env["disp"]) == n0:
        if not env["run"](INVOKE_SLICE):
            env["emu_err"] = True        # emu_start raised -- see note below
        spent += INVOKE_SLICE

    spent, quiet = 0, 0
    sig = (len(env["ta"]), len(env["faults"]))
    while spent < INVOKE_BUDGET and quiet < SETTLE_SLICES:
        if not env["run"](INVOKE_SLICE):
            env["emu_err"] = True
        spent += INVOKE_SLICE
        now = (len(env["ta"]), len(env["faults"]))
        quiet = quiet + 1 if now == sig else 0
        sig = now
    return len(env["disp"]) > n0


def main():
    iters = int(os.environ.get("ITERS", "80"))
    random.seed(int(os.environ.get("SEED", "1234")))
    env = setup()
    c = env["c"]

    # ---------------- controls -------------------------------------------
    env["ta"].clear(); env["faults"].clear()
    ok = invoke(env, b"\x00" * 0x60)
    base_cov = set(env["ta"])
    entry_pas = sorted({pa for va, pa in base_cov if va == HDCP_INVOKE and pa})
    pages = sorted({pa & ~0xFFF for _va, pa in base_cov if pa})
    print("control (benign invoke): dispatched=%s TA blocks=%d faults=%d"
          % (ok, len(base_cov), len(env["faults"])))
    print("   distinct TA code pages executing: %s" % [hex(x) for x in pages])
    print("   VA %#x resolved live to PA: %s" % (HDCP_INVOKE, [hex(x) for x in entry_pas]))
    if not entry_pas:
        print("NEGATIVE CONTROL FAILED: TA_InvokeCommandEntryPoint did not execute")
        return 2
    if env["faults"]:
        print("NEGATIVE CONTROL FAILED: fault handler fired on a benign invoke")
        return 2

    # Positive control: corrupt the TA entry point so it MUST fault.
    #
    # Patch exactly the PHYSICAL page the invoked session executed, taken from
    # the live translation observed during the benign control. Two earlier
    # versions failed here: resolving the VA through a cached TTBR0 hit another
    # task's page, and searching for the image bytes found whichever resident
    # copy came first. Every TA lives at VA 0x400000 in its own address space,
    # so only the live (VA -> PA) pair identifies the right instance.
    print("\npositive control: corrupting the TA entry point to prove the oracle fires")
    img_all = open(os.path.join(BM.PSP_OUT, "hdcp_full.bin"), "rb").read()
    foff = 0x100 + (HDCP_INVOKE - 0x400000)          # 0x100 container header
    needle = img_all[foff:foff + 16]
    # Do NOT filter regions by size: the live TA page sits in the 64 MB
    # 0x4000000-0x8000000 map, and skipping large regions is what made an
    # earlier version of this control fail to find it. Instead exclude the
    # addresses we know are *copies* of the image rather than the running code:
    # our host staging buffer, its aliased view, and the boot staging slot.
    cands = entry_pas
    print("   patching executing entry PA(s): %s" % [hex(x) for x in cands])
    saved = [(x, bytes(c.uc.mem_read(x, 4))) for x in cands]
    for x in cands:
        c.uc.mem_write(x, b"\xff\xff\xff\xff")   # undefined instruction
    c.uc.ctl_flush_tb()

    env["faults"].clear(); env["ta"].clear()
    nexc0 = sum(1 for i, _pc in env["exc"] if i != 2)
    got = invoke(env, b"\x00" * 0x60)
    nexc1 = sum(1 for i, _pc in env["exc"] if i != 2)
    print("   dispatched=%s  TA blocks=%d  entry executed=%s"
          % (got, len(env["ta"]), HDCP_INVOKE in env["ta"]))
    print("   bytes now at entry copies: %s"
          % [bytes(c.uc.mem_read(x, 4)).hex() for x in cands])
    # Three channels, because a TA fault does not always reach the guest:
    #   * sos_fault_handler          -- the firmware's own handler
    #   * non-SWI guest exception    -- routed abort/undef
    #   * emu_start error            -- Unicorn reports an UNDEFINED INSTRUCTION
    #                                   as UC_ERR_INSN_INVALID and stops, rather
    #                                   than delivering intno 1 to the guest.
    # Watching only the guest channels makes an undefined-instruction crash
    # invisible, which is what made two earlier versions of this control fail.
    fired = len(env["faults"]) > 0 or nexc1 > nexc0 or env.get("emu_err")
    print("   non-SWI guest exceptions: %d   emu_start error: %s"
          % (nexc1 - nexc0, env.get("emu_err")))
    for x, orig in saved:
        c.uc.mem_write(x, orig)
    c.uc.ctl_flush_tb()
    print("   sos_fault_handler fired: %s" % fired)
    if not fired:
        print("POSITIVE CONTROL FAILED: crash oracle cannot detect a TA fault;")
        print("a 'no crashes' result from this harness would be meaningless.")
        return 2

    # ---------------- fuzz ------------------------------------------------
    print("\nfuzzing %d host-driven invokes..." % iters)
    cov = set(base_cov)
    crashes = []
    dead = 0
    t0 = time.time()
    reached = 0
    for i in range(iters):
        n = random.choice([0x10, 0x40, 0x80, 0x200, 0x3C0])
        payload = bytearray(random.getrandbits(8) for _ in range(max(n, 8)))
        # Payload +0x04 (command buffer +0x24) GATES TA ENTRY. Measured by
        # single-byte bisection over the payload: only bytes +0x04..+0x07 change
        # the outcome, and of the values scanned (0..47, 0x100, 0xFFFF,
        # 0xFFFFFFFF) ONLY 0 is accepted -- anything else is rejected before
        # TA_InvokeCommandEntryPoint runs. Leaving it random is why this fuzzer
        # reached ZERO TA blocks in 80 iterations while still reporting
        # "0 faults". Pin it so the invokes actually enter the TA; set FUZZ_RAW=1
        # to restore the old (unreachable) behaviour.
        if not os.environ.get("FUZZ_RAW"):
            payload[4:8] = struct.pack("<I", 0)
        payload = bytes(payload)
        params = (random.choice([0, 1, 0xFFFFFFFF, random.getrandbits(32)]),
                  random.getrandbits(32), random.getrandbits(32))
        env["faults"].clear(); env["ta"].clear()
        e0 = sum(1 for k, _pc in env["exc"] if k != 2)
        got = invoke(env, payload, params)
        e1 = sum(1 for k, _pc in env["exc"] if k != 2)
        if not got:
            dead += 1
        if env["faults"] or e1 > e0 or env.get("emu_err"):
            crashes.append((i, len(payload), params,
                            "fault_handler" if env["faults"]
                            else ("guest_exc" if e1 > e0 else "emu_err")))
        if env["ta"]:
            reached += 1
        cov |= env["ta"]
    dt = time.time() - t0

    new_cov = cov - base_cov
    print("\n== result ==")
    print("  iterations            : %d in %.1fs (%.2fs each)" % (iters, dt, dt / max(1, iters)))
    print("  TA blocks covered     : %d (control alone: %d)" % (len(cov), len(base_cov)))
    print("  iterations reaching TA: %d / %d" % (reached, iters))
    print("  new blocks vs control : %d" % len(new_cov))
    print("  invokes not dispatched: %d" % dead)
    print("  TA faults detected    : %d" % len(crashes))
    for i, n, p, kind in crashes[:10]:
        print("     iter %-4d %-14s payload=%#x params=%s"
              % (i, kind, n, [hex(x) for x in p]))
    print("  ccp: %s errors=%d" % (dict(env["ccp"].stats()), len(env["ccp"].errors)))

    # REACHABILITY GATE. Same discipline as the positive control above: a clean
    # negative is only a result if the harness demonstrably exercised the code.
    # Two ways this fuzzer can look clean while measuring nothing:
    #   1. iterations never enter the TA at all (the +0x24 gate, see above);
    #   2. they enter but never vary the path, so every invoke re-walks the
    #      control's blocks and no input can possibly matter.
    # Both were TRUE before the pin was added, and neither was reported.
    if not crashes:
        if reached == 0:
            print("\nRESULT NOT REPORTABLE: 0/%d iterations executed ANY TA block."
                  % iters)
            print("'0 faults' here means the TA was never entered, not that it is safe.")
            return 2
        if not new_cov:
            print("\nRESULT NOT REPORTABLE: %d/%d iterations entered the TA, but they"
                  % (reached, iters))
            print("covered NO blocks beyond the benign control -- the mutated bytes do")
            print("not influence TA execution, so this negative is not evidence of")
            print("absence. The fuzzer is mutating the wrong surface.")
            return 2
    return 0


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