import usb.core
import usb.util
import time
import subprocess
import struct

# Tegra K1 Constants
ID_VENDOR = 0x0955
ID_PRODUCT = 0x7740

# Memory Range to Dump
START_ADDR = 0x40000000
END_ADDR   = 0x40040000 # 256KB total
STEP       = 4          # 4 bytes per reboot

def get_jibo():
    dev = usb.core.find(idVendor=ID_VENDOR, idProduct=ID_PRODUCT)
    if dev is not None:
        try:
            # Check if a kernel driver is using the device and detach it
            if dev.is_kernel_driver_active(0):
                dev.detach_kernel_driver(0)

            # Reset the device to clear any 'Busy' states from shofel2
            dev.reset()

            # Claim the interface
            usb.util.claim_interface(dev, 0)
        except Exception as e:
            print(f"Warning: Could not detach/claim: {e}")

    return dev

def update_config(addr):
    """Write the new target address to the header file."""
    with open("target_config.h", "w") as f:
        f.write(f"#define TARGET_ADDR {hex(addr)}\n")

def compile_payload():
    """Run the compilation script."""
    result = subprocess.run(["./compile4jibo.sh", "jibo_leaker.c"],
                            capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Compilation Failed:\n{result.stderr}")
        return False
    return True

def run_dump():
    print(f"--- Starting Jibo IRAM Deep Leak ---")
    current_addr = START_ADDR

    # Open file in 'append' mode so we don't lose data if we crash
    with open("jibo_iram_brute.bin", "ab") as f:
        while current_addr < END_ADDR:
            print(f"Targeting: {hex(current_addr)}...", end="\r")

            # 1. Prepare Payload
            update_config(current_addr)
            if not compile_payload():
                break

            # 2. Wait for Device
            dev = get_jibo()

            # 3. Inject Payload using your shofel2 tool
            # Note: We use subprocess to call the existing binary
            try:
                subprocess.run(["sudo", "./shofel2_t124", "jibo_leaker.bin"],
                               capture_output=True, timeout=5)
            except Exception as e:
                print(f"\nInjection timeout at {hex(current_addr)}, retrying...")
                continue

            # 4. Catch the data
            # Since the payload writes to PMC and resets, we need to read the
            # USB response that Jibo sends when he re-enters RCM.
            try:
                # Give it a moment to execute and reboot
                time.sleep(0.5)
                dev = get_jibo()

                # Reading the 16 bytes of the RCM header which may contain our leak
                # or status info. On T124, data is often on 0x81.
                data = dev.read(0x81, 16, timeout=1000)

                if data:
                    f.write(data)
                    f.flush() # Force write to disk

                current_addr += STEP

            except usb.core.USBError as e:
                # If it times out, Jibo might not have reset yet
                print(f"\nUSB Error at {hex(current_addr)}: {e}")
                time.sleep(1)
                continue

    print(f"\nDump finished. Data saved to jibo_iram_brute.bin")

if __name__ == "__main__":
    run_dump()