#!/usr/bin/env python3
"""Диагностическая сборка: зеркалит все кадры от DMR-модуля в USB-порт.

Перехватывается единственный вызов `fm100b_on_frame` (0x08006348) из разборщика
кадров. Наша функция выдаёт в компьютерный порт маркер F1 F2, длину и сам кадр,
после чего передаёт управление штатному обработчику — логика радио не меняется.

    python3 build_sniffer.py            # -> stock-fw/rt4d_sniffer.bin
    python3 tools/flash_ps.py stock-fw/rt4d_sniffer.bin
    python3 tools/sniff.py 60           # слушать 60 секунд

Сборка временная: для обычной работы вернуть stock-fw/rt4d_ru_batt.bin.
"""
import struct
from keystone import Ks, KS_ARCH_ARM, KS_MODE_THUMB

BASE = 0x08002800
SRC  = "stock-fw/rt4d_ru_batt.bin"
OUT  = "stock-fw/rt4d_sniffer.bin"

ON_FRAME  = 0x08006348      # void fm100b_on_frame(u8 *frame)
CALL_SITE = 0x08018DD8      # единственный bl на него
TX_BYTE   = 0x08018BFC      # void usart6_tx_byte(u8) — порт к компьютеру


def thumb_bl(src, dst):
    off = (dst - (src + 4)) >> 1
    s = (off >> 23) & 1; i1 = (off >> 22) & 1; i2 = (off >> 21) & 1
    imm10 = (off >> 11) & 0x3FF; imm11 = off & 0x7FF
    j1 = (~i1 ^ s) & 1; j2 = (~i2 ^ s) & 1
    return struct.pack("<HH", 0xF000 | (s << 10) | imm10,
                              0xD000 | (j1 << 13) | (j2 << 11) | imm11)


def main():
    img = bytearray(open(SRC, "rb").read())
    end = BASE + len(img)
    code_at = (end + 3) & ~3
    img += b"\xff" * (code_at - end)
    print(f"основа: {SRC}, {len(img)} Б; свой код с 0x{code_at:08X}")

    o = CALL_SITE - BASE
    assert bytes(img[o:o + 4]) == thumb_bl(CALL_SITE, ON_FRAME), \
        f"по 0x{CALL_SITE:08X} не вызов обработчика кадров"
    print("точка перехвата подтверждена")

    ks = Ks(KS_ARCH_ARM, KS_MODE_THUMB)

    # --- пролог: длина кадра = 9 + BE16(len), маркер F1 F2 + длина ---
    partA = bytes(ks.asm(f"""
        push  {{r4, r5, r6, lr}}
        mov   r4, r0
        ldrb  r5, [r4, #6]
        ldrb  r6, [r4, #7]
        orr.w r5, r6, r5, lsl #8
        adds  r5, #9
        uxtb  r5, r5
        movs  r0, #0xF1
        bl    #{TX_BYTE}
        movs  r0, #0xF2
        bl    #{TX_BYTE}
        mov   r0, r5
        bl    #{TX_BYTE}
        movs  r6, #0
    """, code_at)[0])

    # --- тело цикла (без обратного перехода) ---
    loop_at = code_at + len(partA)
    body = bytes(ks.asm(f"""
        ldrb  r0, [r4, r6]
        bl    #{TX_BYTE}
        adds  r6, #1
        cmp   r6, r5
    """, loop_at)[0])

    # --- blt назад на начало цикла (keystone не умеет метки) ---
    br_at = loop_at + len(body)
    off = (loop_at - (br_at + 4)) >> 1
    blt = struct.pack("<H", 0xDB00 | (off & 0xFF))

    # --- эпилог: передать кадр штатному обработчику ---
    tail_at = br_at + 2
    tail = bytes(ks.asm(f"""
        mov   r0, r4
        bl    #{ON_FRAME}
        pop   {{r4, r5, r6, pc}}
    """, tail_at)[0])

    code = partA + body + blt + tail
    img += code
    img[o:o + 4] = thumb_bl(CALL_SITE, code_at)
    print(f"снифер: {len(code)} Б, цикл на 0x{loop_at:08X}, blt смещение {off}")

    open(OUT, "wb").write(bytes(img))
    print(f"готово: {OUT}, {len(img)} Б")


if __name__ == "__main__":
    main()