#!/usr/bin/env python3
"""
JLink 现场调试 0x82 阶段 —— 验证 deepseek 的 gMcuUpgradeMsgQueue NULL 假设

断点设置：
- 0x1A3A2: 0x82 命令分发入口 (cmp r0, #130)
- 0x1A4A4: 0x82 处理入口
- 0x1A4F6: 0x82 状态检查 (ldrb r0, [r4, #0])
- 0x1A50C: 0x82 错误分支 (状态非法)
- 0x1A1F8: gMcuUpgradeMsgQueue NULL 检查 (cbnz r0, ...)
"""
import pexpect
import sys
import time
import re

JLINK_EXE = "/Applications/SEGGER/JLink_V942/JLinkExe"

def send_cmd(child, cmd, timeout=10):
    child.sendline(cmd)
    child.expect(r'J-Link>', timeout=timeout)
    return child.before

def parse_reg(regs_output, reg_name):
    for line in regs_output.splitlines():
        if line.startswith(reg_name + ' ='):
            val = line.split('=')[1].strip().split(',')[0]
            return int(val, 16)
    return None

def read_mem32(child, addr):
    out = send_cmd(child, f"mem32 0x{addr:08X},1", timeout=5)
    for line in out.splitlines():
        m = re.search(rf'{addr:08X}\s*=\s*([0-9A-Fa-f]+)', line)
        if m:
            return int(m.group(1), 16)
        m = re.search(rf'{addr:08X}\s+([0-9A-Fa-f]+)', line)
        if m:
            return int(m.group(1), 16)
    return None

def main():
    print("[*] 启动 JLinkExe...")
    child = pexpect.spawn(JLINK_EXE, encoding='utf-8', timeout=30)
    child.logfile_read = sys.stdout

    child.expect(r'J-Link>', timeout=15)
    send_cmd(child, "device HC32A4A8xI")
    send_cmd(child, "si SWD")
    send_cmd(child, "speed 4000")
    child.sendline("connect")
    idx = child.expect([r'Cortex-M4 identified', r'J-Link>'], timeout=30)
    if idx == 0:
        child.expect(r'J-Link>', timeout=10)
    print("\n[+] 目标已连接")

    send_cmd(child, "h")
    try:
        send_cmd(child, "ClrBP *", timeout=3)
    except:
        pass

    # 设置关键断点
    bps = [
        (0x0001A3A2, "0x82 命令分发 cmp r0,#130"),
        (0x0001A4A4, "0x82 处理入口"),
        (0x0001A4F6, "0x82 状态检查 ldrb r0,[r4,#0]"),
        (0x0001A50C, "0x82 错误分支 (状态非法)"),
        (0x0001A1F8, "gMcuUpgradeMsgQueue NULL 检查"),
    ]
    for addr, desc in bps:
        send_cmd(child, f"SetBP 0x{addr:08X}")
        print(f"[+] 断点: 0x{addr:08X} ({desc})")

    # 读取关键静态内存
    print("\n[*] === 读取关键内存（运行前快照）===")
    state_addr = 0x1FFE999C
    val = read_mem32(child, state_addr)
    print(f"[MEM] 0x{state_addr:08X} (upgrade state ptr) = 0x{val:08X}")
    if val:
        state_byte = read_mem32(child, val)
        print(f"[MEM] [0x{val:08X}] (state byte) = 0x{state_byte:08X}")

    # r5加载的地址 = 0x1FFE19B8 (从0x1A354的ldr r5,[pc,#...]得知)
    msgq_ptr_addr = 0x1FFE19B8
    val = read_mem32(child, msgq_ptr_addr)
    print(f"[MEM] 0x{msgq_ptr_addr:08X} (gMcuUpgradeMsgQueue ptr) = 0x{val:08X}")
    if val:
        msgq_val = read_mem32(child, val)
        print(f"[MEM] [0x{val:08X}] (queue handle) = 0x{msgq_val:08X}")
    else:
        print("[!] >>> gMcuUpgradeMsgQueue 指针为 NULL！")

    # 运行目标，等待断点命中
    send_cmd(child, "g")
    print("\n[*] 目标运行中，请执行: ./mcu_upgrade_test")
    print("[*] 等待断点命中...\n")

    hit_count = 0
    max_hits = 10
    while hit_count < max_hits:
        try:
            child.expect(r'(Breakpoint|Halt| halted)', timeout=120)
        except pexpect.exceptions.TIMEOUT:
            print("\n[!] 超时：120秒内无断点命中。")
            break

        time.sleep(0.3)
        regs = send_cmd(child, "regs", timeout=5)
        pc = parse_reg(regs, 'PC')

        desc = "未知"
        for addr, d in bps:
            if abs(pc - addr) <= 2:
                desc = d
                break

        print(f"\n[!] ====== 断点 #{hit_count+1} 命中: PC=0x{pc:08X} ({desc}) ======")
        r0 = parse_reg(regs, 'R0')
        r1 = parse_reg(regs, 'R1')
        r4 = parse_reg(regs, 'R4')
        r5 = parse_reg(regs, 'R5')
        lr = parse_reg(regs, 'R14(LR)')
        print(f"[REG] R0=0x{r0:08X} R1=0x{r1:08X} R4=0x{r4:08X} R5=0x{r5:08X} LR=0x{lr:08X}")

        # 如果命中状态检查或错误分支，读取状态字节
        if pc in [0x1A4F6, 0x1A50C] and r4 is not None:
            state_val = read_mem32(child, r4)
            print(f"[MEM] [R4]=[0x{r4:08X}] (state byte) = 0x{state_val:08X}")
            if state_val and (state_val & 0xFF) == 0:
                print("[!] >>> 状态为 0，0x82 将被拒绝！")

        # 如果命中NULL检查
        if pc == 0x1A1F8 and r5 is not None:
            msgq = read_mem32(child, r5)
            print(f"[MEM] [R5]=[0x{r5:08X}] (queue ptr) = 0x{msgq:08X}")
            if msgq == 0:
                print("[!] >>> 确认：gMcuUpgradeMsgQueue 为 NULL！")

        hit_count += 1
        send_cmd(child, "g")
        print("[*] 继续运行...")

    print("\n[*] === 清理 ===")
    send_cmd(child, "h")
    send_cmd(child, "ClrBP *", timeout=3)
    send_cmd(child, "g", timeout=3)
    child.sendline("q")
    try:
        child.expect(pexpect.EOF, timeout=5)
    except:
        child.terminate(force=True)
    print("[+] 完成")

if __name__ == "__main__":
    main()