import struct
import os
import lz4.block
import re

def extract():
    print("Reading payload.bin...")
    with open('payload.bin', 'rb') as f:
        data = f.read()

    # Find all XALZ offsets
    print("Scanning for XALZ blocks...")
    offsets = []
    pos = 0
    while True:
        pos = data.find(b'XALZ', pos)
        if pos == -1:
            break
        offsets.append(pos)
        pos += 4

    print(f"Found {len(offsets)} XALZ blocks.")
    
    os.makedirs('extracted_final', exist_ok=True)
    
    success_count = 0
    
    for i, offset in enumerate(offsets):
        # Determine end of compressed data
        if i < len(offsets) - 1:
            end = offsets[i+1]
        else:
            end = len(data)
            
        # Header parsing
        # XALZ (4) + ID (4) + UncompressedSize (4)
        header = data[offset:offset+12]
        magic, blob_id, uncomp_size = struct.unpack('<4sII', header)
        
        # print(f"Block {i}: Offset {offset}, Size {uncomp_size}")
        
        compressed_data = data[offset+12:end]
        
        try:
            # lz4.block.decompress expects the compressed data.
            # If uncompressed_size is provided, it helps allocation.
            decompressed = lz4.block.decompress(compressed_data, uncompressed_size=uncomp_size)
        except Exception as e:
            print(f"Error extracting block {i} at {offset}: {e}")
            # Try extracting just uncompressed if size matches?
            # If failed, maybe it wasn't compressed? But we saw MZ in literals.
            continue
                
        # Determine name
        name = f"assembly_{i:03d}.dll"
        
        # Scan for internal name
        # Search first 20KB for .dll names
        search_area = decompressed[:min(20000, len(decompressed))]
        # Regex: alphanumeric + . _ -
        matches = re.finditer(b'([a-zA-Z0-9_.-]+\.dll)', search_area)
        best_name = None
        for match in matches:
             try:
                 cand = match.group(1).decode('utf-8')
                 # Filter
                 if cand.lower() != 'mscoree.dll' and len(cand) > 4:
                     # Heuristic: usually starts with capital or specific words
                     if 'System' in cand or 'Microsoft' in cand or 'OneControl' in cand or 'Lippert' in cand or 'Plugin' in cand or 'Xamarin' in cand or 'Android' in cand:
                         best_name = cand
                         break
                     # Fallback
                     if not best_name:
                         best_name = cand
             except:
                 continue
        
        if best_name:
            name = best_name
            
        # Handle duplicates
        output_path = os.path.join('extracted_final', name)
        if os.path.exists(output_path):
             counter = 2
             base_name = name
             while True:
                 name = f"{base_name[:-4]}_{counter}.dll"
                 output_path = os.path.join('extracted_final', name)
                 if not os.path.exists(output_path):
                     break
                 counter += 1
                 
        with open(output_path, 'wb') as out:
            out.write(decompressed)
        
        success_count += 1
        if i % 50 == 0:
            print(f"Processed {i}/{len(offsets)}...")
            
    print(f"Successfully extracted {success_count} assemblies to extracted_final/")

    # List Key Assemblies
    print("\nKey Assemblies Found:")
    for fname in sorted(os.listdir('extracted_final')):
        if any(x in fname for x in ['OneControl', 'Plugin.BLE', 'IdsCan', 'MyRvLink']):
             sz = os.path.getsize(os.path.join('extracted_final', fname))
             print(f"  {fname} ({sz} bytes)")

if __name__ == '__main__':
    extract()