"""Ring-level fuzzer for the PSP GFX command surface + disclosure detector.

Prior campaigns only fuzzed TA shared-buffer CONTENTS. The host also fully
controls the 1024-byte psp_gfx_cmd_resp struct itself (cmd_id and every
address/length/session field). That is the most directly host-reachable PSP
surface and had never been fuzzed.

SAFETY (psp_out/CONSTRAINTS.md: no flash / no TSU):
  cmd_id 7 (DESTROY_TMR -> sysdrv_get_system_info, flash ops) is EXCLUDED.
  cmd_id 5/6/8/9/10 (TMR/VMR/IP_FW/SAVE_RESTORE) excluded from field fuzzing --
  they reprogram memory protection or load firmware. We fuzz the TA lifecycle
  (LOAD_TA/LOAD_ASD/UNLOAD/INVOKE), attestation, and undefined cmd_ids.

DISCLOSURE DETECTOR: the cmd buffer is filled with a marker before each submit.
Any byte the PSP overwrites with something that is neither marker nor a known
response field is candidate leaked PSP memory.
"""
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/'
MARK=0xA5
FORBIDDEN={7,5,6,8,9,10}          # flash / memory-protection / fw-load
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True); return pa,adev.paddr2mc(pa)
app_pa,app_mc=alloc(0x40000); shm_pa,shm_mc=alloc(0x8000)
blob=open(P+'hdcp_full.bin','rb').read()
adev.vram.view(app_pa,len(blob))[:]=blob
adev.gmc.flush_hdp()
def submit(raw:bytes, tmo_note=""):
    adev.vram.view(psp.cmd_paddr,SZ)[:]=raw
    c=am.struct_psp_gfx_cmd_resp.from_buffer_copy(raw)
    t0=time.time()
    try:
        r=psp._ring_submit(c); return "OK",bytes(adev.vram.view(psp.cmd_paddr,SZ)[:]),time.time()-t0
    except RuntimeError as e:
        return "ERR",bytes(adev.vram.view(psp.cmd_paddr,SZ)[:]),time.time()-t0
    except TimeoutError:
        return "TIMEOUT",None,time.time()-t0
def mkcmd(cid, fields:dict):
    b=bytearray(SZ)   # zero-filled: a marker fill corrupts reserved/resp fields and everything is rejected at parse
    struct.pack_into('<I',b,8,cid)          # cmd_id lives at offset 8
    for off,val in fields.items(): struct.pack_into('<I',b,off,val&0xffffffff)
    return bytes(b)
rng=random.Random(int(sys.argv[1]) if len(sys.argv)>1 else 11)
ROUNDS=int(sys.argv[2]) if len(sys.argv)>2 else 4000
EX=[0,1,2,3,0xffffffff,0x7fffffff,0x80000000,0xfffffffe,0xffff,0x10000,0x100000,0x100001,
    0xfffff,0x200,0x8000,0x8001,0x40000,0x40001,255,256,0x2817,0x2818,64,63]
# byte offsets of the cmd union payload inside psp_gfx_cmd_resp
UNION=28   # verified by ctypes sentinel probing: cmd union starts at byte 28
stats=collections.Counter(); anoms=[]; leaks=[]
t0=time.time(); lastrep=0
print("ring-level GFX fuzz, forbidden cmd_ids=%s"%sorted(FORBIDDEN))
for r in range(ROUNDS):
    mode=rng.randrange(4)
    if mode==0:      cid=rng.choice([1,4])                 # LOAD_TA / LOAD_ASD
    elif mode==1:    cid=rng.choice([2,3])                 # UNLOAD / INVOKE
    elif mode==2:    cid=rng.choice([0xf,0x20,0x22,0x27])  # attestation / toc / bootcfg / sriov
    else:            cid=rng.choice([0,0xc,0xd,0xe,0x11,0x17,0x1f,0x23,0x26,0x28,0x3f,0xff,0x10000,0xf0000])
    if cid in FORBIDDEN: continue
    fields={}
    for k in range(6):
        off=UNION+4*rng.randrange(0,12)
        fields[off]=rng.choice(EX)
    # sometimes make it a plausible well-formed load/invoke so we get past the gate
    if cid in (1,4) and rng.random()<0.5:
        fields={UNION+0:hi32(app_mc),UNION+4:lo32(app_mc),UNION+8:rng.choice([len(blob)]+EX),
                UNION+12:hi32(shm_mc),UNION+16:lo32(shm_mc),UNION+20:rng.choice([0x8000]+EX)}
    raw=mkcmd(cid,fields)
    st,out,dt=submit(raw)
    code=struct.unpack_from('<I',out,864)[0] if out else None
    stats[(cid,code)]+=1
    if st=="TIMEOUT" or dt>0.5:
        anoms.append((r,cid,st,dt)); print("  [r=%d] ANOMALY cid=%#x st=%s dt=%.2f alive=%s"%(r,cid,st,dt,psp.is_sos_alive()))
        open('gfxanom_%02x_r%d.bin'%(cid,r),'wb').write(mkcmd(cid,fields))
        if not psp.is_sos_alive(): print("  !! sOS DEAD"); break
        if len(anoms)>5: break
        continue
    # disclosure: bytes the PSP wrote that are not our marker, outside the header (first 0x40)
    # disclosure: non-zero bytes outside the header+status window we know the PSP writes
    # disclosure: bytes that DIFFER from what we submitted, outside the known response window
    changed=[i for i,(a,b) in enumerate(zip(raw,out)) if a!=b and not (860<=i<876)]
    if changed:
        sig=(cid,len(changed),changed[0],changed[-1])
        if sig not in [l[0] for l in leaks] and len(leaks)<25:
            leaks.append((sig,out))
            print("  [r=%d] PSP WROTE %d bytes at 0x%x..0x%x for cid=%#x"%(r,len(changed),changed[0],changed[-1],cid))
            open('gfxleak_%02x_r%d.bin'%(cid,r),'wb').write(out)
    if time.time()-lastrep>60:
        lastrep=time.time(); print("  ..r=%d %.0fs alive=%s"%(r,time.time()-t0,psp.is_sos_alive()))
print("\nrounds=%d elapsed=%.0fs anomalies=%d distinct-write-patterns=%d"%(ROUNDS,time.time()-t0,len(anoms),len(leaks)))
print("status by cmd_id:")
agg=collections.defaultdict(collections.Counter)
for (cid,st),n in stats.items(): agg[cid][st]=n
for cid in sorted(agg): print("   %#07x : %s"%(cid,dict(agg[cid])))
print("sos_alive:",psp.is_sos_alive())