#!/usr/bin/env python3
"""Roll back the WHOLE PSP firmware set to an older signed container, then exercise it.

The historical `psp_13_0_0_sos.bin` containers hold all six components, not just SYS_DRV:
SYS_DRV, SOS, TOC, KDB, SPL, RL.

sOS implements the host-facing surface (LOAD_TA, INVOKE, sessions, descriptors), and every
hardware fuzz in this audit ran against sOS 49.0.54 only. Substituting the whole older set
turns that entire validated corpus into a fresh search space.

All six are substituted together so the set stays self-consistent; mixing an old sOS with a
new SYS_DRV would risk attributing a version-mismatch failure to a security gate.

  argv[1]: container path, or 'identity' (re-inject current bytes -- positive control)
  argv[2]: optional 'fuzz' to run an out-of-bounds-write probe against the booted firmware
"""
import ctypes, glob, struct, sys, time, random
sys.path.insert(0, '/root/psp/hw')
import tinygrad.runtime.support.am.amdev as amdev
import tinygrad.runtime.autogen.am.am as am

H = am.struct_psp_firmware_header_v2_0
DS = ctypes.sizeof(am.struct_psp_fw_bin_desc)
FO = H.psp_fw_bin.offset
NAMES = {getattr(am, n): n.replace('PSP_FW_TYPE_PSP_', '')
         for n in dir(am) if n.startswith('PSP_FW_TYPE_PSP_')}

def load_container(path):
    b = open(path, 'rb').read()
    h = H.from_buffer_copy(b)
    uoff = h.header.ucode_array_offset_bytes
    comps = {}
    for i in range(h.psp_fw_bin_count):
        d = am.struct_psp_fw_bin_desc.from_buffer_copy(b, FO + i * DS)
        s = d.offset_bytes + uoff
        if d.size_bytes and s + d.size_bytes <= len(b):
            comps[d.fw_type] = (b[s:s + d.size_bytes], d.fw_version)
    return comps

target = sys.argv[1]
do_fuzz = len(sys.argv) > 2 and sys.argv[2] == 'fuzz'
newset = None if target == 'identity' else load_container(target)
if newset is not None:
    print('CONTAINER %s comps=%d %s' % (target.split('/')[-1], len(newset),
          ' '.join('%s=%#x' % (NAMES.get(k, hex(k)), v[1]) for k, v in sorted(newset.items()))),
          flush=True)

orig = amdev.AMFirmware.__init__
def patched(self, adev):
    orig(self, adev)
    if newset is None:
        print('IDENTITY no substitution', flush=True); return
    subbed = []
    for t, (blob, ver) in newset.items():
        if t in self.sos_fw:
            self.sos_fw[t] = memoryview(bytearray(blob))
            subbed.append('%s=%#x' % (NAMES.get(t, hex(t)), ver))
    print('SUBSTITUTED %s' % ' '.join(subbed), flush=True)
amdev.AMFirmware.__init__ = patched

from tinygrad import Device
adev = Device["AMD"].iface.dev_impl
psp = adev.psp
P = psp.reg_pref
print('ALIVE %s' % bool(psp.is_sos_alive()), flush=True)
print('SOSVER _58=%#x _59=%#x' % (adev.reg(f"{P}_58").read(), adev.reg(f"{P}_59").read()),
      flush=True)

def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
    return pa, adev.paddr2mc(pa)
blob = open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0], 'rb').read()
ap_, amc = alloc(0x80000)
n = (len(blob) + 3) & ~3
adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob)); adev.gmc.flush_hdp()
WIN = 0x4000; TOTAL = 0x8000; PAT = 0xD5D5D5D5
cp, cmc = alloc(TOTAL)
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
try:
    r = psp._ring_submit(c)
    st = r.resp.status & 0xFFFFFFFF; sid = r.resp.session_id
    print('TALOAD status=%#x sid=%s' % (st, sid), flush=True)
except Exception as e:
    print('TALOAD ERR %s %s' % (type(e).__name__, str(e)[:100]), flush=True); sys.exit(1)
if st != 0: sys.exit(0)

def invoke(cid, mutate=None):
    b = bytearray(struct.pack('<I', PAT) * (WIN // 4)); struct.pack_into('<I', b, 0, cid)
    if mutate: mutate(b)
    adev.vram.view(cp, WIN)[:] = bytes(b)
    adev.vram.view(cp + WIN, TOTAL - WIN)[:] = bytes(
        struct.pack('<I', PAT) * ((TOTAL - WIN) // 4))
    adev.gmc.flush_hdp()
    iv = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
    iv.cmd.cmd_invoke_cmd.session_id = sid; iv.cmd.cmd_invoke_cmd.ta_cmd_id = cid
    t0 = time.time()
    try: s = psp._ring_submit(iv).resp.status & 0xFFFFFFFF
    except Exception as e:
        try: s = int(str(e).strip().split()[-1]) & 0xFFFFFFFF
        except Exception: s = -1
    dt = time.time() - t0
    adev.gmc.flush_hdp()
    tail = bytes(adev.vram.view(cp + WIN, TOTAL - WIN)[:])
    oob = [o for o in range(0, len(tail), 4) if struct.unpack_from('<I', tail, o)[0] != PAT]
    return s, dt, oob

s, dt, oob = invoke(1)
print('TAINVOKE status=%#x oob=%d SYSDRV_SERVICING=%s' % (s, len(oob), s == 0), flush=True)
if do_fuzz:
    rng = random.Random(0x5011)
    hits = []
    for i in range(600):
        cid = rng.randrange(0, 12)
        def m(b, rr=rng):
            for _ in range(rr.randrange(1, 5)):
                struct.pack_into('<I', b, rr.randrange(1, WIN // 4) * 4,
                                 rr.choice([0, 1, 0xFFFFFFFF, 0x1000, 0x8c, 0x500]))
        s2, dt2, oob2 = invoke(cid, m)
        if oob2:
            hits.append(('OOB', i, cid, [hex(x) for x in oob2[:3]])); break
        if dt2 > 5:
            hits.append(('STALL', i, cid)); break
    print('FUZZ 600 invokes rolled-back: %s' % (hits if hits else 'no OOB no stall'), flush=True)