"""READ-ONLY: can the HOST read the 230 registers on the PSP's own SLWR allowlist?

The SLWR table in the SOS bundle (`drivers/rl.bin`) is the firmware's own list of
SMN registers that a TA may touch through `ta_reg_access`. Those 230 addresses
are therefore a *firmware-sanctioned* probe set: interesting registers, known to
decode on this part, with no guessing and no sweeping. Reading them from the
host answers a question the static analysis cannot: does host SMN access
(security level 7) reach the same registers the PSP (level 0) is allowed to?

WHY THIS IS NOT A BLIND SWEEP
    Every address comes from a table shipped in the firmware we run. We are not
    walking an address space hoping to find something.

WHY VALUE ALONE IS NOT AN ANSWER (STALEUS / CVE-2025-54509)
    An SMN access refused for insufficient privilege returns 0x0 or 0xFFFFFFFF,
    which is exactly what an absent register returns. So this script classifies
    with regguard's denial-aware verdicts and never reports "not present" on the
    basis of a value alone.

CONTROLS RUN BEFORE ANY DATA IS BELIEVED (all must pass):
    1. distinct addresses must not all return one value  (catches ALIAS)
    2. repeated reads must agree                         (catches UNSTABLE)
    3. DECISIVE cross-check: read PSP mailbox C2PMSG_81 by NAME through the BAR
       and at its absolute SMN address through the index/data window; the two
       must agree, proving the window actually addresses what we think.

Strictly read-only: writes only the PCIE_INDEX2 selector, which is what a read
through this window requires.
"""
import sys, os, time, collections

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from slwr_parse import parse as slwr_parse

from tinygrad import Device

IDX, DAT = "regBIF_BX0_PCIE_INDEX2", "regBIF_BX0_PCIE_DATA2"
DENIAL = (0x00000000, 0xFFFFFFFF)

d = Device["AMD"]
adev = d.iface.dev_impl


def smn_read(addr):
    adev.reg(IDX).write(addr & 0xFFFFFFFF)
    return adev.reg(DAT).read()


def controls():
    """Return True only if the read window is proven to address real registers."""
    ok = True

    # (3) decisive: mailbox by name vs by absolute SMN address
    # C2PMSG_N (N>=64) = MP0 block 0x03800000 + 0x10A00 + 4*(N-64)
    smn_mbox = 0x03800000 + 0x10A00 + 4 * (81 - 64)
    try:
        by_name = adev.reg("regMP0_SMN_C2PMSG_81").read()
    except Exception as e:
        print("  control 3: cannot read C2PMSG_81 by name (%s)" % e)
        return False
    by_smn = smn_read(smn_mbox)
    agree = by_name == by_smn
    print("  control 3 (window addresses what we think):")
    print("      C2PMSG_81 by name %#010x   via SMN %#010x  -> %s"
          % (by_name, by_smn, "AGREE" if agree else "DISAGREE"))
    ok &= agree

    # (1) several distinct addresses must not all read the same
    probe = [smn_mbox, smn_mbox + 4, smn_mbox + 8, 0x03800000 + 0x10A00]
    vals = [smn_read(a) for a in probe]
    distinct = len(set(vals)) > 1
    print("  control 1 (not all-alias): %s -> %s"
          % ([hex(v) for v in vals], "OK" if distinct else "ALIAS SUSPECTED"))

    # (2) stability
    a = smn_mbox
    reps = [smn_read(a) for _ in range(5)]
    stable = len(set(reps)) == 1
    print("  control 2 (stable repeats): %s -> %s"
          % ([hex(v) for v in reps], "OK" if stable else "UNSTABLE (may be a counter)"))
    return ok


def main():
    _d, _base, _n, entries = slwr_parse()
    addrs = [a for a, _attr in entries]
    print("SLWR allowlist: %d entries, %#x .. %#x\n" % (len(addrs), min(addrs), max(addrs)))

    print("== controls ==")
    if not controls():
        print("\nCONTROLS FAILED -- refusing to report register data.")
        return 2
    print("  -> window validated\n")

    print("== reading the %d allowlisted registers (read-only) ==" % len(addrs))
    results = {}
    t0 = time.time()
    for a in addrs:
        v1 = smn_read(a)
        v2 = smn_read(a)
        results[a] = (v1, v2)
    print("  %d reads in %.2fs\n" % (2 * len(addrs), time.time() - t0))

    # classify
    cls = collections.Counter()
    nonzero = []
    for a, (v1, v2) in results.items():
        if v1 != v2:
            k = "UNSTABLE"
        elif v1 == 0x00000000:
            k = "DENIED?-or-ZERO"
        elif v1 == 0xFFFFFFFF:
            k = "DENIED?-or-ABSENT"
        elif v1 == int.from_bytes(a.to_bytes(4, "little"), "big"):
            k = "MIRROR"
        else:
            k = "REAL"
            nonzero.append((a, v1))
        cls[k] += 1

    print("== classification ==")
    for k, n in cls.most_common():
        print("  %-20s %d" % (k, n))

    allvals = [v for v, _ in results.values()]
    if len(set(allvals)) == 1:
        print("\n  !! every address returned %#010x -- ALIAS/DEAD window, not data"
              % allvals[0])

    if nonzero:
        print("\n== registers returning real-looking data (%d) ==" % len(nonzero))
        for a, v in sorted(nonzero)[:64]:
            print("  %#010x = %#010x   (IP %#05x off %#07x)"
                  % (a, v, a >> 20, a & 0xFFFFF))
        if len(nonzero) > 64:
            print("  ... %d more" % (len(nonzero) - 64))
    else:
        print("\n  No allowlisted register returned a non-denial value from the host.")

    # per-IP summary
    print("\n== per-IP summary ==")
    byip = collections.defaultdict(list)
    for a, (v1, _) in results.items():
        byip[a >> 20].append(v1)
    for ip in sorted(byip):
        vs = byip[ip]
        real = sum(1 for v in vs if v not in DENIAL)
        print("  IP %#05x : %3d regs, %3d non-denial, distinct values %d"
              % (ip, len(vs), real, len(set(vs))))
    return 0


if __name__ == "__main__":
    sys.exit(main())