import os
import time
import subprocess

# Settings
START_ADDR = 0xFFF00000
END_ADDR   = 0xFFF10000
OUTPUT_FILE = os.path.abspath("jibo_bootrom_full.bin")
CHUNK_FILE = os.path.abspath("chunk.bin")
TOOL_PATH = os.path.abspath("./shofel2_t124")

def run_cmd(cmd):
    # Added stderr capture to help debug if the tool itself errors out
    return subprocess.run(cmd, shell=True, capture_output=True, text=True)

print(f"[*] Starting leak from 0x{START_ADDR:08X} to 0x{END_ADDR:08X}")

with open(OUTPUT_FILE, "ab") as f:
    for addr in range(START_ADDR, END_ADDR, 4):
        # ...

        # 1. PUSH PAYLOAD
        # Add a 1-second delay before pushing to let USB stabilize
        time.sleep(1.0)
        run_cmd(f"{TOOL_PATH} PAYLOAD jibo_leaker.bin arm")

        # 2. WAIT FOR RESET
        print("[*] Waiting for device to cycle...", end="\r")
        # Ensure the device actually disappears and reappears
        while b"0955:7740" in subprocess.run("lsusb", shell=True, capture_output=True).stdout:
            time.sleep(0.1) # Wait for it to disconnect

        while b"0955:7740" not in subprocess.run("lsusb", shell=True, capture_output=True).stdout:
            time.sleep(0.5) # Wait for it to reconnect

        time.sleep(1.5) # Crucial "Cool Down" for the BootROM stack

        # 3. DUMP
        dump_res = run_cmd(f"{TOOL_PATH} MEM_DUMP 0x7000e450 0x4 {CHUNK_FILE}")
        # 4. Save to file
        if os.path.exists(CHUNK_FILE):
            with open(CHUNK_FILE, "rb") as chunk:
                data = chunk.read()
                if len(data) == 4:
                    f.write(data)
                    f.flush()
            os.remove(CHUNK_FILE)
        else:
            print(f"\n[!] Error: chunk.bin not found at 0x{addr:08X}")
            print(f"[!] Tool Output: {dump_res.stdout}")
            print(f"[!] Tool Error: {dump_res.stderr}")
            break