AeThex-OS/direct_flash.py

"""
Direct USB fastboot flash using PyUSB
Bypasses fastboot.exe which can't find MediaTek devices
"""
import os
import sys
import struct
import time

# Set libusb DLL path before importing pyusb
os.environ['PATH'] = r'C:\Users\PCOEM\AppData\Local\Programs\Python\Python312\Lib\site-packages\libusb\_platform\windows\x86_64' + os.pathsep + os.environ.get('PATH', '')

import usb.core
import usb.util
import usb.backend.libusb1 as libusb1

# Find the backend explicitly
backend = libusb1.get_backend(find_library=lambda x: r'C:\Users\PCOEM\AppData\Local\Programs\Python\Python312\Lib\site-packages\libusb\_platform\windows\x86_64\libusb-1.0.dll')

# MediaTek Fastboot
VID = 0x0E8D
PID = 0x201C

BOOT_IMG = r"C:\Users\PCOEM\AeThexOS\magisk_patched_boot.img"

def find_device():
    """Find the MediaTek device"""
    print(f"Looking for USB device VID={VID:04X} PID={PID:04X}...")
    dev = usb.core.find(idVendor=VID, idProduct=PID, backend=backend)
    if dev is None:
        print("Device not found!")
        # List all USB devices
        print("\nAll USB devices:")
        for d in usb.core.find(find_all=True, backend=backend):
            print(f"  VID={d.idVendor:04X} PID={d.idProduct:04X} - {d.manufacturer or '?'} {d.product or '?'}")
        return None
    print(f"Found: {dev.manufacturer or '?'} {dev.product or '?'}")
    return dev

def fastboot_command(dev, ep_out, ep_in, cmd):
    """Send a fastboot command and get response"""
    print(f"  >> {cmd}")
    dev.write(ep_out, cmd.encode('utf-8'))

    responses = []
    while True:
        try:
            data = dev.read(ep_in, 64, timeout=5000)
            resp = bytes(data).decode('utf-8', errors='replace')
            print(f"  << {resp}")

            if resp.startswith('OKAY'):
                return ('OKAY', resp[4:])
            elif resp.startswith('FAIL'):
                return ('FAIL', resp[4:])
            elif resp.startswith('DATA'):
                size = int(resp[4:], 16)
                return ('DATA', size)
            elif resp.startswith('INFO'):
                responses.append(resp[4:])
                continue
            else:
                return ('UNKNOWN', resp)
        except usb.core.USBError as e:
            if e.errno == 110:  # timeout
                return ('TIMEOUT', '')
            raise

def fastboot_flash(dev, ep_out, ep_in, partition, img_path):
    """Flash an image to a partition"""
    img_size = os.path.getsize(img_path)
    print(f"\nFlashing {partition} ({img_size} bytes)...")

    # Step 1: Get max download size
    status, val = fastboot_command(dev, ep_out, ep_in, "getvar:max-download-size")
    if status == 'OKAY':
        max_size = int(val, 16)
        print(f"  Max download size: {max_size}")
    else:
        max_size = img_size  # assume it fits

    # Step 2: Download command
    status, val = fastboot_command(dev, ep_out, ep_in, f"download:{img_size:08x}")
    if status != 'DATA':
        print(f"  Download command failed: {status} {val}")
        return False

    # Step 3: Send image data in chunks
    CHUNK_SIZE = 512 * 1024  # 512KB chunks
    sent = 0
    with open(img_path, 'rb') as f:
        while sent < img_size:
            chunk = f.read(CHUNK_SIZE)
            if not chunk:
                break
            dev.write(ep_out, chunk, timeout=30000)
            sent += len(chunk)
            pct = sent * 100 // img_size
            print(f"\r  Sending: {pct}% ({sent}/{img_size})", end='', flush=True)
    print()

    # Step 4: Wait for OKAY after data transfer
    status, val = fastboot_command(dev, ep_out, ep_in, "")
    # Actually, the device should respond OKAY after receiving all data
    try:
        data = dev.read(ep_in, 64, timeout=10000)
        resp = bytes(data).decode('utf-8', errors='replace')
        print(f"  << {resp}")
        if not resp.startswith('OKAY'):
            print(f"  Data transfer response: {resp}")
    except:
        pass

    # Step 5: Flash command
    status, val = fastboot_command(dev, ep_out, ep_in, f"flash:{partition}")
    if status == 'OKAY':
        print(f"  Flash {partition} SUCCESS!")
        return True
    else:
        print(f"  Flash failed: {status} {val}")
        return False

