#!/usr/bin/env python3
"""
Genpix SkyWalker-1 firmware probe and dump tool.

The SkyWalker-1 uses a Cypress FX2 (EZ-USB) microcontroller.
FX2 devices support reading internal RAM (8KB at 0x0000-0x1FFF)
and external RAM via standard vendor requests:
  - bRequest=0xA0 (FX2 firmware load/read)
  - wValue=address, wIndex=0

This tool also queries Genpix-specific vendor commands to gather
device info before attempting a firmware dump.
"""

import sys
import struct
import argparse
from datetime import datetime

try:
    import usb.core
    import usb.util
except ImportError:
    print("pyusb required: pip install pyusb")
    sys.exit(1)

VENDOR_ID = 0x09C0
PRODUCT_ID = 0x0203

# Genpix vendor commands (from SkyWalker1Control.h)
CMD_GET_USB_SPEED = 0x07
CMD_FW_VERSION_READ = 0x0B
CMD_VENDOR_STRING_READ = 0x0C
CMD_PRODUCT_STRING_READ = 0x0D
CMD_RESET_FX2 = 0x13
CMD_FW_BCD_VERSION_READ = 0x14
CMD_GET_8PSK_CONFIG = 0x80
CMD_GET_SIGNAL_STRENGTH = 0x87
CMD_GET_SIGNAL_LOCK = 0x90
CMD_GET_SERIAL_NUMBER = 0x93

# FX2 standard vendor request for RAM access
FX2_RAM_REQUEST = 0xA0

# FX2 memory map
FX2_INTERNAL_RAM_SIZE = 0x2000  # 8KB internal RAM
FX2_EXTERNAL_RAM_SIZE = 0x10000  # Up to 64KB external

# Config status bits
CONFIG_BITS = {
    0x01: "8PSK Started",
    0x02: "BCM4500 FW Loaded",
    0x04: "Intersil LNB On",
    0x08: "DVB Mode",
    0x10: "22kHz Tone",
    0x20: "18V Selected",
    0x40: "DC Tuned",
    0x80: "Armed (streaming)",
}


def find_device():
    dev = usb.core.find(idVendor=VENDOR_ID, idProduct=PRODUCT_ID)
    if dev is None:
        print("SkyWalker-1 not found. Is it plugged in?")
        sys.exit(1)
    return dev


def vendor_in(dev, request, value=0, index=0, length=64, timeout=2000):
    """Send a vendor IN control transfer (device-to-host)."""
    try:
        return dev.ctrl_transfer(
            usb.util.CTRL_TYPE_VENDOR | usb.util.CTRL_IN,
            request, value, index, length, timeout
        )
    except usb.core.USBError as e:
        return None


def detach_kernel_driver(dev):
    """Detach kernel driver if attached."""
    for cfg in dev:
        for intf in cfg:
            if dev.is_kernel_driver_active(intf.bInterfaceNumber):
                try:
                    dev.detach_kernel_driver(intf.bInterfaceNumber)
                    print(f"  Detached kernel driver from interface {intf.bInterfaceNumber}")
                    return intf.bInterfaceNumber
                except usb.core.USBError as e:
                    print(f"  Warning: Could not detach kernel driver: {e}")
                    print("  Try running with sudo, or: sudo modprobe -r dvb_usb_gp8psk")
                    sys.exit(1)
    return None


def probe_device_info(dev):
    """Query all known Genpix info commands."""
    print("\n=== Genpix SkyWalker-1 Device Info ===\n")

    # Firmware version (6 bytes)
    data = vendor_in(dev, CMD_FW_VERSION_READ, length=6)
    if data is not None and len(data) == 6:
        fw_int = (data[2] << 16) | (data[1] << 8) | data[0]
        build_date = f"20{data[5]:02d}-{data[4]:02d}-{data[3]:02d}"
        print(f"  FW Version:  {data[2]}.{data[1]:02d}.{data[0]} (0x{fw_int:06x})")
        print(f"  FW Build:    {build_date}")
    else:
        print(f"  FW Version:  (failed: {data})")

    # BCD version
    data = vendor_in(dev, CMD_FW_BCD_VERSION_READ, length=2)
    if data is not None:
        print(f"  BCD Version: {bytes(data).hex()}")

    # Vendor string
    data = vendor_in(dev, CMD_VENDOR_STRING_READ, length=64)
    if data is not None:
        s = bytes(data).rstrip(b'\x00').decode('ascii', errors='replace')
        print(f"  Vendor:      {s}")

    # Product string
    data = vendor_in(dev, CMD_PRODUCT_STRING_READ, length=64)
    if data is not None:
        s = bytes(data).rstrip(b'\x00').decode('ascii', errors='replace')
        print(f"  Product:     {s}")

    # USB speed
    data = vendor_in(dev, CMD_GET_USB_SPEED, length=1)
    if data is not None:
        speeds = {0: "Low", 1: "Full (12Mbps)", 2: "High (480Mbps)"}
        print(f"  USB Speed:   {speeds.get(data[0], f'Unknown ({data[0]})')}")

    # Serial number
    data = vendor_in(dev, CMD_GET_SERIAL_NUMBER, length=8)
    if data is not None:
        print(f"  Serial:      {bytes(data).hex()} ({bytes(data).rstrip(b'\\x00').decode('ascii', errors='replace')})")

    # 8PSK config/status
    data = vendor_in(dev, CMD_GET_8PSK_CONFIG, length=1)
    if data is not None:
        status = data[0]
        print(f"  Config:      0x{status:02x}")
        for bit, desc in CONFIG_BITS.items():
            state = "ON" if status & bit else "off"
            print(f"               [{state:>3}] {desc}")

    print()


