#!/usr/bin/env python3
"""Phase-0 UI hook PoC, built ON TOP of the russified image.

Injects a custom draw-router at 0x08029000 (free flash, after the russification
cave@0x08028860 + font@0x08028900). Redirects ONLY the draw dispatch bl at
0x080207E6 to the router, which:
  - tail-calls the stock draw dispatcher 0x0801E1BC (screen renders normally),
  - then, if g_screen==0 (home/standby), paints a custom "МОД" marker via the
    stock draw_string SDK (also exercising our CP1251 Cyrillic cave).
Key dispatch is left 100% stock. One 4-byte code patch; fully recoverable.
"""
import struct
from keystone import Ks, KS_ARCH_ARM, KS_MODE_THUMB

BASE      = 0x08002800
ROUTER    = 0x08029000
MARKER    = 0x08029040
DRAW_BL   = 0x080207E6      # bl 0x0801E1BC in ui_tick_normal
STOCK_DRAW= 0x0801E1BC
DRAW_STR  = 0x08008A50
G_SCREEN  = 0x200008B3
ROOT      = "C:/Users/vikto/Documents/Claude/rt-4d"

def main():
    img = bytearray(open(f"{ROOT}/stock-fw/rt4d_ru_full.bin", "rb").read())
    if len(img) < (MARKER+8-BASE):
        img += b"\xff" * ((MARKER+8-BASE) - len(img))
    def w(va, b): img[va-BASE:va-BASE+len(b)] = b
    ks = Ks(KS_ARCH_ARM, KS_MODE_THUMB)

    router = f"""
        push {{r4, lr}}
        movw r4, #{STOCK_DRAW+1 & 0xFFFF}
        movt r4, #{(STOCK_DRAW+1) >> 16}
        blx  r4                      // stock draw dispatch (renders current screen)
        movw r4, #{G_SCREEN & 0xFFFF}
        movt r4, #{G_SCREEN >> 16}
        ldrb r4, [r4]
        cmp  r4, #0
        bne  done                    // only decorate the home screen
        sub  sp, #8
        movs r0, #0
        str  r0, [sp]                // mode = 0 (arg5 -> [sp+0x28] inside draw_string)
        movs r0, #0                  // page (top)
        movs r1, #104                // x
        movw r2, #{MARKER & 0xFFFF}
        movt r2, #{MARKER >> 16}
        movs r3, #3                  // len
        movw r4, #{DRAW_STR+1 & 0xFFFF}
        movt r4, #{(DRAW_STR+1) >> 16}
        blx  r4                      // draw_string(0,104,marker,3,0)
        add  sp, #8
      done:
        pop  {{r4, pc}}
    """
    code = bytes(ks.asm(router, ROUTER)[0])
    assert len(code) <= (MARKER-ROUTER), f"router too big: {len(code)}"
    w(ROUTER, code)
    w(MARKER, "МОД".encode("cp1251") + b"\x00")

    # redirect the draw dispatch bl -> router
    old = bytes(img[DRAW_BL-BASE:DRAW_BL-BASE+4])
    new = bytes(ks.asm(f"bl #{ROUTER}", DRAW_BL)[0]); assert len(new) == 4
    w(DRAW_BL, new)

    out = f"{ROOT}/stock-fw/rt4d_ui_poc.bin"
    open(out, "wb").write(img)
    print(f"router {len(code)}B @0x{ROUTER:08X} | marker МОД @0x{MARKER:08X}")
    print(f"draw bl @0x{DRAW_BL:08X}: {old.hex()} -> {new.hex()}")
    print(f"image {len(img)}B -> {out}")

if __name__ == "__main__":
    main()