#!/usr/bin/env python3
"""
JLink runtime patching for MCU upgrade CAN bypass.
Uses subprocess (no pexpect needed).
Usage:
  1. Flash ORIGINAL mcu_app.hex (MCU + Android boots)
  2. Connect J-Link
  3. python3 jlink_runtime_patch.py
  4. In another terminal: adb shell su root /data/local/tmp/mcu_upgrade_test
"""

import subprocess
import sys
import time

JLINK = "/Applications/SEGGER/JLink_V942/JLinkExe"
DEVICE = "HC32A4A8xI"

def jlink(commands, timeout=15):
    """Send commands to JLinkExe and return output"""
    full = f"device {DEVICE}\nsi SWD\nspeed 4000\nconnect\nh\n{commands}\ng\nq\n"
    result = subprocess.run(
        [JLINK],
        input=full,
        capture_output=True,
        text=True,
        timeout=timeout
    )
    return result.stdout

def read32(addr):
    """Read 32-bit value from MCU memory"""
    out = jlink(f"mem32 0x{addr:08X},1\nsleep 50")
    for line in out.splitlines():
        parts = line.strip().split()
        if len(parts) >= 2:
            try:
                a = int(parts[0].rstrip(':'), 16)
                if a == addr:
                    return int(parts[1], 16)
            except ValueError:
                continue
    return None

def write16(addr, val):
    """Write 16-bit value"""
    return jlink(f"w2 0x{addr:08X},0x{val:04X}\nsleep 50")

def write32(addr, val):
    """Write 32-bit value"""
    return jlink(f"w4 0x{addr:08X},0x{val:08X}\nsleep 50")

def main():
    print("=" * 60)
    print("JLink Runtime Patch - MCU Upgrade CAN Bypass")
    print("=" * 60)
    print()
    print("[*] Step 1: Reading MCU RAM state...")
    
    # Read key state variables
    g_queue = read32(0x1FFE19B8)
    print(f"    gMcuUpgradeMsgQueue @ 0x1FFE19B8 = 0x{g_queue:08X}" if g_queue else "    gMcuUpgradeMsgQueue = FAILED TO READ")
    
    can_handle_ptr = read32(0x1FFE5CF8)
    print(f"    CAN handle ptr @ 0x1FFE5CF8    = 0x{can_handle_ptr:08X}" if can_handle_ptr else "    CAN handle ptr = FAILED TO READ")
    
    if can_handle_ptr and can_handle_ptr != 0xFFFFFFFF and can_handle_ptr != 0:
        print(f"    CAN handle structure at 0x{can_handle_ptr:08X}")
        # Read handle fields
        for offset in [0, 4, 8, 12, 16, 18, 20, 24]:
            val = read32(can_handle_ptr + offset)
            if val is not None:
                print(f"      +0x{offset:02X}: 0x{val:08X}")
    else:
        print(f"    CAN handle ptr is NULL (0x{can_handle_ptr:08X})")
        print("    This means the handle was never allocated.")
        print("    Need to create a fake handle structure.")
    
    print()
    print("[*] Step 2: Reading current flash state at patch addresses...")
    for addr in [0x125B4, 0x124E8, 0x125D4, 0x125DE, 0x125FC]:
        out = jlink(f"mem16 0x{addr:08X},1\nsleep 50")
        for line in out.splitlines():
            if f"{addr:08X}" in line:
                print(f"    0x{addr:05X}: {line.strip()}")
    
    print()
    print("[*] Step 3: Would apply these patches:")
    print("    Flash (5 patches, safe after boot):")
    print("      0x125B4: movs r0,#1; bx lr")
    print("      0x124E8: cbz → nop")
    print("      0x125D4: cbz → nop")
    print("      0x125DE: beq → nop")
    print("      0x125FC: beq → nop")
    if can_handle_ptr and can_handle_ptr != 0 and can_handle_ptr != 0xFFFFFFFF:
        print(f"    RAM: CAN handle at 0x{can_handle_ptr:08X} → set fields to 'ready'")
    print()
    print("[*] Wait for user to review above, then type 'yes' to apply")
    
    # Don't auto-apply - let user see the state first
    print("\n[*] Done reading state. Run with --apply to patch.")
    print(f"    Or run: {JLINK} -device {DEVICE} -If SWD -Speed 4000 -CommanderScript runtime_patch.jlink")

if __name__ == '__main__':
    main()