"""STEP 1: detect how many times the PSP reads the image buffer (read multiplicity).

Why statistical rather than "load a corrupt image and watch it run": corrupting a
real IP firmware (MEC/SMU/...) to observe the loaded result risks hanging or
mis-programming the GPU. Verification *failure*, by contrast, is inert -- the PSP
just returns an error and loads nothing. So we measure read multiplicity purely
from verify pass/fail statistics, which is safe and still answers the security
question directly.

Method
------
The bridge flipper alternates one signature-covered byte between A (genuine) and
B (corrupt) with a controllable duty cycle d = fraction of time the byte is A.
Issue LOAD_IP_FW many times and measure the pass rate p.

  * PSP hashes the image ONCE   -> p ~= d
  * PSP reads it TWICE (e.g. hash pass then load/copy pass, both covered by the
    check) -> both reads must see A -> p ~= d^2
  * N reads -> p ~= d^N

Sweeping d and fitting log(p)/log(d) gives N. N > 1 is the double-fetch that makes
the bridge-SRAM TOCTOU exploitable; N == 1 means verify and load are a single read
and the attack does not apply to this path.

Control arms (must hold or the measurement is meaningless):
  * byte pinned to A -> p must be ~1.0
  * byte pinned to B -> p must be ~0.0   (proves the byte is actually hashed)

SAFETY: only ever submits an image whose bytes are either the genuine signed ones
or a one-byte corruption. A corrupted image FAILS verification and is not loaded.
Run with a fw_type that is already loaded, so a success merely reloads identical
content.
"""
import time, math, ctypes
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from psp_out.toctou_bridge import Flipper, M_ARM, M_DISARM


class Detector:
  def __init__(self, flip_off=0x800):
    self.f = Flipper()
    self.adev = self.f.adev
    self.psp = self.f.psp
    self.flip_off = flip_off
    self.f.select_method(0)          # XDATA window: does not disturb the 0x7000 staging path

  def pick_image(self):
    """Choose an already-loaded IP firmware: reloading identical content is inert."""
    for types, blob in self.adev.fw.descs:
      if len(blob) > self.flip_off + 16:
        return types[0], blob
    raise RuntimeError("no suitable firmware descriptor found")

  def submit(self, fw_type, blob) -> int:
    """One LOAD_IP_FW. Returns the PSP status word (0 == verified+loaded)."""
    self.psp._prep_msg1(blob)
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
    from tinygrad.helpers import data64
    cmd.cmd.cmd_load_ip_fw.fw_phy_addr_hi, cmd.cmd.cmd_load_ip_fw.fw_phy_addr_lo = data64(self.psp.msg1_addr)
    cmd.cmd.cmd_load_ip_fw.fw_size = len(blob)
    cmd.cmd.cmd_load_ip_fw.fw_type = fw_type
    resp = self.psp._ring_submit(cmd)
    return resp.resp.status

  def trial(self, fw_type, blob, mode, trials=12):
    """mode: 'A' pin genuine, 'B' pin corrupt, 'flip' alternate."""
    genuine = blob[self.flip_off]
    corrupt = genuine ^ 0xFF
    ok = 0
    for _ in range(trials):
      self.psp._prep_msg1(blob)                       # restage pristine image
      if mode == "A":    self.f.poke(self.flip_off, genuine)
      elif mode == "B":  self.f.poke(self.flip_off, corrupt)
      else:              self.f.arm(self.flip_off, genuine, corrupt)
      try:
        st = self.submit(fw_type, blob)
      except Exception:
        st = -1
      finally:
        if mode == "flip": self.f.disarm()
      ok += (st == 0)
    return ok / trials


if __name__ == "__main__":
  d = Detector()
  fw_type, blob = d.pick_image()
  print("target fw_type=%d size=%#x  flip_off=%#x" % (fw_type, len(blob), d.flip_off))
  print("\n[controls] these must hold or the measurement is invalid:")
  pa = d.trial(fw_type, blob, "A")
  print("  byte pinned GENUINE : pass rate %.2f  (expect ~1.00)" % pa)
  pb = d.trial(fw_type, blob, "B")
  print("  byte pinned CORRUPT : pass rate %.2f  (expect ~0.00 => byte IS hashed)" % pb)
  if not (pa > 0.8 and pb < 0.2):
    print("\nCONTROLS FAILED -> offset %#x is probably not covered by the signature check;"
          " pick a different flip_off before trusting any flip result." % d.flip_off)
  else:
    pf = d.trial(fw_type, blob, "flip", trials=24)
    print("\n[flip] free-running A/B : pass rate %.2f" % pf)
    print("  single-read model predicts ~0.50; two reads predict ~0.25")
    if pf > 0.05:
      print("  -> verification passed at least once while the byte was flipping.")
      print("     Next: prove verify-saw-A while load-saw-B (the actual TOCTOU).")
    else:
      print("  -> never passed; window may be too tight, or reads are single+atomic.")