#!/usr/bin/env python3
"""Does a host INVOKE payload change what the OC TA sends to SYS_DRV opcodes 0x33/3?

Static conclusion (AUDIT-TA-MEMSAFETY-HUNT pass 60-62, REPORT-SYSDRV-COUNT-UNDERFLOW):
`oc_TA_InvokeCommandEntryPoint__h2` builds the SYS_DRV descriptor from a firmware GLOBAL with
a CONSTANT length 0x30, so the host cannot influence the length or content that reaches the
staging buffer. Therefore the count field at buf+0x88 (offset 140) is unreachable.

Prediction if that reading is correct:
    OC's INVOKE status is INDEPENDENT of the host shared-buffer contents -- filling the buffer
    with large size fields and patterns changes nothing.

Overturn condition:
    a status differing between profiles, or a fault/hang, would mean host data does reach that
    path and the static reading is wrong.

Safety: OC does not emit SYS_DRV opcode 0xd, so the underflow trigger is never issued here.
This only sends ordinary INVOKEs.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad ./oc_payload_probe.py --dir /root/psp/contall
"""
import argparse, glob, os, struct, sys

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--dir', default='/root/psp/contall')
    ap.add_argument('--cmd-ids', default='0,1,2,3')
    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, cmd_id, profile):
        buf = bytearray(0x2000)
        struct.pack_into('<I', buf, 0, cmd_id)
        if profile == 'big':
            # plausible {ptr,len} and size fields all set large, plus a marker at 0x88
            for off in range(4, 0x80, 4):
                struct.pack_into('<I', buf, off, 0x00000800)
            struct.pack_into('<I', buf, 0x88, 0xCAFEF00D)
            for off in range(0x100, 0x400, 4):
                struct.pack_into('<I', buf, off, 0x41414141)
        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, cmd_id
        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, 16)[:])

    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 = [('OC','ta_50-0-0-1_cbd47c59'), ('AUC','ta_16-0-0-9_9e9044cf'),
           ('FP','ta_19-0-0-8_7e8be282'), ('OPM','ta_25-0-0-81_dc596ba6'),
           ('MERIT','ta_24-0-0-70_5cd4afda'), ('GCP','ta_22-0-0-77_7c426e7f')]
    ids = [int(x,0) for x in a.cmd_ids.split(',')]
    overall = []
    for name, stem in TAS:
      cands = glob.glob(os.path.join(a.dir, stem + '*'))
      if not cands:
        print('%-6s image not found' % name); continue
      blob = open(cands[0], 'rb').read()
      print('%s: %s (%d bytes)' % (name, os.path.basename(cands[0]), len(blob)))
      results = {}
      for profile in ('zero', 'big'):
        st, sid, shm = load(blob)
        if st != 0 or sid is None:
            print('   LOAD failed (%s): %#x' % (profile, st)); continue
        row = []
        for cid in ids:
            s2, out = invoke(sid, shm, cid, profile)
            row.append((cid, s2, out[:8].hex()))
        unload(sid)
        results[profile] = row
        print('   %-5s: %s' % (profile, ['%#x->%#x' % (c, s2) for c, s2, _ in row]))
      if 'zero' in results and 'big' in results:
        pairs = list(zip(results['zero'], results['big']))
        diff = [(z[0], z[1], bg[1]) for z, bg in pairs if z[1] != bg[1]]
        dout = [(z[0], z[2], bg[2]) for z, bg in pairs if z[2] != bg[2]]
        print('   identical statuses: %s ; identical outputs: %s'
              % ('YES' if not diff else 'NO %s' % diff, 'YES' if not dout else 'NO %s' % dout))
        overall.append((name, not diff, not dout))
    print('\nSUMMARY (host payload must NOT change behaviour if the static reading holds):')
    for n, sdok, outok in overall:
        print('   %-6s statuses same: %-3s  outputs same: %s' % (n, 'YES' if sdok else 'NO', 'YES' if outok else 'NO'))
    bad = [n for n, sdok, outok in overall if not (sdok and outok)]
    print('\n%s' % ('All TAs insensitive to host payload -> corroborates non-reachability.'
                    if not bad else 'SENSITIVE: %s -> static reading overturned for these.' % bad))
    return 0

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