def dump_fx2_ram(dev, output_file, start=0x0000, size=FX2_INTERNAL_RAM_SIZE, chunk=64):
    """
    Attempt to read FX2 internal RAM using the standard FX2 vendor request 0xA0.

    The Cypress FX2 bootloader/firmware typically supports:
    - bRequest = 0xA0
    - wValue = start address
    - wIndex = 0
    - Direction = IN (device to host)
    """
    print(f"=== Attempting FX2 RAM dump: 0x{start:04X} - 0x{start+size-1:04X} ({size} bytes) ===\n")

    firmware = bytearray()
    addr = start
    errors = 0
    consecutive_errors = 0

    while addr < start + size:
        remaining = (start + size) - addr
        read_len = min(chunk, remaining)

        data = vendor_in(dev, FX2_RAM_REQUEST, value=addr, index=0, length=read_len)

        if data is None:
            errors += 1
            consecutive_errors += 1
            firmware.extend(b'\xff' * read_len)
            if consecutive_errors >= 5:
                print(f"\n  Stopped: {consecutive_errors} consecutive read failures at 0x{addr:04X}")
                print("  Device may not support FX2 RAM readback (EEPROM firmware)")
                break
        else:
            consecutive_errors = 0
            firmware.extend(data)

        if (addr - start) % 0x400 == 0:
            pct = ((addr - start) / size) * 100
            print(f"  0x{addr:04X} [{pct:5.1f}%]  {'OK' if data is not None else 'FAIL'}", end='\r')

        addr += read_len

    print(f"\n\n  Read {len(firmware)} bytes, {errors} chunk errors")

    if firmware and any(b != 0xFF for b in firmware):
        with open(output_file, 'wb') as f:
            f.write(firmware)
        print(f"  Saved to: {output_file}")

        # Quick analysis
        non_ff = sum(1 for b in firmware if b != 0xFF)
        non_zero = sum(1 for b in firmware if b != 0x00)
        print(f"  Non-0xFF bytes: {non_ff}/{len(firmware)}")
        print(f"  Non-0x00 bytes: {non_zero}/{len(firmware)}")

        # Check for FX2 reset vector
        if len(firmware) >= 3:
            print(f"  First 16 bytes: {firmware[:16].hex(' ')}")
            if firmware[0] == 0x02:
                jump_addr = (firmware[1] << 8) | firmware[2]
                print(f"  Reset vector: LJMP 0x{jump_addr:04X} (typical FX2 firmware)")
    else:
        print("  No valid data read — dump appears empty")

    return firmware


def scan_vendor_commands(dev, start=0x00, end=0xFF):
    """Brute-force scan all vendor IN commands to find undocumented ones."""
    print(f"=== Scanning vendor commands 0x{start:02X}-0x{end:02X} ===\n")
    found = []
    for cmd in range(start, end + 1):
        data = vendor_in(dev, cmd, length=64, timeout=500)
        if data is not None and len(data) > 0:
            preview = bytes(data[:16]).hex(' ')
            is_known = cmd in (
                CMD_GET_USB_SPEED, CMD_FW_VERSION_READ, CMD_VENDOR_STRING_READ,
                CMD_PRODUCT_STRING_READ, CMD_FW_BCD_VERSION_READ, CMD_GET_8PSK_CONFIG,
                CMD_GET_SIGNAL_STRENGTH, CMD_GET_SIGNAL_LOCK, CMD_GET_SERIAL_NUMBER,
                FX2_RAM_REQUEST,
            )
            marker = " [KNOWN]" if is_known else " [NEW!]"
            print(f"  0x{cmd:02X}: [{len(data):3d} bytes] {preview}...{marker}")
            found.append((cmd, data))
    print(f"\n  Found {len(found)} responding commands")
    return found


def main():
    parser = argparse.ArgumentParser(description="Genpix SkyWalker-1 firmware probe/dump tool")
    parser.add_argument('--info', action='store_true', help="Query device info")
    parser.add_argument('--dump', metavar='FILE', help="Dump FX2 RAM to file")
    parser.add_argument('--scan', action='store_true', help="Scan all vendor commands")
    parser.add_argument('--start', type=lambda x: int(x, 0), default=0x0000,
                        help="RAM dump start address (default: 0x0000)")
    parser.add_argument('--size', type=lambda x: int(x, 0), default=FX2_INTERNAL_RAM_SIZE,
                        help=f"RAM dump size (default: 0x{FX2_INTERNAL_RAM_SIZE:X})")
    parser.add_argument('--external', action='store_true',
                        help="Try to dump external RAM (64KB)")
    args = parser.parse_args()

    if not any([args.info, args.dump, args.scan]):
        args.info = True
        args.scan = True

    print(f"Genpix SkyWalker-1 Firmware Tool")
    print(f"{'=' * 40}")

    dev = find_device()
    print(f"\nFound device: Bus {dev.bus} Addr {dev.address}")
    intf = detach_kernel_driver(dev)

    try:
        dev.set_configuration()
    except usb.core.USBError:
        pass  # May already be configured

    try:
        if args.info:
            probe_device_info(dev)

        if args.scan:
            scan_vendor_commands(dev)
            print()

        if args.dump:
            if args.external:
                dump_fx2_ram(dev, args.dump, args.start, FX2_EXTERNAL_RAM_SIZE)
            else:
                dump_fx2_ram(dev, args.dump, args.start, args.size)

    finally:
        if intf is not None:
            try:
                usb.util.release_interface(dev, intf)
                dev.attach_kernel_driver(intf)
                print("\nRe-attached kernel driver")
            except usb.core.USBError:
                print("\nNote: Run 'sudo modprobe dvb_usb_gp8psk' to reload driver")


if __name__ == '__main__':
    main()