"""Hardware fuzzer for the PSP TA invoke interface (HDCP / DTM / RAS).

Every audit in this project concluded the TA surface is bounded. This tests that
empirically: for each ta_cmd_id, push structured-random payloads into the shared
buffer and watch for anomalies static analysis would have missed.

Anomaly = a status code never seen for well-formed input, a timeout/hang, or a
PSP that stops responding. A health check runs between rounds; the fuzz aborts
the moment the PSP stops answering so we never pile damage on a wedged ring.
"""
import ctypes, struct, time, random, sys, collections
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32

d = Device["AMD"]; adev = d.iface.dev_impl; psp = adev.psp
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
P = '/Users/user/dev/amd-psp-re/psp_out/'
SHM = 0x8000

def stat(): return int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr, SZ)[:])[864:868], "little")
def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True); return pa, adev.paddr2mc(pa)

def load(cmd, blob):
    ap, amc = alloc(0x40000); cp, cmc = alloc(SHM)
    adev.vram.view(ap, len(blob))[:] = blob; adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    c = am.struct_psp_gfx_cmd_resp(cmd_id=cmd); t = c.cmd.cmd_load_ta
    t.app_phy_addr_hi, t.app_phy_addr_lo = hi32(amc), lo32(amc); t.app_len = len(blob)
    t.cmd_buf_phy_addr_hi, t.cmd_buf_phy_addr_lo = hi32(cmc), lo32(cmc); t.cmd_buf_len = SHM
    r = psp._ring_submit(c)
    return r.resp.session_id, cp

def invoke(sid, shm, tc, payload):
    b = bytearray(SHM)
    b[:len(payload)] = payload
    adev.vram.view(shm, SHM)[:] = bytes(b); adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    ic = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
    iv = ic.cmd.cmd_invoke_cmd
    iv.session_id, iv.ta_cmd_id = sid, tc
    t0 = time.time()
    try:
        psp._ring_submit(ic); ps = 0
    except RuntimeError:
        ps = stat()
    except TimeoutError:
        return "TIMEOUT", time.time()-t0
    return ps, time.time()-t0

def healthy(sid, shm):
    s,_ = invoke(sid, shm, 1, b"\x00"*64)   # benign cmd
    return s != "TIMEOUT"

TA = sys.argv[1] if len(sys.argv)>1 else 'hdcp'
BLOB = {'hdcp':'hdcp_full.bin','dtm':'dtm_full.bin','ras':'ras_full.bin'}[TA]
CMDS = {'hdcp': list(range(1,16)), 'dtm': list(range(1,12)), 'ras': list(range(1,12))}[TA]
ROUNDS = int(sys.argv[2]) if len(sys.argv)>2 else 60

# ASD must be resident first: HDCP/DTM RPC into it.
asd_sid, asd_shm = load(am.GFX_CMD_ID_LOAD_ASD, open(P+'asd_full.bin','rb').read())
print("ASD loaded: sid=%#x"%asd_sid)
sid, shm = load(am.GFX_CMD_ID_LOAD_TA, open(P+BLOB,'rb').read())
print("%s loaded: sid=%#x shm=%#x"%(TA, sid, shm))
if sid == 0: print("!! sid 0"); sys.exit(1)

rng = random.Random(0xC0FFEE)
def payload(n):
    mode = rng.randrange(6)
    if mode==0: return bytes(rng.randrange(256) for _ in range(n))
    if mode==1: return b"\xff"*n
    if mode==2: return b"\x00"*n
    if mode==3:  # extreme 32-bit values at dword slots
        vals=[0,1,0xffffffff,0x7fffffff,0x80000000,0xfffffffe,0x100,0xffff]
        return b"".join(struct.pack('<I', rng.choice(vals)) for _ in range(n//4))
    if mode==4:  # small counts/lengths
        return b"".join(struct.pack('<I', rng.choice([0,1,2,23,24,25,31,32,33,63,64,255,256])) for _ in range(n//4))
    return bytes(rng.choice([0,0xff,0x7f,0x80]) for _ in range(n))

seen = collections.defaultdict(collections.Counter)
anoms = []
t_start=time.time()
for r in range(ROUNDS):
    for tc in CMDS:
        p = payload(rng.choice([16,64,256,1024,4096]))
        s, dt = invoke(sid, shm, tc, p)
        seen[tc][s]+=1
        if s=="TIMEOUT" or dt>0.5:
            anoms.append((tc, s, dt, p[:32].hex()))
            print("  ANOMALY cmd=%#x status=%s dt=%.2fs payload=%s"%(tc,s,dt,p[:32].hex()))
            if not healthy(sid, shm):
                print("  !! PSP unresponsive after cmd %#x -- aborting"%tc); r=ROUNDS; break
    if r%20==0 and r: print("  ...round %d, %.0fs"%(r, time.time()-t_start))
print("\n=== status histogram per ta_cmd_id ===")
for tc in sorted(seen):
    print("  cmd %#04x : %s"%(tc, dict(seen[tc])))
print("\nanomalies: %d"%len(anoms))
print("elapsed %.0fs"%(time.time()-t_start))