#!/usr/bin/env python3
"""Full-resolution per-dword sweep of the WHOLE shared buffer.

Payload fuzzing so far mutated offsets 4..0xfc (deep fuzz) or filled broad regions uniformly.
A command that reads a parameter deeper in the buffer would be missed by both. With the ring
ceiling removed (pass 85) a per-dword sweep of the entire 0x4000 window is affordable.

For each dword offset: set it to 0xFFFFFFFF, invoke, compare status and the first 64 bytes of
output against the unmutated baseline. Report offsets that matter, stalls, or writes beyond
the window.
"""
import glob, struct, sys, time
sys.path.insert(0,'/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
from ring_wrap import install
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
submit=install(psp,adev)
PAT=0x11111111; WIN=0x4000; TOTAL=0x8000
TAS={'AUC':('ta_16-0-0-9_9e9044cf',5),'CSU':('ta_42-0-0-6_1c874c3c',1),'OC':('ta_50-0-0-1_cbd47c59',1)}
name=sys.argv[1]; stem,cmd=TAS[name]
base=open(glob.glob('/root/psp/contall/'+stem+'*')[0],'rb').read()
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
ap_,amc=alloc(0x80000); cp,cmc=alloc(TOTAL)
n=(len(base)+3)&~3
adev.vram.view(ap_,n)[:]=base+b'\x00'*(n-len(base)); adev.gmc.flush_hdp()
c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t=c.cmd.cmd_load_ta
t.app_phy_addr_lo,t.app_phy_addr_hi=amc&0xFFFFFFFF,amc>>32; t.app_len=len(base)
t.cmd_buf_phy_addr_lo,t.cmd_buf_phy_addr_hi=cmc&0xFFFFFFFF,cmc>>32; t.cmd_buf_len=WIN
r=submit(c); sid=r.resp.session_id
print('%s load %#x sid=%s cmd=%d'%(name,r.resp.status&0xFFFFFFFF,sid,cmd), flush=True)
def inv(off=None):
    buf=bytearray(WIN); struct.pack_into('<I',buf,0,cmd)
    if off is not None: struct.pack_into('<I',buf,off,0xFFFFFFFF)
    adev.vram.view(cp,WIN)[:]=bytes(buf)
    adev.vram.view(cp+WIN,TOTAL-WIN)[:]=bytes(struct.pack('<I',PAT)*((TOTAL-WIN)//4))
    adev.gmc.flush_hdp()
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
    iv=c.cmd.cmd_invoke_cmd; iv.session_id,iv.ta_cmd_id=sid,cmd
    t0=time.time()
    try: st=submit(c).resp.status&0xFFFFFFFF
    except Exception as e:
        try: st=int(str(e).strip().split()[-1])&0xFFFFFFFF
        except Exception: st=-1
    dt=time.time()-t0
    adev.gmc.flush_hdp()
    o=bytes(adev.vram.view(cp,TOTAL)[:])
    bey=sum(1 for x in range(WIN,TOTAL,4) if struct.unpack_from('<I',o,x)[0]!=PAT)
    return st, dt, o[4:0x44], bey
bst,_,bout,_ = inv()
print('  baseline status %#x'%bst, flush=True)
hits=[]; stalls=0
for off in range(4, WIN, 4):
    st,dt,out,bey = inv(off)
    if dt>5: print('  STALL at offset %#x'%off, flush=True); stalls=1; break
    if bey: hits.append(('OOB',hex(off),bey)); break
    if st!=bst or out!=bout: hits.append((hex(off),hex(st),'out-differs' if out!=bout else ''))
print('%s: %d dwords swept; offsets that matter: %s'%(name,(WIN-4)//4,
      hits[:10] if hits else 'NONE'), flush=True)