"""HDCP cmd 0x12 -> ASD cmd 0x308 reachability probe.

hdcp_invoke_entry case 0x12 -> sub_403FD0(shm+0xC, shm+0x1410)
    R4 = *(shm+0xC)   -> ASD 0x308 a1 = LENGTH   (attacker)
    R1 = shm+0x10     -> ASD 0x308 a2 = BUFFER   (attacker)
NO session gate on this path (top-level switch in hdcp_invoke_entry).

ASD sub_D13340:
    requires 48 <= len <= 0x1400, validate_ptr(buf,len)
    asd_memcpy(v23[5120], buf, len)
    len2 = (b5<<16)|(b6<<8)|b7      <- re-derived FROM OUR COPIED BYTES
    if (len2 < 0x13FC) len2 += 5
    ... gated ...
    sub_D010E0(global, v23, len2)   <- OOB write, len2 up to 0xFFFFFF

This probe ONLY checks reachability and response codes with benign len2 values.
"""
import ctypes, struct, time
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32

d=Device["AMD"]; adev=d.iface.dev_impl; psp=adev.psp
SZ=ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
def stat(): return int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr,SZ)[:])[864:868],"little")
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True); return pa,adev.paddr2mc(pa)
P='/Users/user/dev/amd-psp-re/psp_out/'
def load(cmd,blob):
    ap,amc=alloc(0x40000); cp,cmc=alloc(0x8000)
    adev.vram.view(ap,len(blob))[:]=blob; adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr,SZ)[:]=b"\x00"*SZ
    c=am.struct_psp_gfx_cmd_resp(cmd_id=cmd); t=c.cmd.cmd_load_ta
    t.app_phy_addr_hi,t.app_phy_addr_lo=hi32(amc),lo32(amc); t.app_len=len(blob)
    t.cmd_buf_phy_addr_hi,t.cmd_buf_phy_addr_lo=hi32(cmc),lo32(cmc); t.cmd_buf_len=0x8000
    r=psp._ring_submit(c); return r.resp.session_id,cp
load(am.GFX_CMD_ID_LOAD_ASD, open(P+'asd_full.bin','rb').read())
hdcp_sid,hdcp_shm= load(am.GFX_CMD_ID_LOAD_TA, open(P+'hdcp_full.bin','rb').read())
print("HDCP sid=%#x  (NO topology / session needed for cmd 0x12)\n" % hdcp_sid)

def inv(buf,tc):
    adev.vram.view(hdcp_shm,0x8000)[:]=bytes(buf); adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr,SZ)[:]=b"\x00"*SZ
    ic=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD); iv=ic.cmd.cmd_invoke_cmd
    iv.session_id,iv.ta_cmd_id=hdcp_sid,tc
    t0=time.time()
    try: psp._ring_submit(ic); return 0, bytes(adev.vram.view(hdcp_shm,0x2000)[:]), time.time()-t0
    except RuntimeError: return stat(), bytes(adev.vram.view(hdcp_shm,0x2000)[:]), time.time()-t0
    except TimeoutError: return "TIMEOUT", None, time.time()-t0

def cmd12(length, payload=None):
    b=bytearray(0x8000)
    struct.pack_into('<I',b,0,0x12)          # top-level command
    struct.pack_into('<I',b,0xC,length)      # -> ASD 0x308 a1 (len)
    if payload: b[0x10:0x10+len(payload)]=payload
    ps,out,dt=inv(b,0x12)
    ta = struct.unpack_from('<I',out,4)[0] if out is not None else None
    return ps,ta,dt

print("%-34s %-9s %-8s %s" % ("case","psp","ta_ret","elapsed"))
# benign bodies: len2 derived from bytes 5,6,7 of the copied data -> keep them small
body = bytearray(b'\x00'*0x1400)
body[0]=0x80          # (v23[0] & 0xF0) == 0x80 branch check
for L in (0, 16, 47, 48, 64, 256, 0x1400, 0x1401, 0x2000):
    ps,ta,dt = cmd12(L, bytes(body[:min(L,0x1400)]) if L else None)
    note=""
    if L < 48: note="(below min 48 -> expect reject)"
    elif L > 0x1400: note="(above max 0x1400 -> expect reject)"
    print("%-34s %-9s %-8s %6.2fs %s" % ("len=%#x"%L, ps, ta, dt, note))
    if not psp.is_sos_alive(): print("  sOS DIED"); break
print("\nsOS alive:", psp.is_sos_alive())