#!/usr/bin/env python3
"""Sweep VRAM for ranges the PSP refuses to read from.

Uses the validated LOAD_TA classification oracle
(psp_out/FINDING-LOADTA-ADDRESS-CLASSIFICATION-ORACLE.md):

    0x34  readable, contents invalid   -> ordinary VRAM
    0xa   refused                      -> protected / reserved
    0x0   a real TA image              -> only at our control buffer

app_len is kept SMALL (0x1000) so the classified extent is one page rather than
the whole blob -- the TMR sweep showed the check covers [start, start+len), so a
large length smears boundaries by the blob size.

Every probe is inside backed VRAM (CONSTRAINTS.md: never aim LOAD_TA outside
VRAM; it wedges the SMU).

Known ground truth used as a live control: the TMR at PA [0x300000,0x1600000)
must come back refused. If it does not, the oracle is not working on this boot
and the sweep's negatives are meaningless.
"""
import sys, time
sys.path.insert(0, "/root/psp/work")
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

LO   = int(sys.argv[1], 0) if len(sys.argv) > 1 else 0x0
HI   = int(sys.argv[2], 0) if len(sys.argv) > 2 else 0x10000000
STEP = int(sys.argv[3], 0) if len(sys.argv) > 3 else 0x100000
APP_LEN = 0x1000

def main():
    adev = Device["AMD"].iface.dev_impl
    hb = adev.reg("regMP0_SMN_C2PMSG_81")
    def alive():
        a = hb.read(); time.sleep(0.2); return hb.read() != a
    if not alive(): raise SystemExit("PSP not healthy before start")
    cmd_pa = adev.mm.palloc(0x10000, align=0x1000, zero=True, boot=False)
    mcc = adev.paddr2mc(cmd_pa)

    def probe(pa):
        mc = adev.paddr2mc(pa)
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA)
        c.cmd.cmd_load_ta.app_phy_addr_lo = mc & 0xFFFFFFFF
        c.cmd.cmd_load_ta.app_phy_addr_hi = mc >> 32
        c.cmd.cmd_load_ta.app_len = APP_LEN
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_lo = mcc & 0xFFFFFFFF
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_hi = mcc >> 32
        c.cmd.cmd_load_ta.cmd_buf_len = 0x10000
        try:
            r = adev.psp._ring_submit(c); return r.resp.status
        except Exception as e:
            try: return int(str(e).strip().split()[-1])
            except Exception: return -1

    # live control: the TMR must classify as refused on THIS boot
    ctl = probe(0x900000)
    print("control: TMR PA 0x900000 -> %#x  %s" % (ctl, "REFUSED (oracle live)" if ctl == 0xA else "NOT REFUSED"))
    if ctl != 0xA:
        print("RESULT NOT REPORTABLE: oracle control failed on this boot"); return 1

    t0 = time.time()
    hits, n = [], 0
    pa = LO
    print("sweeping PA [%#x, %#x) step %#x, app_len %#x" % (LO, HI, STEP, APP_LEN))
    while pa < HI:
        st = probe(pa); n += 1
        if st == 0xA: hits.append(pa)
        elif st not in (0x34, 0x0):
            print("   %#-12x unexpected status %#x" % (pa, st))
        if n % 64 == 0:
            print("   ... %#x  (%d probed, %d refused, %.0fs)" % (pa, n, len(hits), time.time() - t0))
            if not alive():
                print("PSP STALLED at %#x -- stopping" % pa); break
        pa += STEP
    print()
    print("probed %d addresses in %.0fs; %d refused" % (n, time.time() - t0, len(hits)))
    runs = []
    for h in hits:
        if runs and h == runs[-1][1] + STEP: runs[-1][1] = h
        else: runs.append([h, h])
    print("refused ranges (step-granular):")
    for a, b in runs:
        known = " <- known TMR" if not (b < 0x300000 or a >= 0x1600000) else "  *** NOT the TMR ***"
        print("   %#010x .. %#010x  (%d pages of step)%s" % (a, b + STEP, (b - a) // STEP + 1, known))
    print("PSP alive after sweep: %s" % alive())
    return 0

sys.exit(main())