"""Deep fuzz of hdcp2_process_auth_message_v2 (cmd 0x0F) -- the full 580-byte struct.

0x0F memcpy_aligned's 0x244 bytes of host data into a stack struct then
interprets it. Earlier runs only varied subcmd(+0x10) and N(+0x2B); this fuzzes
every dword of the struct. Uses cmd 0x11 (close_all_sessions) to recycle slots
so the run doesn't die of slot exhaustion (the defect that silently truncated
earlier campaigns).
"""
import ctypes,struct,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/'; S=0x8000; MARK=0x5A
RANGES=[(0x00020000,0x00070000,'sOS'),(0x00400000,0x00412000,'HDCP'),
        (0x00D00000,0x00D50000,'ASD'),(0x00E10000,0x00E80000,'SYSDRV'),(0x0020A000,0x0020C000,'TAstage')]
def classify(v):
    for lo,hi,nm in RANGES:
        if lo<=v<hi: return nm
    return None
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(S)
    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=S
    return psp._ring_submit(c).resp.session_id, cp
def inv(sid,shm,buf,tc):
    adev.vram.view(shm,S)[:]=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
    try: psp._ring_submit(ic)
    except RuntimeError: pass
    except TimeoutError: return None
    return bytes(adev.vram.view(shm,S)[:])
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())
def topo(h):
    b=bytearray(S); struct.pack_into('<I',b,0,2); struct.pack_into('<10I',b,0x10,h,1,1,0,0,0,0,0,0,0)
    o=inv(dtm,dsh,b,2); return struct.unpack_from('<I',o,8)[0] if o else None
print("bring-up:",[topo(h) for h in range(8)])
def simple(cmd):
    b=bytearray(S); struct.pack_into('<I',b,0,cmd); return inv(hdcp,hsh,b,cmd)
def open_slot():
    b=bytearray(S); struct.pack_into('<I',b,0,0x0E); struct.pack_into('<I',b,0xC,0)
    o=inv(hdcp,hsh,b,0x0E); return struct.unpack_from('<I',o,4)[0], struct.unpack_from('<I',o,0x1410)[0]
simple(0x11)
g,SLOT=open_slot(); print("gate ret=%s slot=%s"%(g,SLOT))
if g!=0: print("GATE CLOSED - cold cycle needed"); sys.exit(1)
rng=random.Random(int(sys.argv[1]) if len(sys.argv)>1 else 13)
ROUNDS=int(sys.argv[2]) if len(sys.argv)>2 else 6000
EX=[0,1,2,3,4,5,6,7,8,0xff,0x100,0xffff,0x10000,0xffffffff,0x7fffffff,0x80000000,0xfffffffe,
    16,24,25,31,32,33,46,47,52,56,63,64,103,128,232,255,256,0x244,0x245,5120,5121]
rets=collections.Counter(); leaks={}; hangs=[]; sizes=collections.Counter()
for r in range(ROUNDS):
    if r%64==0:
        simple(0x11); g,SLOT=open_slot()
        if g!=0: print("  [r=%d] gate closed after recycle (ret=%s) - stopping"%(r,g)); break
    b=bytearray(S)
    for i in range(0x1400,S): b[i]=MARK
    struct.pack_into('<I',b,0,0x0F); struct.pack_into('<I',b,0xC,SLOT)
    # the 0x244-byte struct starts at +0x0C (a1 = buf+12); fuzz it wholesale
    for off in range(0x0C,0x0C+0x244,4):
        if off==0x0C: continue                        # keep a valid session index
        if rng.random()<0.55: struct.pack_into('<I',b,off,rng.choice(EX))
        else: struct.pack_into('<I',b,off,rng.getrandbits(32))
    sent=bytes(b); o=inv(hdcp,hsh,b,0x0F)
    if o is None:
        hangs.append(r); print("  [r=%d] HANG alive=%s"%(r,psp.is_sos_alive()))
        open('hang0f_struct_r%d.bin'%r,'wb').write(sent)
        if not psp.is_sos_alive(): break
        if len(hangs)>3: break
        continue
    rets[struct.unpack_from('<I',o,4)[0]]+=1
    diff=[i for i in range(0x1400,S) if sent[i]!=o[i]]
    sizes[len(diff)]+=1
    if len(diff)>40:
        print("  [r=%d] OUTLIER write %d bytes %#06x..%#06x ret=%s"%(r,len(diff),diff[0],diff[-1],struct.unpack_from('<I',o,4)[0]))
        print("       payload: %s"%o[diff[0]:diff[0]+80].hex())
        open('outlier_0f_r%d_in.bin'%r,'wb').write(sent)
        open('outlier_0f_r%d_out.bin'%r,'wb').write(o)
    for off in range(0x1400,S-3,4):
        if sent[off:off+4]==o[off:off+4]: continue
        v=struct.unpack_from('<I',o,off)[0]; cl=classify(v)
        if cl and (cl,off) not in leaks:
            leaks[(cl,off)]=v
            print("  [r=%d] *** PSP POINTER LEAK %s %#010x at shm+%#x ***"%(r,cl,v,off))
            open('leak0fstruct_r%d.bin'%r,'wb').write(o)
    if r%1500==0 and r: print("  ..r=%d alive=%s leaks=%d hangs=%d"%(r,psp.is_sos_alive(),len(leaks),len(hangs)))
print("\nret histogram: %s"%dict(rets))
print("write-size histogram: %s"%dict(sizes))
print("PSP-pointer leaks: %d   hangs: %d"%(len(leaks),len(hangs)))
print("sos_alive:",psp.is_sos_alive())