#!/usr/bin/env python3
"""Out-of-bounds write test across the CORPUS TA images (not just the 12 in use).

All hardware work so far used the twelve images this part currently loads. The corpus holds
78 distinct TA builds, and pass 68 shelved the question "does any other build pass a
caller-chosen length?" because static analysis lacked per-image function boundaries. It can
be asked dynamically instead: load each build and check whether it writes past the declared
command-buffer window.

Per image: LOAD_TA with cmd_buf_len = 0x1000 over an 0x8000 guard-filled allocation, invoke
commands 0..7, report in-window writes (the control) and any dword beyond 0x1000.
"""
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)
import os
PAT=0xA5A5A5A5; DECL=int(os.environ.get('DECL','0x1000'),0); TOTAL=0x8000
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
def load(blob):
    ap_,amc=alloc(0x80000); cp,cmc=alloc(TOTAL)
    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=DECL
    try: r=submit(c); return r.resp.status&0xFFFFFFFF, r.resp.session_id, cp
    except Exception as e:
        try: return int(str(e).strip().split()[-1])&0xFFFFFFFF,None,cp
        except Exception: return -1,None,cp
def probe(sid,shm,cid):
    buf=bytearray(struct.pack('<I',PAT)*(TOTAL//4)); struct.pack_into('<I',buf,0,cid)
    adev.vram.view(shm,TOTAL)[:]=bytes(buf); 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()
    out=bytes(adev.vram.view(shm,TOTAL)[:])
    ins=sum(1 for o in range(4,DECL,4) if struct.unpack_from('<I',out,o)[0]!=PAT)
    bey=[(o,struct.unpack_from('<I',out,o)[0]) for o in range(DECL,TOTAL,4)
         if struct.unpack_from('<I',out,o)[0]!=PAT]
    return st,dt,ins,bey
files=sorted(glob.glob('/root/psp/contall/ta_*.bin'))
lo,hi=int(sys.argv[1]),int(sys.argv[2])
oobs=[]; void=0; exercised=0
SKIP=('ta_19-0-0-8_7e8be282','ta_16-0-0-9_9e9044cf')   # known wedge triggers (FP cmd2/7, AUC cmd6)
for p in files[lo:hi]:
    if any(k in p for k in SKIP): print('%-34s skipped (known wedge)'%p.rsplit('/',1)[1], flush=True); continue
    nm=p.rsplit('/',1)[1]
    st,sid,shm=load(open(p,'rb').read())
    if st!=0 or sid is None:
        print('%-34s LOAD %#x'%(nm,st), flush=True); continue
    tot=0; bey_all=[]
    for cid in range(int(os.environ.get('CLO','0')), int(os.environ.get('CHI','6'))):
        s,dt,ins,bey=probe(sid,shm,cid)
        tot+=ins; bey_all+=bey
        if dt>5:
            print('%-34s STALL at cmd %d  <<< WEDGE, aborting batch'%(nm,cid), flush=True)
            print('\nrange %d..%d aborted at %s'%(lo,hi,nm), flush=True); sys.exit(3)
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id=sid
    try: submit(c)
    except Exception: pass
    if bey_all: oobs.append((nm,len(bey_all),bey_all[:3])); print('%-34s OOB %d'%(nm,len(bey_all)), flush=True)
    elif tot: exercised+=1
    else: void+=1
print('\nrange %d..%d : exercised=%d void=%d OOB=%s'%(lo,hi,exercised,void,oobs[:3] or 'none'), flush=True)