#!/usr/bin/env python3
"""One-shot non-ROP equality check for signed GCP .77 word at 0x404A74.

No selector-4 / pivot / LSDMA / MP1 security register request. FE0A diagnostic
may issue MP1 PMI3 event writes. Three signed auth sessions in one clean boot:
local positive gate control, then two thresholds bracketing one code dword.
A fault or unexpected result aborts the sequence; 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_exact.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_app,asd_mc=alloc(0x80000)
asd_cmd,asd_cmd_mc=alloc(0x10000)
app,app_mc=alloc(0x80000)
cmd,cmd_mc=alloc(0x40000)
helper.update({'cmd':cmd})
asd_blob=open('/root/psp/contall/asd_33-0-0-163_18debcc4.bin','rb').read()
gcp_blob=open(glob.glob('/root/psp/contall/ta_22-0-0-77*')[0],'rb').read()
a_status,_=load(asd_blob,am.GFX_CMD_ID_LOAD_ASD,asd_app,asd_mc,asd_cmd_mc,0x10000)
print('ASD load',hex(a_status),flush=True)
assert a_status==0
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

last_ring_ok=False
try:
    # Counter=0x460db5af; in-buffer word 0x60000000 must pass the same
    # 0xffffffff declared-length / diagnostic-selector gates.
    helper['SEED']=0x460DB5B0
    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 0x404a74 is 0x460db5b0. These two
    # comparisons distinguish exact equality, not merely >0x40000000.
    for seed,expected in ((0x460DB5B0,0x241),(0x460DB5B1,0x100)):
        helper['SEED']=seed
        last_ring_ok=False
        h=setup_handle()
        print('ONE comparator threshold counter=%#x'%(seed-1),flush=True)
        last_ring_ok=False
        ring,ta,_=invoke(0xFE0A,diagnostic_buf(h,0xFC3E0A60,0x11111111))
        last_ring_ok=ring==0
        print('comparator ring=%#x status=%#x expected=%#x'%(ring,ta,expected),flush=True)
        assert ring==0 and ta==expected, 'nonmatching/faulted threshold: stop, do not infer params VA'
    print('two thresholds match exact signed codeword 0x460db5b0 in this boot;',
          'this does not measure the params VA in another boot/session',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)