#!/usr/bin/env python3
"""Build the Smatfarm GD32F130 UART-control patch (v2) for YOUR firmware dump.

    python3 apply_patch_v2.py my_dump.bin tramp.bin impl.bin

v2 differs from v1 in two important ways:

 1. v1 overwrote the function at 0x08000678 believing it was dead touch-handling
    code. It is NOT -- it is the FAN CONTROLLER (fan on if any channel setpoint
    >= 20, driving PB9). v1 therefore leaves the fan permanently off, because
    boot init drives PB9 low and nothing else ever raises it. v2 reimplements
    the fan logic alongside the UART parser.

 2. Fan + UART logic does not fit the 124-byte hole, so v2 places a 4-byte
    tail-call trampoline at 0x08000678 and the real code at 0x0800a000, in the
    unused flash region.

This writes TWO SMALL PAGE FILES rather than a full 64KB image, on purpose:
reflashing all 64KB also rewrites the upper 32KB, and a dump of that region
read through a CH347 adapter is not reliable (see the Verify your dump
section of the writeup). Flash only what actually changed.

Outputs:
    page_08000400.bin   1KB page containing the trampoline -> flash at 0x08000400
    page_0800a000.bin   the implementation                 -> flash at 0x0800a000
"""
import sys

TRAMP_ADDR = 0x08000678
TRAMP_END  = 0x080006F4     # end of the original function; pad to here
IMPL_ADDR  = 0x0800A000
PAGE_SIZE  = 0x400
FLASH_BASE = 0x08000000
NOP        = b"\x00\xbf"

def main():
    if len(sys.argv) != 4:
        print(__doc__); sys.exit(1)
    dump_p, tramp_p, impl_p = sys.argv[1:4]
    fw    = bytearray(open(dump_p, "rb").read())
    tramp = open(tramp_p, "rb").read()
    impl  = open(impl_p, "rb").read()

    if len(fw) != 0x10000:
        print(f"error: expected a 65536-byte dump, got {len(fw)}"); sys.exit(1)

    off = TRAMP_ADDR - FLASH_BASE
    if fw[off:off+2] != b"\x10\xb5":
        print("error: bytes at 0x08000678 are not the expected push{r4,lr} (10 b5).\n"
              "       Either this dump is a different firmware version, or it is\n"
              "       ALREADY PATCHED. Do not apply this on top of v1 -- reflash\n"
              "       your original backup first.")
        sys.exit(1)

    hole = TRAMP_END - TRAMP_ADDR
    if len(tramp) > hole:
        print(f"error: trampoline is {len(tramp)} bytes, only {hole} available"); sys.exit(1)

    # splice trampoline + NOP padding across the original function
    fw[off:off+len(tramp)] = tramp
    pad = TRAMP_END - (off + len(tramp)) - (TRAMP_ADDR - FLASH_BASE - off)
    pad_start, pad_end = off + len(tramp), TRAMP_END - FLASH_BASE
    fw[pad_start:pad_end] = (NOP * ((pad_end - pad_start) // 2 + 1))[:pad_end - pad_start]

    # page containing the trampoline
    page_start = (off // PAGE_SIZE) * PAGE_SIZE
    page = fw[page_start:page_start + PAGE_SIZE]
    p1 = f"page_{FLASH_BASE + page_start:08x}.bin"
    open(p1, "wb").write(page)

    # implementation target must be unused flash
    io = IMPL_ADDR - FLASH_BASE
    region = fw[io:io + len(impl)]
    if any(b not in (0x00, 0xFF) for b in region):
        print(f"warning: 0x{IMPL_ADDR:08x} is not blank in your dump -- it contains\n"
              f"         real data. Pick a different free address before flashing.")
    p2 = f"page_{IMPL_ADDR:08x}.bin"
    open(p2, "wb").write(impl)

    print(f"wrote {p1}  ({len(page)} bytes)  -> flash at 0x{FLASH_BASE + page_start:08x}")
    print(f"wrote {p2}  ({len(impl)} bytes)  -> flash at 0x{IMPL_ADDR:08x}")
    print(f"\ntrampoline {len(tramp)}B + {pad_end - pad_start}B NOP padding at 0x{TRAMP_ADDR:08x}")
    print("\nFlash with (one page at a time, verify each):")
    print(f'  flash erase_address 0x{FLASH_BASE + page_start:08x} 0x{PAGE_SIZE:x}')
    print(f'  flash write_image {p1} 0x{FLASH_BASE + page_start:08x}')
    print(f'  flash erase_address 0x{IMPL_ADDR:08x} 0x{PAGE_SIZE:x}')
    print(f'  flash write_image {p2} 0x{IMPL_ADDR:08x}')

if __name__ == "__main__":
    main()