#!/usr/bin/env python3
"""Do TA builds rejected by 49.0.54 load on rolled-back firmware?

The corpus holds 78 distinct TA builds; 63 load on the shipping firmware and the rest are
refused. Those refusals are a property of the *running* firmware's loader, and the rollback
finding means the attacker chooses that firmware. Any build that loads only on an older set is
attack surface this audit has never exercised, since every TA fuzz ran against the builds
49.0.54 accepts.

  argv[1]: container path to roll back to, or 'identity' for the shipping firmware
  argv[2], argv[3]: index range into the sorted TA list (session wall is ~28 loads per boot)

Every successful load is unloaded immediately so the session wall is not hit; a row showing
status 0xb is the wall, not a refusal, and is labelled as such.
"""
import ctypes, glob, os, sys
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

def load_container(path):
    b = open(path, 'rb').read()
    h = H.from_buffer_copy(b)
    uoff = h.header.ucode_array_offset_bytes
    out = {}
    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):
            out[d.fw_type] = b[s:s + d.size_bytes]
    return out

target = sys.argv[1]
lo = int(sys.argv[2]) if len(sys.argv) > 2 else 0
hi = int(sys.argv[3]) if len(sys.argv) > 3 else 24
newset = None if target == 'identity' else load_container(target)
orig = amdev.AMFirmware.__init__
def patched(self, adev):
    orig(self, adev)
    if newset is None:
        print('FW identity', flush=True); return
    for t, blob in newset.items():
        if t in self.sos_fw:
            self.sos_fw[t] = memoryview(bytearray(blob))
    print('FW substituted %d components' % len(newset), flush=True)
amdev.AMFirmware.__init__ = patched

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

def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
    return pa, adev.paddr2mc(pa)
files = sorted(glob.glob('/root/psp/contall/ta_*'))
print('TACOUNT %d range %d..%d' % (len(files), lo, hi), flush=True)
ap_, amc = alloc(0x200000); cp, cmc = alloc(0x8000)
for path in files[lo:hi]:
    blob = open(path, 'rb').read()
    n = (len(blob) + 3) & ~3
    if n > 0x200000:
        print('  TA %-42s SKIP_TOO_BIG' % os.path.basename(path)[:42], flush=True); continue
    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 = 0x4000
    try:
        r = psp._ring_submit(c)
        st = r.resp.status & 0xFFFFFFFF; sid = r.resp.session_id
    except Exception as e:
        try: st = int(str(e).strip().split()[-1]) & 0xFFFFFFFF
        except Exception: st = -1
        sid = None
    tag = ' SESSION_WALL' if st == 0xb else ''
    print('  TA %-42s status=%#-11x sid=%s%s'
          % (os.path.basename(path)[:42], st, sid, tag), flush=True)
    if st == 0 and sid is not None:
        u = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        u.cmd.cmd_unload_ta.session_id = sid
        try: psp._ring_submit(u)
        except Exception: pass