#!/usr/bin/env python3
"""Detect the unsigned VCN write by diffing the VCN MMIO window.

No register names needed: snapshot the window, perform a type-49 write carrying
a distinctive value, snapshot again, diff. A register holding our magic proves
(a) the writes really happen and (b) the index -> register mapping.

Controls:
  * a no-write baseline diff (snapshot twice with no command in between), to
    measure how much the window moves on its own;
  * a rejected command (missing index 195) between snapshots, which must not
    produce our magic.

Progress to /tmp/vcndet.txt.
"""
import struct, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

OUT = "/tmp/vcndet.txt"
MAGIC = 0xA5A50000


def say(s):
    with open(OUT, "a") as f:
        f.write(s + "\n")
    print(s, flush=True)


adev = Device["AMD"].iface.dev_impl
psp = adev.psp
hb = adev.reg("regMP0_SMN_C2PMSG_81")


def alive():
    a = hb.read(); time.sleep(0.25); return hb.read() != a


# VCN MMIO base from the card's own IP discovery
BASE, END = 0x7800, 0x7F00


def snap():
    out = {}
    for r in range(BASE, END):
        try:
            out[r] = adev.rreg(r)
        except Exception:
            pass
    return out


def diff(a, b):
    return {k: (a.get(k), b.get(k)) for k in b if a.get(k) != b.get(k)}


def submit(payload, fw_type=49):
    psp._prep_msg1(memoryview(bytes(payload)))
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
    cmd.cmd.cmd_load_ip_fw.fw_phy_addr_lo = psp.msg1_addr & 0xFFFFFFFF
    cmd.cmd.cmd_load_ip_fw.fw_phy_addr_hi = psp.msg1_addr >> 32
    cmd.cmd.cmd_load_ip_fw.fw_size = len(payload)
    cmd.cmd.cmd_load_ip_fw.fw_type = fw_type
    try:
        return psp._ring_submit(cmd).resp.status & 0xFFFFFFFF
    except RuntimeError as e:
        return int(str(e).rsplit(" ", 1)[-1]) & 0xFFFFFFFF


def ent(i, v):
    return struct.pack("<II", i, v)


say("\n######## VCN write detection ########")
say("window: MMIO dword %#x..%#x (%d regs)" % (BASE, END, END - BASE))
say("gate: %#010x (0xffff300f = open)" % submit(b"\x00" * 0x208))

s0 = snap()
nz = sum(1 for v in s0.values() if v not in (0, 0xFFFFFFFF))
say("snapshot0: %d regs read, %d non-trivial" % (len(s0), nz))

s1 = snap()
base_noise = diff(s0, s1)
say("CONTROL A (no command between snapshots): %d regs differ %s"
    % (len(base_noise), list(base_noise)[:5]))

say("\nCONTROL B: rejected command (no index 195) between snapshots")
st = submit(b"".join(ent(0x100 + i, MAGIC | i) for i in range(8)))
s2 = snap()
d2 = diff(s1, s2)
say("   status %#010x ; %d regs differ ; any holding magic: %s"
    % (st, len(d2), [hex(k) for k, (a, b) in d2.items() if b is not None and (b & 0xFFFF0000) == MAGIC] or "none"))

say("\nTHE WRITE: index 195 + 63 host entries carrying magic")
tbl = ent(195, 0) + b"".join(ent(0x100 + i, MAGIC | i) for i in range(63))
st = submit(tbl)
say("   status %#010x" % st)
s3 = snap()
d3 = diff(s2, s3)
hits = {k: b for k, (a, b) in d3.items() if b is not None and (b & 0xFFFF0000) == MAGIC}
say("   %d regs changed ; %d hold our magic" % (len(d3), len(hits)))
for k, v in list(hits.items())[:12]:
    say("      MMIO dword %#06x = %#010x   (entry index %d)" % (k, v, v & 0xFFFF))
if not hits and d3:
    say("   changed regs (first 8): %s" % [(hex(k), hex(a or 0), hex(b or 0)) for k, (a, b) in list(d3.items())[:8]])

say("\npsp alive %s / smu alive %s" % (alive(), adev.smu.is_smu_alive()))