#!/usr/bin/env python3
"""
Repack app.asar from the decompiled source directory.
Usage: python3 repack-asar.py [source-dir] [output-file]
Default source-dir: directory containing this script's parent (airfieldz-decompiled)
Default output: app.asar in current directory

Requirement: pip install jsbeautifier (optionnel, pour beautifier le JSON)
"""

import struct, json, os, sys, hashlib

def align_up(x, align=4):
    return (x + align - 1) & ~(align - 1)

def is_leaf(node):
    """Check if a dict represents a file (leaf) or directory"""
    return isinstance(node, dict) and 'size' in node

def build_asar(source_dir, output_path):
    # Walk through source directory and build file list
    file_list = []
    excludes = {'scripts', 'README.md', 'app.asar'}
    
    for root, dirs, files in os.walk(source_dir):
        rel_root = os.path.relpath(root, source_dir)
        if rel_root == '.':
            dirs[:] = [d for d in dirs if d not in excludes]
        
        for name in sorted(files):
            full_path = os.path.join(root, name)
            rel_path = os.path.relpath(full_path, source_dir)
            
            skip = False
            for excl in excludes:
                if rel_path.startswith(excl):
                    skip = True
                    break
            if skip:
                continue
            
            with open(full_path, 'rb') as f:
                content = f.read()
            
            file_list.append((rel_path.replace('\\', '/'), content))
    
    # Build the files tree
    files_tree = {}
    contents = {}
    
    for rel_path, content in file_list:
        parts = rel_path.split('/')
        current = files_tree
        for i, part in enumerate(parts):
            if i == len(parts) - 1:
                sha256 = hashlib.sha256(content).hexdigest()
                current[part] = {
                    "size": len(content),
                    "offset": "0",
                    "integrity": {
                        "algorithm": "SHA256",
                        "hash": sha256
                    }
                }
            else:
                if part not in current:
                    current[part] = {}
                current = current[part]
        contents[rel_path] = content
    
    # Convert flat tree to ASAR format (wrap directories with 'files' key)
    def to_asar_format(node):
        if is_leaf(node):
            return node
        result = {}
        for k, v in node.items():
            result[k] = to_asar_format(v)
        return {"files": result}
    
    header = to_asar_format(files_tree)
    
    # Calculate offsets
    header_json = json.dumps(header, separators=(',', ':'))
    header_bytes = header_json.encode('utf-8')
    
    json_size = len(header_bytes)
    json_start = 16
    payload_start = align_up(json_start + json_size)
    
    # Calculate file offsets in sorted order (alphabetical)
    sorted_files = sorted(file_list, key=lambda x: x[0])
    current_offset = 0
    
    # Build the offset map
    offset_map = {}
    for rel_path, content in sorted_files:
        offset_map[rel_path] = (current_offset, len(content))
        current_offset += len(content)
    
    # Parse and update JSON with offsets and block hashes
    header_obj = json.loads(header_json)
    
    def set_offsets(node):
        if isinstance(node, dict) and 'files' in node:
            for name, child in sorted(node['files'].items()):
                set_offsets(child)
        elif isinstance(node, dict) and 'size' in node:
            # Leaf node - no action needed here, we'll update from offset_map
            pass
    
    def apply_offsets(node, path=''):
        if isinstance(node, dict) and 'files' in node:
            for name, child in sorted(node['files'].items()):
                apply_offsets(child, f"{path}/{name}" if path else name)
        elif isinstance(node, dict) and 'size' in node:
            if path in offset_map:
                off, sz = offset_map[path]
                node['offset'] = str(off)
                block_hash = hashlib.sha256(contents[path]).hexdigest()
                node['integrity'] = {
                    "algorithm": "SHA256",
                    "hash": block_hash,
                    "blocks": [block_hash]
                }
    
    apply_offsets(header_obj)
    
    # Re-serialize with updated offsets
    header_json = json.dumps(header_obj, separators=(',', ':'))
    header_bytes = header_json.encode('utf-8')
    json_size = len(header_bytes)
    payload_start = align_up(json_start + json_size)
    
    # Write ASAR file
    with open(output_path, 'wb') as f:
        # Write 16-byte base header
        f.write(struct.pack('<I', 4))
        f.write(struct.pack('<I', payload_start - 8))
        f.write(struct.pack('<I', payload_start - 12))
        f.write(struct.pack('<I', json_size))
        
        # Write JSON header
        f.write(header_bytes)
        
        # Pad to 4 bytes
        padding = payload_start - (json_start + json_size)
        if padding > 0:
            f.write(b'\x00' * padding)
        
        # Write file contents
        for rel_path, content in sorted_files:
            f.write(content)
    
    total_size = os.path.getsize(output_path)
    print(f"✓ ASAR créé: {output_path}")
    print(f"  Taille: {total_size} bytes ({total_size/1024/1024:.1f} MB)")
    print(f"  Fichiers: {len(sorted_files)}")

if __name__ == "__main__":
    script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    src = sys.argv[1] if len(sys.argv) > 1 else script_dir
    out = sys.argv[2] if len(sys.argv) > 2 else os.path.join(os.getcwd(), "app.asar")
    
    build_asar(src, out)
