#!/usr/bin/env python3
"""Command-surface fuzz across the twelve TAs, on silicon.

Everything before this was static or emulated. The emulator was shown to produce false
positives for these TAs (one fault refuted on hardware), so this runs the real thing.

For each TA: LOAD, then for each command id, invoke with several payload profiles and record
(status, first 32 bytes out). Flags anything that is not one of the two benign outcomes
(0x0 accepted / 0xffff0000 id rejected), plus any case where a TA's response to a FIXED probe
command changes after some other command ran -- that is the signature of state corruption.

Controls:
  * a benign profile must produce the same statuses as the six-TA differential already
    recorded (OC 0->0, AUC 0->0xffff0000, ...); if not, the harness is wrong;
  * the fixed re-probe after each command detects harness-induced drift too.

Safety: ordinary INVOKEs. XGMI is the only TA that emits SYS_DRV opcode 0xd and it is
included -- but with a zeroed staging path the 0xd handler takes its bounded branch, which is
what the emulation and the static reading both show.
"""
import argparse, glob, os, struct, sys, time

PROFILES = {
    'zero':  lambda b: None,
    'big':   lambda b: [struct.pack_into('<I', b, o, 0x00000800) for o in range(4, 0x80, 4)],
    'ones':  lambda b: [struct.pack_into('<I', b, o, 0xFFFFFFFF) for o in range(4, 0x80, 4)],
    'ptrish':lambda b: [struct.pack_into('<I', b, o, 0x7F000000 + o) for o in range(4, 0x80, 4)],
}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--dir', default='/root/psp/contall')
    ap.add_argument('--max-cmd', type=int, default=16)
    ap.add_argument('--min-cmd', type=int, default=0)
    ap.add_argument('--only', default='')
    a = ap.parse_args()
    from tinygrad import Device
    import tinygrad.runtime.autogen.am.am as am
    adev = Device["AMD"].iface.dev_impl; psp = adev.psp
    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(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=0x8000
        try:
            r=psp._ring_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 invoke(sid, shm, cid, prof):
        buf = bytearray(0x2000); struct.pack_into('<I', buf, 0, cid)
        PROFILES[prof](buf)
        adev.vram.view(shm, len(buf))[:] = 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
        try:
            r=psp._ring_submit(c); st=r.resp.status&0xFFFFFFFF
        except Exception as e:
            try: st=int(str(e).strip().split()[-1])&0xFFFFFFFF
            except Exception: st=-1
        adev.gmc.flush_hdp()
        return st, bytes(adev.vram.view(shm,32)[:])
    def unload(sid):
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id=sid
        try: return psp._ring_submit(c).resp.status&0xFFFFFFFF
        except Exception: return -1

    TAS = [('AUC','ta_16-0-0-9_9e9044cf'),('CSU','ta_42-0-0-6_1c874c3c'),('FP','ta_19-0-0-8_7e8be282'),
           ('GCP','ta_22-0-0-77_7c426e7f'),('MERIT','ta_24-0-0-70_5cd4afda'),('OC','ta_50-0-0-1_cbd47c59'),
           ('OPM','ta_25-0-0-81_dc596ba6'),('POWERPLAY','ta_34-0-0-20_e391bdbb'),('RAP','ta_7-0-2-19_f368676e'),
           ('TMZ','ta_29-0-0-4_ce34f538'),('WSPE','ta_30-0-0-38_ec5dee6d'),('XGMI','ta_32-0-0-20_20a3d1a4')]
    if a.only: TAS=[t for t in TAS if t[0] in a.only.split(',')]
    BENIGN = {0x0, 0xffff0000}
    anomalies=[]
    for name, stem in TAS:
        cands = glob.glob(os.path.join(a.dir, stem+'*'))
        if not cands: print('%-10s no image'%name); continue
        blob=open(cands[0],'rb').read()
        st, sid, shm = load(blob)
        if st!=0 or sid is None:
            print('%-10s LOAD %#x'%(name,st)); continue
        ref_st, ref_out = invoke(sid, shm, 0, 'zero')     # fixed re-probe reference
        odd=[]; drift=[]
        for cid in range(a.min_cmd, a.max_cmd):
            for prof in PROFILES:
                t0=time.time(); s2,out = invoke(sid, shm, cid, prof); dt=time.time()-t0
                if s2 not in BENIGN: odd.append((cid,prof,hex(s2)))
                if dt > 2.0: odd.append((cid,prof,'SLOW %.1fs'%dt))
            r2, rout = invoke(sid, shm, 0, 'zero')
            if (r2, rout) != (ref_st, ref_out): drift.append((cid, hex(r2), rout[:8].hex()))
        unload(sid)
        print('%-10s odd=%-3d drift=%-3d %s' % (name, len(odd), len(drift),
              (odd[:3] if odd else '') if odd else ('DRIFT %s'%drift[:3] if drift else 'clean')))
        if odd or drift: anomalies.append((name, odd, drift))
    print('\nTAs with anomalies: %s' % ([n for n,_,_ in anomalies] or 'NONE'))
    for n,odd,drift in anomalies:
        print('  %s odd(%d): %s'%(n,len(odd),odd[:8]))
        print('  %s drift(%d): %s'%(n,len(drift),drift[:8]))
    return 0

if __name__=='__main__': sys.exit(main())