#!/usr/bin/env python3
"""Handle replay at EVERY offset — closing the weak spot recorded in audit pass 94.

`ta_handlefuzz.py` injected harvested handles only at offset +4, so a real handle field
elsewhere was never exercised and its negative was explicitly recorded as weak evidence.
With the ring ceiling removed, the injection point can be swept across the whole window.

Method: harvest handles the TA emits, then for each dword offset inject (a) a harvested
handle and (b) a random value, and compare. An offset where the HANDLE is accepted and the
random value is rejected is a field that is both read and liveness/ownership checked -- and
if the handle was destroyed first, that is a use-after-free.

Sequence per offset: destroy-ish pass first (invoke every command once to give the TA a
chance to free things), then replay.
"""
import glob, random, 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)
WIN=0x4000
TAS={'GCP':'ta_22-0-0-77_7c426e7f','MERIT':'ta_24-0-0-70_5cd4afda','OPM':'ta_25-0-0-81_dc596ba6'}
name=sys.argv[1]; NOFF=int(sys.argv[2]) if len(sys.argv)>2 else 512
blob=open(glob.glob('/root/psp/contall/'+TAS[name]+'*')[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(0x8000)
n=(len(blob)+3)&~3
adev.vram.view(ap_,n)[:]=blob+b'\x00'*(n-len(blob)); 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(blob)
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'%(name,r.resp.status&0xFFFFFFFF,sid), flush=True)
def inv(cid, off=None, val=0):
    b=bytearray(WIN); struct.pack_into('<I',b,0,cid)
    if off is not None: struct.pack_into('<I',b,off,val&0xFFFFFFFF)
    adev.vram.view(cp,WIN)[:]=bytes(b); 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,cid
    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()
    return st, dt, bytes(adev.vram.view(cp,0x80)[:])
# harvest handles, then run every command once so the TA can free whatever it frees
handles=set()
for cid in range(0,12):
    st,dt,out=inv(cid)
    if dt>5: print('  stall harvesting cmd %d'%cid, flush=True); sys.exit(3)
    for o in range(4,0x80,4):
        v=struct.unpack_from('<I',out,o)[0]
        if v not in (0,cid,0xFFFFFFFF): handles.add(v)
for cid in range(0,12): inv(cid)          # give the TA a chance to destroy
print('  harvested %d handles; sweeping %d offsets x 12 cmds'%(len(handles),NOFF), flush=True)
if not handles: print('  no handles -> test is void'); sys.exit(0)
h=sorted(handles)[0]
rng=random.Random(0xA5); hits=[]
for off in range(4, 4+NOFF*4, 4):
    for cid in (1,3,5,9):
        sh,dth,_ = inv(cid, off, h)
        sr,dtr,_ = inv(cid, off, rng.getrandbits(32))
        if dth>5 or dtr>5: hits.append(('STALL',hex(off),cid)); break
        if sh==0 and sr!=0: hits.append((hex(off),cid,'handle-accepted %#x vs random %#x'%(sh,sr)))
    if hits and hits[-1][0]=='STALL': break
print('%s: offsets where a harvested handle is accepted but a random value is not: %s'
      %(name, hits[:6] if hits else 'NONE'), flush=True)