#!/usr/bin/env python3
"""Drive GCP's state machine 0 -> 1 -> 2 and watch the status codes.

Every FE09 invoke so far returned 560 ("wrong state") because we never allocated a DRM slot
first: gcp_cmd01 is the 0 -> 1 transition and its only gate is `arg10 <= 1` (fw/...FE09 doc
Addendum 11). So the sequence to test is:

  step 1  cmd 0x01 (Allocate)  -> expect status 0, and the SLOT INDEX written to params+0x1950
  step 2  cmd 0xFE08           -> 1 -> 2, gated on ctx_slot state==1 and ctx_slot+772 != 1,
                                  then a GP key-object chain built from the host's 256-byte blob
                                  at params+0x14 (gp_object_allocate / populate_attrs /
                                  operation_allocate / set_operation_key / digest_do_final).
                                  Expect 0 on success; 609/610/611/1296..1300 identify which
                                  stage refused.
  step 3  cmd 0xFE09           -> if the slot reached state 2 the status changes from 560.
                                  560 means still-not-2; anything else means we moved.

Status codes are the oracle, all host-visible, and nothing here performs an out-of-bounds write
(the FE09 displacement is aimed inside the buffer), so no fault and no power cycle is needed --
the probe can iterate.

HIBYTE(arg10) selects the slot and arg10 must be < 0x10000000, so only slots 0..15 are
addressable even though cmd01 scans 32.
"""
import struct, sys, glob
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
TA = 'ta_22-0-4-50_81b082b4'
BUFLEN = 0x8000
PAT = 0xA5
ERRS = {0: 'OK', 256: '256 bad-arg', 560: '560 WRONG STATE (not 2)', 561: '561', 562: '562 OMAC/verify',
        609: '609 slot+772==1', 610: '610 state!=1', 611: '611 stage refused',
        1049: '1049 all slots busy', 1050: '1050 arg10>1',
        1296: '1296 object_allocate', 1297: '1297 populate_attrs',
        1298: '1298 operation_allocate', 1299: '1299 set_operation_key',
        1300: '1300 digest_do_final'}

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 alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
    return pa, adev.paddr2mc(pa)

app_pa, app_mc = alloc(0x80000)
cand = glob.glob('/root/psp/contall/%s*' % TA)
if not cand:
    raise SystemExit('missing 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:
        return -1, None

def invoke(sid, cmd, arg10=0, blob256=None, disp=None, extra=None):
    """One GCP TA command. Returns (ring_status, ta_status_word, out1950, raw)."""
    buf = bytearray(bytes([PAT]) * BUFLEN)
    struct.pack_into('<I', buf, 0x000, cmd)          # the dispatcher switches on params[0]
    struct.pack_into('<I', buf, 0x008, 0)            # status slot, cleared
    struct.pack_into('<I', buf, 0x010, arg10)        # arg10: HIBYTE selects the slot
    struct.pack_into('<I', buf, 0x414, 0xFFFFFFFB)
    struct.pack_into('<I', buf, 0x818, 0xFFFFFFFB)
    if blob256 is not None:
        buf[0x014:0x014 + len(blob256)] = blob256    # FE08's key blob lives at params+0x14
    if disp is not None:
        struct.pack_into('<I', buf, 0x820, 1)
        struct.pack_into('<I', buf, 0x828, disp & 0xFFFFFFFF)
    for off, val in (extra or []):
        struct.pack_into('<I', buf, off, val & 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, cmd
    try:
        rs = psp._ring_submit(c).resp.status & 0xFFFFFFFF
    except Exception:
        rs = -1
    adev.gmc.flush_hdp()
    raw = bytes(adev.vram.view(cp, BUFLEN)[:])
    ta_st = struct.unpack_from('<I', raw, 8)[0]
    out = struct.unpack_from('<I', raw, 0x1950)[0]
    return rs, ta_st, out, raw

def show(tag, rs, ta, out):
    print('   %-26s ring=%-5s ta_status=%#-8x %-26s out1950=%#x'
          % (tag, ('FAULT' if rs == -1 else hex(rs)), ta, ERRS.get(ta, '?'), out), flush=True)

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

SAFE_D = (PARAMS_VA + 0x300 - BASE) & 0xFFFFFFFF    # FE09 store lands inside the buffer

print('\n=== baseline: FE09 on an UN-allocated slot (expect 560) ===', flush=True)
rs, ta, out, _ = invoke(sid, 0xFE09, arg10=0, disp=SAFE_D)
show('FE09 (no cmd01)', rs, ta, out)

print('\n=== step 1: cmd 0x01 Allocate (expect 0, slot index at +0x1950) ===', flush=True)
rs, ta, out, raw = invoke(sid, 0x01, arg10=0)
show('cmd01 arg10=0', rs, ta, out)
slot = out if ta == 0 else None
if slot is None or slot > 15:
    print('   cmd01 did not yield a usable slot (need 0..15 for HIBYTE addressing)', flush=True)

print('\n=== step 2: FE08 on that slot, several blob shapes (1 -> 2) ===', flush=True)
if slot is not None and slot <= 15:
    a10 = (slot & 0xF) << 24
    shapes = [('zeros', bytes(256)),
              ('0xFF', b'\xff' * 256),
              ('incrementing', bytes(range(256))),
              ('rsa2048-ish hdr', struct.pack('<10I', 0,0,0,0,0,0,0,0,0,2048) + bytes(256-40))]
    for nm, bl in shapes:
        rs, ta, out, _ = invoke(sid, 0xFE08, arg10=a10, blob256=bl)
        show('FE08 %s' % nm, rs, ta, out)
        if ta == 0:
            print('      *** FE08 SUCCEEDED -> slot should now be in state 2 ***', flush=True)
            break

    print('\n=== step 3: FE09 again on the same slot (560 => still not state 2) ===', flush=True)
    rs, ta, out, _ = invoke(sid, 0xFE09, arg10=a10, disp=SAFE_D)
    show('FE09 after FE08', rs, ta, out)
    if ta != 560:
        print('      *** status moved off 560 -- state advanced ***', flush=True)

print('\nring: %s' % ('WEDGED' if frame() != f0 else 'alive'), flush=True)