#!/usr/bin/env python3
"""Does the PSP's status distinguish UNREACHABLE memory from readable-but-invalid?

Step 2 of STRATEGY-MEASUREMENT-FIRST.md. The emulator can only separate
mapped/unmapped because that is all it models; the axis we need
(protected vs absent vs present) only exists on silicon.

LOAD_TA is the lever: the host chooses app_phy_addr, and the PSP reads and
verifies that buffer using the CCP. So the returned status is a function of
what the PSP found at an address WE choose.

Cheapest kill-test first: if an ABSENT address returns the same status as a
readable-but-invalid buffer, there is no access-vs-content separation at all and
the strategy dies here.

CONTROLS (per pspemu.oracle discipline):
  positive : a real TA image at a real address must give the known-good status,
             proving the path works and the status is meaningful
  negative : two DIFFERENT readable-but-invalid buffers must give the SAME
             status, proving the status is not just noise
"""
import sys, struct, traceback
sys.path.insert(0, "/root/psp/work")
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

def main():
    dev = Device["AMD"].iface.dev_impl
    adev = dev

    import time
    hb = adev.reg("regMP0_SMN_C2PMSG_81")
    def alive():
        a = hb.read(); time.sleep(0.3); return hb.read() != a
    if not alive():
        raise SystemExit("PSP heartbeat not advancing before we start -- aborting")
    print("PSP healthy before we start")

    def rnd(n, a=0x1000): return (n + a - 1) & ~(a - 1)

    import load_ta
    blob, _ver = load_ta.ta_blob("ras")
    good_pa = adev.mm.palloc(rnd(len(blob)), align=0x1000, zero=True, boot=False)
    cmd_pa  = adev.mm.palloc(0x10000, align=0x1000, zero=True, boot=False)
    junk_pa = adev.mm.palloc(rnd(len(blob)), align=0x1000, zero=True, boot=False)
    junk2_pa= adev.mm.palloc(rnd(len(blob)), align=0x1000, zero=True, boot=False)
    adev.vram.view(good_pa, len(blob), "B")[:] = blob
    adev.vram.view(junk_pa, 0x1000, "B")[:] = b"\xA5" * 0x1000
    adev.vram.view(junk2_pa, 0x1000, "B")[:] = b"\x5C" * 0x1000

    # NO absent-address arm. Aiming LOAD_TA outside VRAM produced a VMC VM-fault
    # storm that wedged the SMU and cost a power cycle (psp_out/CONSTRAINTS.md).
    # Every arm below is inside real, backed memory.
    #
    # TMR is PA [0x300000, 0x1600000) on this card (tmr_paddr=0x300000,
    # tmr_size=0x1300000). Backed VRAM, but protected -- which is the axis the
    # emulator cannot model and the whole reason for going to silicon.
    TMR_PA = 0x900000

    def load_at(mc_app, size, tag):
        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_app & 0xFFFFFFFF
        c.cmd.cmd_load_ta.app_phy_addr_hi = mc_app >> 32
        c.cmd.cmd_load_ta.app_len = size
        mc_cmd = adev.paddr2mc(cmd_pa)
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_lo = mc_cmd & 0xFFFFFFFF
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_hi = mc_cmd >> 32
        c.cmd.cmd_load_ta.cmd_buf_len = 0x10000
        try:
            r = adev.psp._ring_submit(c)
            st, sess = r.resp.status, r.resp.session_id
        except Exception as e:
            # tinygrad raises "PSP command failed <cmd> <status>"; the status is
            # the datum, so parse it rather than discarding it as an error.
            msg = str(e); st = None
            try: st = int(msg.strip().split()[-1])
            except Exception: pass
            print("   %-32s status=%-#12x (raised: %s)  psp=%s"
                  % (tag, st if st is not None else -1, msg[:34], "alive" if alive() else "DEAD"))
            return st
        hp = "alive" if alive() else "DEAD"
        print("   %-32s status=%-#12x session=%-#8x psp=%s" % (tag, st, sess or 0, hp))
        return st

    print("addresses: good=%#x junk=%#x junk2=%#x  TMR_MC=%#x"
          % (adev.paddr2mc(good_pa), adev.paddr2mc(junk_pa), adev.paddr2mc(junk2_pa), adev.paddr2mc(TMR_PA)))
    print()
    print("LOAD_TA with the image buffer aimed at different memory:")
    s_good  = load_at(adev.paddr2mc(good_pa),  len(blob), "real TA image (positive control)")
    s_junk  = load_at(adev.paddr2mc(junk_pa),  len(blob), "readable junk A")
    s_junk2 = load_at(adev.paddr2mc(junk2_pa), len(blob), "readable junk B")
    s_tmr   = load_at(adev.paddr2mc(TMR_PA),   len(blob), "inside TMR (protected, backed)")
    print()
    pos = (s_good is not None)
    neg = (s_junk is not None and s_junk2 is not None and s_junk != s_junk2)
    print("CONTROL positive (path works, real image gives a status): %s" % ("PASS" if pos else "FAIL"))
    print("CONTROL negative (two junk buffers agree, status not noise): %s"
          % ("PASS" if not neg else "FAIL"))
    if not pos or neg:
        print("RESULT NOT REPORTABLE"); return 1
    sep = (s_tmr != s_junk)
    print()
    print("readable-invalid = %s   TMR-protected = %s"
          % (hex(s_junk), hex(s_tmr) if s_tmr is not None else None))
    if sep:
        print("=> SEPARATE: a protected source is distinguishable from an unreadable-content one.")
        print("   The classification oracle is REAL on this axis.")
    else:
        print("=> IDENTICAL. Most likely reading: the PSP is PERMITTED to read its own")
        print("   TMR (it owns it), so protection never triggers on a PSP-initiated")
        print("   read. That is a useful fact, but it means LOAD_TA cannot classify")
        print("   protection, and this oracle does not exist via this path.")
    return 0

if __name__ == "__main__":
    try: sys.exit(main())
    except Exception:
        traceback.print_exc(); sys.exit(3)