#!/usr/bin/env python3
"""
XABA v2 Assembly Extractor
Extracts .NET DLLs from Xamarin XABA v2.x format blob files
"""

import struct
import os
import sys

def extract_xaba_v2(blob_path, output_dir):
    """Extract assemblies from XABA v2.x format"""

    with open(blob_path, 'rb') as f:
        # Skip ELF header (first 0x4000 bytes)
        f.seek(0x4000)
        data = f.read()

    # Check XABA magic
    magic = data[0:4]
    if magic != b'XABA':
        print(f"Error: Not a XABA file (magic: {magic})")
        return

    # Parse header
    version = struct.unpack('<I', data[4:8])[0]
    local_entry_count = struct.unpack('<I', data[8:12])[0]
    global_entry_count = struct.unpack('<I', data[12:16])[0]
    store_id = struct.unpack('<I', data[16:20])[0]

    print(f"XABA Format v{version >> 16}.{version & 0xFFFF}")
    print(f"Store ID: {store_id}")
    print(f"Local entries: {local_entry_count}")
    print(f"Global entries: {global_entry_count}")
    print()

    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    # The index starts after header
    # Based on observed format, try different offsets
    # XABA v2.2 appears to have entries listed after header

    # Search for DLL names in the data to find the index
    # Look for common DLL name pattern
    dll_markers = [b'OneControl.dll\x00', b'Plugin.BLE.dll\x00', b'App.LippertConnect.dll\x00']

    index_start = None
    for marker in dll_markers:
        pos = data.find(marker)
        if pos != -1 and pos < 100000:  # Should be early in file
            # Found a name, work backwards to find structure start
            # Names are typically after some fixed-size header fields
            index_start = pos - 200  # Rough estimate
            break

    if index_start:
        print(f"Found assembly index around offset {hex(index_start)}")
    else:
        print("Could not locate assembly index, using fallback method...")
        # Fallback: extract all PE files
        extract_all_pe_files(data, output_dir)
        return

    # Try to parse entries (this is approximate)
    # Format appears to be variable, so we'll use PE extraction as fallback
    extract_all_pe_files(data, output_dir)

def extract_all_pe_files(data, output_dir):
    """Extract all valid PE/DLL files from data"""

    print("Scanning for PE/DLL files...")
    dll_count = 0
    pos = 0
    extracted_names = {}

    while pos < len(data) - 64:
        # Look for MZ header
        pos = data.find(b'MZ', pos + 1)
        if pos == -1:
            break

        try:
            # Verify PE signature
            pe_offset = struct.unpack('<I', data[pos+60:pos+64])[0]
            if pos + pe_offset + 4 < len(data):
                if data[pos+pe_offset:pos+pe_offset+4] == b'PE\x00\x00':
                    # Valid PE file found

                    # Try to find assembly name in the PE
                    # Look for strings near the PE header
                    name_search = data[max(0, pos-1000):pos+5000]
                    assembly_name = None

                    # Look for .dll in nearby strings
                    for match_pos in range(len(name_search)):
                        if name_search[match_pos:match_pos+4] == b'.dll':
                            # Found .dll, extract name before it
                            start = match_pos
                            while start > 0 and name_search[start-1:start].isalnum() or name_search[start-1:start] in [b'.', b'_', b'-']:
                                start -= 1
                            possible_name = name_search[start:match_pos+4].decode('utf-8', errors='ignore')
                            if 20 < len(possible_name) < 150 and possible_name.endswith('.dll'):
                                assembly_name = possible_name
                                break

                    # Estimate size
                    next_mz = data.find(b'MZ', pos + 10000)
                    if next_mz == -1:
                        size = min(5000000, len(data) - pos)
                    else:
                        size = min(next_mz - pos, 5000000)

                    # Extract
                    dll_data = data[pos:pos+size]

                    if assembly_name and assembly_name not in extracted_names:
                        filename = assembly_name
                        extracted_names[assembly_name] = True
                    else:
                        filename = f"assembly_{dll_count:03d}.dll"

                    output_path = os.path.join(output_dir, filename)

                    with open(output_path, 'wb') as out:
                        out.write(dll_data)

                    print(f"[{dll_count+1}] {filename} ({size:,} bytes)")
                    dll_count += 1

        except Exception as e:
            pass

    print(f"\nExtracted {dll_count} assemblies to {output_dir}/")

    # List important ones
    print("\nKey assemblies extracted:")
    for fname in sorted(os.listdir(output_dir)):
        if any(x in fname for x in ['OneControl', 'Plugin.BLE', 'IDS', 'LCI']):
            size = os.path.getsize(os.path.join(output_dir, fname))
            print(f"  {fname:50s} ({size:>10,} bytes)")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 extract_xaba_v2.py <blob_file> [output_dir]")
        sys.exit(1)

    blob_file = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "extracted_dlls"

    extract_xaba_v2(blob_file, output_dir)