#!/usr/bin/env python3
"""
Run a Windows PE under Wine and dump its process memory after unpacking.

Launches the EXE, waits for it to unpack, then reads /proc/PID/mem
guided by /proc/PID/maps to capture the unpacked code and data sections.
Searches the dump for FX2 firmware signatures.
"""
import subprocess
import time
import os
import sys
import signal
import struct
import re
import glob


def find_wine_pid(exe_basename, timeout=10):
    """Find the Wine process PID by looking for the .exe in /proc."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        for pid_dir in glob.glob('/proc/[0-9]*'):
            try:
                cmdline = open(f'{pid_dir}/cmdline', 'rb').read()
                if exe_basename.lower().encode() in cmdline.lower():
                    pid = int(os.path.basename(pid_dir))
                    # Skip if it's our own python process
                    if pid == os.getpid():
                        continue
                    return pid
            except (PermissionError, FileNotFoundError, ProcessLookupError):
                continue
        time.sleep(0.2)
    return None


def dump_process_memory(pid, output_dir):
    """Dump all readable memory regions of a process."""
    maps_path = f'/proc/{pid}/maps'
    mem_path = f'/proc/{pid}/mem'

    regions = []
    try:
        with open(maps_path, 'r') as f:
            for line in f:
                parts = line.split()
                addr_range = parts[0]
                perms = parts[1]
                # Only dump readable regions
                if 'r' not in perms:
                    continue
                start_s, end_s = addr_range.split('-')
                start = int(start_s, 16)
                end = int(end_s, 16)
                size = end - start
                # Skip huge regions (> 64MB) and tiny ones
                if size > 64 * 1024 * 1024 or size < 64:
                    continue
                pathname = parts[5].strip() if len(parts) > 5 else ""
                regions.append((start, end, perms, pathname))
    except (PermissionError, FileNotFoundError) as e:
        print(f"  Cannot read maps: {e}")
        return None

    print(f"  Found {len(regions)} readable memory regions")

    all_data = bytearray()
    region_info = []

    try:
        with open(mem_path, 'rb') as mem:
            for start, end, perms, pathname in regions:
                size = end - start
                try:
                    mem.seek(start)
                    chunk = mem.read(size)
                    offset_in_dump = len(all_data)
                    all_data.extend(chunk)
                    region_info.append({
                        'va_start': start,
                        'va_end': end,
                        'perms': perms,
                        'pathname': pathname,
                        'dump_offset': offset_in_dump,
                        'size': len(chunk)
                    })
                except (OSError, ValueError):
                    pass
    except PermissionError as e:
        print(f"  Cannot read mem: {e}")
        print("  Try running with sudo or as the same user as Wine")
        return None

    # Save full dump
    dump_file = os.path.join(output_dir, 'wine_memdump.bin')
    with open(dump_file, 'wb') as f:
        f.write(all_data)
    print(f"  Saved {len(all_data)} bytes to {dump_file}")

    # Save region map
    map_file = os.path.join(output_dir, 'wine_memdump_regions.txt')
    with open(map_file, 'w') as f:
        for r in region_info:
            f.write(f"0x{r['va_start']:08X}-0x{r['va_end']:08X} "
                    f"{r['perms']:5s} dump_off=0x{r['dump_offset']:08X} "
                    f"size=0x{r['size']:06X} {r['pathname']}\n")
    print(f"  Saved region map to {map_file}")

    return all_data, region_info


def search_firmware(data, region_info):
    """Search dumped memory for FX2 firmware signatures."""
    print(f"\n{'=' * 50}")
    print("Searching for firmware signatures...")
    print(f"{'=' * 50}")

    # 1. C2 EEPROM header with Genpix VID
    print("\n[1] C2 EEPROM headers (C2 C0 09 03 02):")
    c2_genpix = bytes([0xC2, 0xC0, 0x09, 0x03, 0x02])
    pos = 0
    while True:
        idx = data.find(c2_genpix, pos)
        if idx < 0:
            break
        region = find_region(region_info, idx)
        ctx = bytes(data[idx:idx + 32])
        print(f"  0x{idx:08X} (VA: {region}): {ctx.hex(' ')}")
        # Parse the full C2 header
        if idx + 8 <= len(data):
            vid = data[idx + 1] | (data[idx + 2] << 8)
            pid = data[idx + 3] | (data[idx + 4] << 8)
            did = data[idx + 5] | (data[idx + 6] << 8)
            config = data[idx + 7]
            print(f"    VID=0x{vid:04X} PID=0x{pid:04X} DID=0x{did:04X} Config=0x{config:02X}")
        pos = idx + 1

    # 2. FX2 RAM clear init sequence
    print("\n[2] FX2 init sequence (78 7F E4 F6 D8 FD 75 81):")
    fx2_init = bytes([0x78, 0x7F, 0xE4, 0xF6, 0xD8, 0xFD, 0x75, 0x81])
    pos = 0
    while True:
        idx = data.find(fx2_init, pos)
        if idx < 0:
            break
        region = find_region(region_info, idx)
        ctx = bytes(data[max(0, idx - 4):idx + 16])
        print(f"  0x{idx:08X} (VA: {region}): {ctx.hex(' ')}")
        pos = idx + 1

    # 3. Partial RAM clear pattern
    print("\n[3] RAM clear pattern (78 7F E4 F6 D8 FD):")
    ram_clear = bytes([0x78, 0x7F, 0xE4, 0xF6, 0xD8, 0xFD])
    pos = 0
    hits = 0
    while True:
        idx = data.find(ram_clear, pos)
        if idx < 0:
            break
        region = find_region(region_info, idx)
        ctx = bytes(data[max(0, idx - 4):idx + 16])
        print(f"  0x{idx:08X} (VA: {region}): {ctx.hex(' ')}")
        hits += 1
        if hits >= 10:
            break
        pos = idx + 1

    # 4. LJMP at what could be code address 0x0000 (start of firmware)
    # Look for 02 XX XX where XX XX is 0x0100-0x3FFF
    print("\n[4] C2 load records (LEN_H LEN_L 00 00 02 = record at addr 0x0000):")
    for off in range(len(data) - 8):
        rec_len = (data[off] << 8) | data[off + 1]
        if 0x0100 <= rec_len <= 0x4000:
            if data[off + 2] == 0x00 and data[off + 3] == 0x00 and data[off + 4] == 0x02:
                target = (data[off + 5] << 8) | data[off + 6]
                if 0x0100 <= target <= 0x3FFF:
                    # Check if this looks like a valid C2 record chain
                    region = find_region(region_info, off)
                    ctx = bytes(data[off:off + 16])
                    # Also check 8 bytes before for C2 header
                    has_c2_header = (off >= 8 and data[off - 8] == 0xC2)
                    header_note = " ** C2 HEADER 8 BYTES BEFORE! **" if has_c2_header else ""
                    print(f"  0x{off:08X} (VA: {region}): len={rec_len} "
                          f"addr=0x0000 LJMP 0x{target:04X} -- {ctx.hex(' ')}{header_note}")

    # 5. Known VID/PID bytes near potential firmware data
    print("\n[5] VID 0x09C0 references:")
    vid_bytes = b'\xC0\x09'
    pos = 0
    hits = 0
    while True:
        idx = data.find(vid_bytes, pos)
        if idx < 0:
            break
        # Check if followed by PID within 4 bytes
        if idx + 4 < len(data):
            nearby = data[idx:idx + 8]
            if b'\x03\x02' in nearby:
                region = find_region(region_info, idx)
                ctx = bytes(data[max(0, idx - 4):idx + 16])
                print(f"  0x{idx:08X} (VA: {region}): {ctx.hex(' ')}")
                hits += 1
                if hits >= 20:
                    break
        pos = idx + 1

    # 6. Search for the firmware version string "2.13"
    print("\n[6] Version strings:")
    for pattern in [b'2.13', b'2.06', b'2.10', b'SkyWalker', b'Genpix',
                    b'8PSK', b'EEPROM', b'firmware', b'I2C']:
        pos = 0
        while True:
            idx = data.find(pattern, pos)
            if idx < 0:
                break
            region = find_region(region_info, idx)
            # Get surrounding context as ascii
            start = max(0, idx - 16)
            end = min(len(data), idx + 48)
            ctx_bytes = bytes(data[start:end])
            ctx_ascii = ctx_bytes.decode('ascii', errors='replace')
            ctx_ascii = re.sub(r'[^\x20-\x7e]', '.', ctx_ascii)
            print(f"  0x{idx:08X} (VA: {region}): '{ctx_ascii}'")
            pos = idx + 1

    # 7. Look for USB vendor request setup patterns
    # The updater will set bRequest=0x83 (I2C_WRITE) or 0xA0 to write firmware
    print("\n[7] USB transfer setup (IOCTL/vendor request patterns):")
    # WinUSB_ControlTransfer uses WINUSB_SETUP_PACKET:
    # RequestType(1), Request(1), Value(2), Index(2), Length(2)
    # For vendor OUT: RequestType=0x40, Request=0x83/0xA0
    for req_type, req, desc in [(0x40, 0xA0, "FX2 RAM write"),
                                 (0x40, 0x83, "I2C_WRITE"),
                                 (0x40, 0x84, "I2C_READ")]:
        pattern = bytes([req_type, req])
        pos = 0
        hits_count = 0
        while True:
            idx = data.find(pattern, pos)
            if idx < 0:
                break
            # Check if followed by reasonable wValue/wIndex
            if idx + 8 <= len(data):
                wval = struct.unpack_from('<H', data, idx + 2)[0]
                widx = struct.unpack_from('<H', data, idx + 4)[0]
                wlen = struct.unpack_from('<H', data, idx + 6)[0]
                if wlen > 0 and wlen < 0x4000:
                    region = find_region(region_info, idx)
                    print(f"  0x{idx:08X} ({desc}): "
                          f"ReqType=0x{req_type:02X} Req=0x{req:02X} "
                          f"wVal=0x{wval:04X} wIdx=0x{widx:04X} wLen=0x{wlen:04X} "
                          f"(VA: {region})")
                    hits_count += 1
                    if hits_count >= 10:
                        break
            pos = idx + 1


def find_region(region_info, dump_offset):
    """Find the VA region for a given dump offset."""
    for r in region_info:
        if r['dump_offset'] <= dump_offset < r['dump_offset'] + r['size']:
            va = r['va_start'] + (dump_offset - r['dump_offset'])
            return f"0x{va:08X} [{r['pathname'] or 'anon'}]"
    return "unknown"


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Wine memory dump for firmware extraction")
    parser.add_argument('exe', help='Windows PE executable to run under Wine')
    parser.add_argument('-o', '--output-dir', default='.',
                        help='Output directory for dumps')
    parser.add_argument('--wait', type=float, default=3.0,
                        help='Seconds to wait after launch for unpacking (default: 3)')
    parser.add_argument('--skip-launch', action='store_true',
                        help='Skip launching Wine, just attach to existing process')
    args = parser.parse_args()

    exe_path = os.path.abspath(args.exe)
    exe_basename = os.path.basename(exe_path)
    os.makedirs(args.output_dir, exist_ok=True)

    wine_proc = None
    pid = None

    if not args.skip_launch:
        print(f"Launching {exe_basename} under Wine...")
        # Use WINEDEBUG=-all to reduce noise
        env = os.environ.copy()
        env['WINEDEBUG'] = '-all'
        wine_proc = subprocess.Popen(
            ['wine', exe_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env=env
        )
        print(f"  Wine wrapper PID: {wine_proc.pid}")

        # Wait for the actual .exe process to appear
        print(f"  Waiting {args.wait}s for unpacking...")
        time.sleep(args.wait)

    # Find the Windows process PID
    print(f"  Looking for {exe_basename} process...")
    pid = find_wine_pid(exe_basename, timeout=5)
    if pid is None:
        # Try looking for wine-preloader or wine64-preloader
        print("  Couldn't find by exe name, searching all wine processes...")
        for pid_dir in glob.glob('/proc/[0-9]*'):
            try:
                cmdline = open(f'{pid_dir}/cmdline', 'rb').read()
                if b'wine' in cmdline.lower() and pid_dir != f'/proc/{os.getpid()}':
                    p = int(os.path.basename(pid_dir))
                    if wine_proc and p == wine_proc.pid:
                        continue
                    print(f"    Found wine process PID {p}: {cmdline[:100]}")
            except:
                pass

        if pid is None and wine_proc:
            pid = wine_proc.pid
            print(f"  Using Wine wrapper PID: {pid}")

    if pid:
        print(f"\n  Target PID: {pid}")
        result = dump_process_memory(pid, args.output_dir)
        if result:
            data, region_info = result
            search_firmware(data, region_info)
    else:
        print("  ERROR: Could not find process")

    # Cleanup
    if wine_proc:
        print("\nTerminating Wine process...")
        try:
            wine_proc.terminate()
            wine_proc.wait(timeout=5)
        except:
            wine_proc.kill()


if __name__ == '__main__':
    main()