"""Deep TA fuzz campaign -- now that the PSP ring wptr wrap is fixed.

Every previous campaign terminated at exactly 1024 ring submissions because of
the driver's non-wrapping wptr; that bug masked whatever the fuzzer would have
found later. This runs unbounded.

Targets HDCP + DTM + RAS. Tracks the status-code space per command and flags:
  - timeouts / slow calls (real TA hangs, no longer confounded by the ring bug)
  - NOVEL status codes appearing only after long state accumulation
  - sOS death
"""
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/'
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(0x8000)
    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=0x8000
    return psp._ring_submit(c).resp.session_id,cp
NINV=0
def inv(sid,shm,buf,tc):
    global NINV
    adev.vram.view(shm,0x8000)[:]=bytes(buf); 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
    NINV+=1; t0=time.time()
    try: psp._ring_submit(ic); return 0,bytes(adev.vram.view(shm,0x80)[:]),time.time()-t0
    except RuntimeError: return stat(),None,time.time()-t0
    except TimeoutError: return "TIMEOUT",None,time.time()-t0
asd,_   = load(am.GFX_CMD_ID_LOAD_ASD, open(P+'asd_full.bin','rb').read())
dtm,dsh = load(am.GFX_CMD_ID_LOAD_TA,  open(P+'dtm_full.bin','rb').read())
hdcp,hsh= load(am.GFX_CMD_ID_LOAD_TA,  open(P+'hdcp_full.bin','rb').read())
ras,rsh = load(am.GFX_CMD_ID_LOAD_TA,  open(P+'ras_full.bin','rb').read())
print("sids asd=%#x dtm=%#x hdcp=%#x ras=%#x"%(asd,dtm,hdcp,ras))
def topo(h):
    b=bytearray(0x8000); struct.pack_into('<I',b,0,2); struct.pack_into('<10I',b,0x10,h,1,1,0,0,0,0,0,0,0)
    return inv(dtm,dsh,b,2)[0]
for h in range(8): topo(h)
TARGETS=[('hdcp',hdcp,hsh,list(range(1,16))),('dtm',dtm,dsh,list(range(1,12))),('ras',ras,rsh,list(range(1,12)))]
EX=[0,1,2,3,4,5,7,8,0xffffffff,0x7fffffff,0x80000000,0xfffffffe,0xffff,0x10000,255,256,
    24,25,31,32,33,62,63,64,0x27a,0x27b,0x27c,0x2817,0x2818,0x1400,0x1401]
rng=random.Random(int(sys.argv[1]) if len(sys.argv)>1 else 7)
ROUNDS=int(sys.argv[2]) if len(sys.argv)>2 else 20000
seen=collections.defaultdict(set); hist=collections.Counter(); anoms=[]
t0=time.time(); lastreport=0
for r in range(ROUNDS):
    name,sid,shm,cmds=rng.choice(TARGETS); cmd=rng.choice(cmds)
    b=bytearray(0x8000); struct.pack_into('<I',b,0,cmd)
    for off in (0x04,0x08,0x0C,0x10,0x14,0x18,0x1C,0x20,0x24):
        if rng.random()<0.7: struct.pack_into('<I',b,off,rng.choice(EX))
    n=rng.choice([8,32,128,512,2048,5000])
    for k in range(0x28,min(0x28+n,0x8000)): b[k]=rng.randrange(256)
    for _ in range(10):
        off=rng.randrange(0x28,min(0x28+n,0x7ff0))&~3
        struct.pack_into('<I',b,off,rng.choice(EX))
    ps,out,dt=inv(sid,shm,b,cmd)
    ret=struct.unpack_from('<I',out,4)[0] if out else None
    key=(name,cmd); hist[(name,cmd,ret)]+=1
    if ret is not None and ret not in seen[key]:
        seen[key].add(ret)
        if len(seen[key])>1 and r>2000:
            print("  [r=%d] NOVEL status %s for %s cmd %#x (after %d rounds)"%(r,ret,name,cmd,r))
            open('novel_%s_%02x_%s.bin'%(name,cmd,ret),'wb').write(bytes(b))
    if ps=="TIMEOUT" or dt>0.5:
        anoms.append((r,name,cmd,ps,dt))
        print("  [r=%d] ANOMALY %s cmd=%#x ps=%s dt=%.2f sos_alive=%s"%(r,name,cmd,ps,dt,psp.is_sos_alive()))
        open('anom_%s_%02x_r%d.bin'%(name,cmd,r),'wb').write(bytes(b))
        if not psp.is_sos_alive(): print("  !! sOS DEAD"); break
        if len(anoms)>6: print("  too many anomalies, stopping"); break
    if time.time()-lastreport>60:
        lastreport=time.time(); print("  ..r=%d ninv=%d %.0fs alive=%s"%(r,NINV,time.time()-t0,psp.is_sos_alive()))
print("\nrounds=%d ninv=%d elapsed=%.0fs anomalies=%d"%(ROUNDS,NINV,time.time()-t0,len(anoms)))
print("status space per (ta,cmd):")
for k in sorted(seen): print("   %-6s cmd %#04x : %s"%(k[0],k[1],sorted(x for x in seen[k] if x is not None)))
print("sos_alive:",psp.is_sos_alive())