#!/usr/bin/env python3
"""Can the host actually INVOKE the twelve newly-discovered TAs, or only load them?

`fw/AUDIT-NEW-TAS.md` establishes that this Navi 31 authenticates twelve TA types the
project never audited, and that all twelve carry class `TA` (the invokable kind) rather
than ASD's `DR`. That is an inference from the image plus a successful load. ASD is the
standing precedent for why that is not enough: it loads cleanly and then refuses **every**
`INVOKE_CMD` with `psp_status 0x4`
(`FINDING-ASD-IS-DRIVER-CLASS-NOT-INVOKABLE.md`).

So this measures it. For each TA: LOAD_TA -> INVOKE_CMD(ta_cmd_id=0) -> UNLOAD_TA.

Command id 0 is deliberate. The goal is to separate "the PSP will not dispatch to this TA
at all" from "the TA received a command and answered", which the status alone shows. An
unimplemented id is the least invasive way to ask; this is not a command-surface fuzz and
it does not try to make any TA do work.

Controls, both required:
  POS  HDCP from this machine's own psp_13_0_0_ta.bin -- must invoke (not 0x4).
  NEG  ASD, driver class -- must refuse with 0x4, reproducing the known behaviour.
       Without it, a run where everything returns 0x4 would look like a finding
       instead of a broken harness.

Successful loads allocate a session and this rig hits a session wall around 45 loads per
boot, so run under AM_RESET=1; the set here is ~15 loads.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad invoke_newtas_probe.py --dir /root/psp/contall
"""
import argparse
import glob
import os
import struct
import sys

NEW = ['ta_16-0-0-9_9e9044cf', 'ta_42-0-0-6_1c874c3c', 'ta_19-0-0-8_7e8be282',
       'ta_22-0-0-77_7c426e7f', 'ta_24-0-0-70_5cd4afda', 'ta_50-0-0-1_cbd47c59',
       'ta_25-0-0-81_dc596ba6', 'ta_34-0-0-20_e391bdbb', 'ta_7-0-2-19_f368676e',
       'ta_29-0-0-4_ce34f538', 'ta_30-0-0-38_ec5dee6d', 'ta_32-0-0-20_20a3d1a4']
NAMES = dict(zip(NEW, ['AUC', 'CSU', 'FP', 'GCP', 'MERIT', 'OC',
                       'OPM', 'POWERPLAY', 'RAP', 'TMZ', 'WSPE', 'XGMI']))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--dir', default='/root/psp/contall')
    ap.add_argument('--ta-bin', default='/lib/firmware/amdgpu/psp_13_0_0_ta.bin')
    ap.add_argument('--asd', default='/root/psp/asd_full.bin')
    ap.add_argument('--cmd-ids', default='0,1,2',
                    help='ids to try; a TA that answers ANY of them is invokable')
    args = ap.parse_args()

    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)

    def load(blob, cmd_id=None):
        ap_, amc = alloc(0x80000)
        cp, cmc = alloc(0x8000)
        n = (len(blob) + 3) & ~3
        adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob))
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=cmd_id or am.GFX_CMD_ID_LOAD_TA)
        t = c.cmd.cmd_load_ta
        t.app_phy_addr_lo, t.app_phy_addr_hi = amc & 0xFFFFFFFF, amc >> 32
        t.app_len = len(blob)
        t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
        t.cmd_buf_len = 0x8000
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id, cp
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None, cp
            except Exception:
                return -1, None, cp

    def invoke(sid, shm, cmd_id):
        buf = bytearray(0x2000)
        struct.pack_into('<I', buf, 0, cmd_id)          # many TAs read msg[0] too
        adev.vram.view(shm, len(buf))[:] = 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, cmd_id
        try:
            r = psp._ring_submit(c)
            st = r.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, 16)[:])
        return st, out

    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

    ids = [int(x, 0) for x in args.cmd_ids.split(',')]

    def run(tag, blob, load_cmd=None):
        """Calibrated on this machine's own bundle TAs:
             0x4        -> the PSP will not dispatch to this TA at all (ASD/driver class)
             0xffff0000 -> dispatched; the TA rejected the command id
             0x0        -> dispatched and accepted
           So "invokable" means: at least one id returns something other than 0x4."""
        st, sid, shm = load(blob, load_cmd)
        if st != 0:
            print('  %-12s LOAD failed %#x' % (tag, st), flush=True)
            return None
        res = []
        for cid in ids:
            ist, out = invoke(sid, shm, cid)
            res.append((cid, ist, out[:8].hex()))
        unload(sid)
        invokable = any(r[1] != 4 for r in res)
        print('  %-12s %-26s %s'
              % (tag, 'INVOKABLE' if invokable else 'refused (0x4 on every id)',
                 ' '.join('%d:%#x' % (c, s2) for c, s2, _ in res)), flush=True)
        return 0 if invokable else 4

    # ---- controls -------------------------------------------------------
    print('controls:')
    import subprocess
    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'
    pos = None
    data = open(tab, 'rb').read()
    p = 0
    while True:
        m = data.find(b'$PS1', p)
        if m < 0:
            break
        p = m + 4
        cont = m - 0x10
        if cont < 0:
            continue
        n = struct.unpack_from('<I', data, cont + 0x14)[0]
        if not (0 < n < 0x400000) or cont + 0x200 + n > len(data):
            continue
        if data[cont + 0x100:cont + 0x104] != b'\x00\x02\x00\x00':
            continue
        r = run('POS bundle#%d' % (p // max(n, 1) + 1), data[cont:cont + 0x200 + n])
        # bundle TA #1 is ASD (driver class) and refuses everything -- keep going
        # until one is invokable, which is the control we actually need.
        if r is not None and r != 4:
            pos = r
            break
    if pos is None or pos == 4:
        print('  POSITIVE CONTROL FAILED (bundle TA not invokable) -- stopping.')
        return 1
    if os.path.exists(args.asd):
        neg = run('NEG asd(DR)', open(args.asd, 'rb').read(), am.GFX_CMD_ID_LOAD_ASD)
        if neg != 4:
            print('  note: ASD did not return 0x4; the 0x4 signature may not be reliable')

    print('\nnewly discovered TAs:')
    for stem in NEW:
        hits = glob.glob(os.path.join(args.dir, stem + '*.bin'))
        if not hits:
            print('  %-12s (container not found)' % NAMES[stem])
            continue
        run(NAMES[stem], open(hits[0], 'rb').read())
    return 0


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