#!/usr/bin/env python3
"""Scan firmware for command-dispatch CMP-immediate and branch-switch patterns."""
from capstone import Cs, CS_ARCH_ARM, CS_MODE_THUMB, CS_MODE_LITTLE_ENDIAN
import struct, collections

IMG = "app/firmware-PP_da14585_4.2_CH.img"
BASE = 0x07FC0000
HDR = 0x40
CODE_END = 0x119e4  # code region size

CMDS = {0xdd:"time", 0xe2:"commit", 0x03:"img_blk_full", 0x04:"img_blk_small",
        0x01:"refresh", 0xAA:"refresh_small", 0xEF:"activate", 0xE6:"lut_cal",
        0xFB:"sleep_time", 0xEA:"slot", 0xFA:"countdown_off", 0xE1:"mode",
        0xE3:"invert", 0xE4:"font", 0xE5:"rotate", 0x23:"bw_3color", 0x14:"12_24h",
        0x15:"clk_color", 0x99:"cal_text", 0x00:"refresh_step1", 0xE9:"carousel"}

def main():
    data = open(IMG, "rb").read()[HDR:HDR+CODE_END]
    md = Cs(CS_ARCH_ARM, CS_MODE_THUMB | CS_MODE_LITTLE_ENDIAN)
    # collect all instructions via same sweep
    insns = []
    off = 0
    while off < len(data):
        chunk = data[off:off+0x2000]
        addr = BASE + off
        consumed = 0
        for ins in md.disasm(chunk, addr):
            if ins.address < addr + consumed:
                continue
            gap = ins.address - (addr + consumed)
            if gap:
                consumed += (gap // 2) * 2
            insns.append(ins)
            consumed = ins.address + ins.size - addr
        off += consumed if consumed else 2

    # CMP Rn, #imm with imm in CMDS
    hits = collections.defaultdict(list)
    for ins in insns:
        if ins.mnemonic == "cmp":
            parts = ins.op_str.split(",")
            if len(parts) == 2 and parts[1].strip().startswith("#"):
                try:
                    v = int(parts[1].strip()[1:], 0)
                except ValueError:
                    continue
                if v in CMDS:
                    hits[v].append((ins.address, ins.op_str))
    for v in sorted(hits):
        print(f"--- CMP #0x{v:02x} ({CMDS[v]}): {len(hits[v])} hits ---")
        for a, ops in hits[v]:
            print(f"  0x{a:08x}: cmp {ops}")

if __name__ == "__main__":
    main()