#!/usr/bin/env python3
"""One-shot non-ROP params-VA discriminator for GCP session ID 2.

Run ONLY immediately after the successful SID1 adjacent-word test left ASD
loaded and GCP unloaded. No selector-4, pivot or DMA control writes. Selector
11 may emit MP1 PMI3 diagnostic events. Preserve this boot: no ATX beforehand.
A fault or unexpected result stops; ATX recovery is external.
"""
import ast
import glob
import hashlib
import json
import os
import struct
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

# Extract explicitly whitelisted pure CMAC functions and seven probe helpers,
# never execute either old script's top-level hardware requests.
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)
        elif 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 {n.name for n in body if isinstance(n,ast.FunctionDef)}==set(names)
    assert len(body)==len(names)+int(include_sb)
    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)
helper=extract('/root/psp/hw/gcp_fe0a_params_va_nonrop.py',
    '110650cc738fa3301594c8060d7c26d1757bfed0968cc96e2e50417cfaf26073',
    ('alloc','submit','load','build','invoke','setup_handle','diagnostic_buf'))
# Compiled functions' globals are the helper dict; resolve cross-function names.
helper.update(pure)
helper.update({'am':am,'Device':Device,'glob':glob,'struct':struct,'BUFLEN':0x40000,'PAT':0xA5})
g=json.load(open('/root/psp/hw/fe08_good.json'))
helper.update({'ct':bytes.fromhex(g['ct']),'sk':bytes.fromhex(g['sk']),
               'seedoff':g['seedoff']})
assert pure['cmac'](helper['sk'],bytes([1,0,0,0])+bytes(range(4,0x20))).hex()==g['mac']
print('CMAC self-test OK; no key material printed',flush=True)
marker='/root/psp/hw/.gcp_fe0a_params_va_sid2.attempted'
fd=os.open(marker,os.O_WRONLY|os.O_CREAT|os.O_EXCL,0o600);os.close(fd)

adev=Device['AMD'].iface.dev_impl
psp=adev.psp
helper.update({'adev':adev,'psp':psp})
alloc,load,invoke,setup_handle,diagnostic_buf=(helper[k] for k in
    ('alloc','load','invoke','setup_handle','diagnostic_buf'))
# ASD from the preceding successful test remains loaded; re-loading would
# change the experiment. Only a fresh GCP load is issued here.
app,app_mc=alloc(0x80000)
cmd,cmd_mc=alloc(0x40000)
helper.update({'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)
helper['sid']=sid
print('GCP load',hex(status),'sid',sid,flush=True)
assert status==0 and sid==2, 'not the expected SID2 mapping scenario'

last_ring_ok=False
try:
    # Counter=0x4604a120; in-buffer word 0x60000000 must pass the same
    # 0xffffffff declared-length / diagnostic-selector gates.
    helper['SEED']=0x4604A121
    h=setup_handle()
    last_ring_ok=False
    ring,ta,_=invoke(0xFE0A,diagnostic_buf(h,0,0x60000000))
    last_ring_ok=ring==0
    print('in-buffer gate ring=%#x status=%#x'%(ring,ta),flush=True)
    assert ring==0 and ta==0x241, 'gate control failed, no OOB probe'
    # The signed .77 dword at 0x404a78 is 0x4604a121, unique in its ELF.
    # Old-base displacement reads it if params=0x04024000; shifted-base
    # displacement reads it if params=0x04035000. Only follow the alternative
    # after a clean old-base rejection and a working in-buffer control.
    helper['SEED']=0x4604A121
    last_ring_ok=False
    h=setup_handle()
    last_ring_ok=False
    ring,ta,_=invoke(0xFE0A,diagnostic_buf(h,0xFC3E0A64,0x11111111))
    last_ring_ok=ring==0
    print('old-base threshold counter=0x4604a120 ring=%#x status=%#x'%(ring,ta),flush=True)
    assert ring==0 and ta in (0x100,0x241), 'unexpected old-base result; stop'
    if ta==0x241:
        disp=0xFC3E0A64; label='OLD_BASE'
    else:
        disp=0xFC3CFA64; label='SHIFTED_BASE'
        helper['SEED']=0x4604A121
        last_ring_ok=False
        h=setup_handle()
        last_ring_ok=False
        ring,ta,_=invoke(0xFE0A,diagnostic_buf(h,disp,0x11111111))
        last_ring_ok=ring==0
        print('shifted-base threshold counter=0x4604a120 ring=%#x status=%#x'%(ring,ta),flush=True)
        assert ring==0 and ta==0x241, 'shifted-base did not match; stop'
    # Exact equality: the boundary counter must reject the same target.
    helper['SEED']=0x4604A122
    last_ring_ok=False
    h=setup_handle()
    last_ring_ok=False
    ring,ta,_=invoke(0xFE0A,diagnostic_buf(h,disp,0x11111111))
    last_ring_ok=ring==0
    print(label,'equality boundary counter=0x4604a121 ring=%#x status=%#x'%(ring,ta),flush=True)
    assert ring==0 and ta==0x100, 'second threshold did not confirm the unique word'
    print(label,'matches unique signed codeword 0x4604a121 in this SID2 boot;',
          'this is still a source read, not an independently verified destination write',flush=True)
finally:
    # On an exception, don't pile an unload onto a potentially faulted ring.
    if last_ring_ok:
        try:
            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)
        except Exception as exc:
            print('GCP unload unavailable; ATX needed:',type(exc).__name__,flush=True)