#!/usr/bin/env python3
"""One-shot signed GCP diagnostic event versus host-only PMI3 neighbor readback.

This is a CONTROL for candidate fabric target 0x03B30600, not an ASD 0x403
write. Only FE0A selector 11's signed diagnostic messages are issued. No
read/write of the candidate data register itself; host SMN INDEX2 writes
select the read address only. No repeat of the quarantined ROP chain.
"""
import ast
import glob
import hashlib
import json
import os
import struct
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

def extract(path,sha,names,include_sb=False):
    src=open(path,'rb').read()
    assert hashlib.sha256(src).hexdigest()==sha
    body=[]
    for n in ast.parse(src.decode()).body:
        if isinstance(n,ast.FunctionDef) and n.name in names: body.append(n)
        if include_sb and isinstance(n,ast.Assign) and len(n.targets)==1 and isinstance(n.targets[0],ast.Name) and n.targets[0].id=='SB': body.append(n)
    assert len(body)==len(names)+int(include_sb)
    assert {n.name for n in body if isinstance(n,ast.FunctionDef)}==set(names)
    ns={}
    exec(compile(ast.fix_missing_locations(ast.Module(body=body,type_ignores=[])),path,'exec'),ns)
    return ns
pure=extract('/root/psp/hw/gcp_fe0a_sel8_probe.py',
    '7e0a36457d7841e027d94ca42742ab5702e325b6ce719d1e261149173d9e64a9',
    ('xt','expand','enc','dbl','cmac'),True)
h=extract('/root/psp/hw/gcp_fe0a_params_va_nonrop.py',
    '110650cc738fa3301594c8060d7c26d1757bfed0968cc96e2e50417cfaf26073',
    ('alloc','submit','load','build','invoke','setup_handle','diagnostic_buf'))
h.update(pure);h.update({'am':am,'glob':glob,'struct':struct,'BUFLEN':0x40000,'PAT':0xA5,'SEED':1})
g=json.load(open('/root/psp/hw/fe08_good.json'))
h.update({'ct':bytes.fromhex(g['ct']),'sk':bytes.fromhex(g['sk']),'seedoff':g['seedoff']})
assert pure['cmac'](h['sk'],bytes([1,0,0,0])+bytes(range(4,0x20))).hex()==g['mac']
print('CMAC self-test OK; key material not printed',flush=True)
fd=os.open('/root/psp/hw/.gcp_pmi3_signed_event_control.attempted',os.O_WRONLY|os.O_CREAT|os.O_EXCL,0o600);os.close(fd)
adev=Device['AMD'].iface.dev_impl
psp=adev.psp
h.update({'adev':adev,'psp':psp})
alloc,load,invoke,setup_handle,diagnostic_buf=(h[k] for k in ('alloc','load','invoke','setup_handle','diagnostic_buf'))

def smn(a):
    adev.wreg(14,a) # PCIE_INDEX2, numeric read selector only
    return adev.rreg(15)&0xffffffff

def snapshot(tag):
    vals=tuple(smn(a) for a in (0x03B30604,0x03B30608,0x03B3060C))
    print(tag,'PMI3 adjacent status,rptr,wptr',tuple(hex(v) for v in vals),
          'MP0',hex(smn(0x03810ACC)),'bogus',hex(smn(0x0FEEDC00)),flush=True)
    return vals

# ASD is still loaded after the preceding signed ASD-first run. GCP TA
# unload was independently verified zero; do not disturb ASD or ATX.
app,app_mc=alloc(0x80000)
cmd,cmd_mc=alloc(0x40000)
h['cmd']=cmd
gcp_blob=open(glob.glob('/root/psp/contall/ta_22-0-0-77*')[0],'rb').read()
status,sid=load(gcp_blob,am.GFX_CMD_ID_LOAD_TA,app,app_mc,cmd_mc,0x40000)
h['sid']=sid
print('GCP load',hex(status),'sid',sid,flush=True)
assert status==0 and sid
ring_ok=False
try:
    handle=setup_handle()
    before=snapshot('BEFORE')
    b=diagnostic_buf(handle,0,0x60000000)
    struct.pack_into('<I',b,0x414,0x1000) # ordinary in-buffer declared length
    ring,ta,_=invoke(0xFE0A,b)
    ring_ok=ring==0
    print('signed selector-11 FE0A ring=%#x ta=%#x'%(ring,ta),flush=True)
    assert ring==0 and ta==0x241
    after=snapshot('AFTER')
    print('adjacent_changed',before!=after,flush=True)
finally:
    if ring_ok:
        u=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        u.cmd.cmd_unload_ta.session_id=sid
        ru=psp._ring_submit(u)
        print('GCP unload',hex(ru.resp.status&0xffffffff),flush=True)