#!/usr/bin/env python3
"""
Patch fw_tinygrad.bin to add the GCP TOCTOU hammer behind vendor opcode 0xEB.

Mirrors app/patch_e7_fastread.py but with the OFFSETS FOR fw_tinygrad.bin (the
stock E7 patcher's 0x1411c is wrong for this image; the dispatch is at 0x140E0).
It repurposes the unused E2 vendor command (tinygrad uses E1/E3/E4/E5/8A):

  1. 0x140E1: cjne a,#0xE2  ->  cjne a,#0xEB            (select our opcode)
  2. 0x140E4: lcall 0xD945  ->  lcall 0x7D24            (-> handler entry)
  3. 0x7D00 : inject the linked handler blob (237 B, free/zero region)

Handler entry is 0x7D24 (tlp_write_fixed sits first at 0x7D00; see build/h.map).

Every patch verifies the expected original bytes first and REFUSES on mismatch,
so running it against a different image is inert rather than corrupting.
"""
import sys, hashlib
from pathlib import Path

ROOT = Path(__file__).parent
SRC  = ROOT / "fw_tinygrad.bin"
OUT  = ROOT / "fw_tinygrad_hammer.bin"
BLOB = ROOT / "build" / "handler_bytes.bin"

# expected sha256 prefix of the pristine fw_tinygrad.bin this patcher was built against
EXPECT_SRC_SHA16 = "e496176650b83985"

HANDLER_BASE   = 0x7D00
HANDLER_ENTRY  = 0x7D24

# (offset, expected_original, replacement)
PATCH_OPCODE = (0x140E1, b"\xE2", b"\xEB")
PATCH_LCALL  = (0x140E4, b"\xD9\x45", b"\x7D\x24")   # lcall 0xD945 -> lcall 0x7D24

def main():
    data = bytearray(SRC.read_bytes())
    got = hashlib.sha256(data).hexdigest()[:16]
    if got != EXPECT_SRC_SHA16:
        print(f"REFUSE: source sha256 {got} != expected {EXPECT_SRC_SHA16}.")
        print("  This patcher's offsets were derived from a specific fw_tinygrad.bin.")
        print("  Re-derive the dispatch offset (cjne a,#0xe2) before patching another image.")
        sys.exit(1)

    blob = BLOB.read_bytes()
    assert HANDLER_BASE + len(blob) <= 0x7E00, "handler overflows free slot 0x7D00..0x7E00"

    ok = True
    for off, old, new in (PATCH_OPCODE, PATCH_LCALL):
        if data[off:off+len(old)] != old:
            print(f"REFUSE: at {off:#x} expected {old.hex()} found {data[off:off+len(old)].hex()}")
            ok = False
    # handler slot must be empty
    if any(data[HANDLER_BASE:HANDLER_BASE+len(blob)]):
        print(f"REFUSE: handler slot {HANDLER_BASE:#x} is not zero-filled")
        ok = False
    if not ok:
        sys.exit(1)

    for off, old, new in (PATCH_OPCODE, PATCH_LCALL):
        data[off:off+len(new)] = new
        print(f"patched {off:#x}: {old.hex()} -> {new.hex()}")
    data[HANDLER_BASE:HANDLER_BASE+len(blob)] = blob
    print(f"injected {len(blob)}-byte handler at {HANDLER_BASE:#x} (entry {HANDLER_ENTRY:#x})")

    OUT.write_bytes(bytes(data))
    print(f"wrote {OUT}  sha256={hashlib.sha256(bytes(data)).hexdigest()[:16]}  ({len(data)} bytes)")
    # sanity: re-decode the hooked dispatch
    d = bytes(data)
    assert d[0x140E0:0x140E6] == b"\xB4\xEB\x0E\x12\x7D\x24", d[0x140E0:0x140E6].hex()
    print("verify: dispatch now 'cjne a,#0xEB,+14 ; lcall 0x7D24'  OK")

if __name__ == "__main__":
    main()