def main():
    print("=== Direct USB Fastboot Flash ===")
    print()

    if not os.path.exists(BOOT_IMG):
        print(f"ERROR: {BOOT_IMG} not found")
        return 1

    dev = find_device()
    if dev is None:
        return 1

    # Print device config
    print(f"\nDevice configuration:")
    for cfg in dev:
        print(f"  Config {cfg.bConfigurationValue}: {cfg.bNumInterfaces} interface(s)")
        for intf in cfg:
            print(f"    Interface {intf.bInterfaceNumber}: Class={intf.bInterfaceClass:02X} SubClass={intf.bInterfaceSubClass:02X} Protocol={intf.bInterfaceProtocol:02X}")
            for ep in intf:
                direction = "IN" if usb.util.endpoint_direction(ep.bEndpointAddress) == usb.util.ENDPOINT_IN else "OUT"
                print(f"      EP {ep.bEndpointAddress:02X} ({direction}): MaxPacket={ep.wMaxPacketSize}")

    # Claim the interface
    try:
        if dev.is_kernel_driver_active(0):
            dev.detach_kernel_driver(0)
    except:
        pass

    # Set configuration
    try:
        dev.set_configuration()
    except usb.core.USBError:
        pass

    # Get the interface
    cfg = dev.get_active_configuration()
    intf = cfg[(0, 0)]

    # Find endpoints
    ep_out = usb.util.find_descriptor(intf, custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_OUT)
    ep_in = usb.util.find_descriptor(intf, custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_IN)

    if ep_out is None or ep_in is None:
        print("ERROR: Could not find USB endpoints")
        return 1

    print(f"\nEndpoints: OUT={ep_out.bEndpointAddress:02X} IN={ep_in.bEndpointAddress:02X}")

    # Test connection with getvar
    print("\nTesting fastboot protocol...")
    status, val = fastboot_command(dev, ep_out, ep_in, "getvar:product")
    if status == 'TIMEOUT':
        print("No response - device may not support fastboot protocol on this interface")
        print("Trying ADB protocol instead...")
        # Try ADB handshake
        ADB_CNXN = b'CNXN'
        ADB_VERSION = 0x01000001
        ADB_MAXDATA = 256 * 1024
        msg = struct.pack('<4sIII', ADB_CNXN, ADB_VERSION, ADB_MAXDATA, len(b'host::\0'))
        # ... this would be complex, skip for now
        return 1

    print(f"\nProduct: {val}")

    # Get more info
    for var in ['serialno', 'product', 'secure', 'unlocked', 'slot-count', 'current-slot']:
        status, val = fastboot_command(dev, ep_out, ep_in, f"getvar:{var}")

    # Flash boot_b
    success = fastboot_flash(dev, ep_out, ep_in, "boot_b", BOOT_IMG)

    if success:
        print("\n=== Rebooting ===")
        fastboot_command(dev, ep_out, ep_in, "reboot")
        print("\nDone! Magisk root should be active after boot.")
    else:
        # Try just "boot" without slot suffix
        print("\nRetrying with partition name 'boot'...")
        success = fastboot_flash(dev, ep_out, ep_in, "boot", BOOT_IMG)
        if success:
            print("\n=== Rebooting ===")
            fastboot_command(dev, ep_out, ep_in, "reboot")

    return 0 if success else 1

if __name__ == '__main__':
    sys.exit(main())