#!/usr/bin/env python3
"""
Static analysis tool for PE and ELF binaries.

Parses headers, imports/exports, sections with entropy, packer detection,
resource extraction, compiler identification, and signature checks.
Outputs structured JSON.

Usage:
    python3 static_analyzer.py --file <binary> [--output report.json]
    python3 static_analyzer.py --file <binary> --headers
    python3 static_analyzer.py --file <binary> --imports
    python3 static_analyzer.py --file <binary> --entropy
    python3 static_analyzer.py --file <binary> --packer-detect
    python3 static_analyzer.py --file <binary> --resources
    python3 static_analyzer.py --file <binary> --signature
    python3 static_analyzer.py --file <binary> --compiler
"""

from __future__ import annotations

import argparse
import json
import math
import os
import struct
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path

# ---------------------------------------------------------------------------
# Optional imports
# ---------------------------------------------------------------------------

try:
    import pefile
    HAS_PEFILE = True
except ImportError:
    HAS_PEFILE = False

try:
    from elftools.elf.elffile import ELFFile
    from elftools.elf.sections import SymbolTableSection
    from elftools.elf.dynamic import DynamicSection
    HAS_ELFTOOLS = True
except ImportError:
    HAS_ELFTOOLS = False

# ---------------------------------------------------------------------------
# Entropy calculation
# ---------------------------------------------------------------------------

def calculate_entropy(data: bytes) -> float:
    """Calculate Shannon entropy of a byte sequence (0.0 - 8.0 scale)."""
    if not data:
        return 0.0
    counts = Counter(data)
    length = len(data)
    entropy = 0.0
    for count in counts.values():
        if count == 0:
            continue
        p = count / length
        entropy -= p * math.log2(p)
    return round(entropy, 4)


def entropy_classification(entropy: float) -> str:
    """Classify entropy value."""
    if entropy < 1.0:
        return "sparse/empty"
    elif entropy < 3.5:
        return "plain text/resources"
    elif entropy < 5.0:
        return "code with data"
    elif entropy < 6.5:
        return "compiled code (normal)"
    elif entropy < 7.5:
        return "compressed/dense"
    else:
        return "encrypted/packed"


# ---------------------------------------------------------------------------
# PE analysis
# ---------------------------------------------------------------------------

PACKER_SECTION_NAMES = {
    ".upx0": "UPX", ".upx1": "UPX", ".upx2": "UPX",
    "upx0": "UPX", "upx1": "UPX",
    ".aspack": "ASPack", ".adata": "ASPack",
    ".themida": "Themida",
    ".vmp0": "VMProtect", ".vmp1": "VMProtect", ".vmp2": "VMProtect",
    ".enigma1": "Enigma", ".enigma2": "Enigma",
    ".mpress1": "MPRESS", ".mpress2": "MPRESS",
    ".nsp0": "NsPack", ".nsp1": "NsPack",
    ".petite": "Petite",
    ".yp": "Y0da Protector",
    ".perplex": "Perplex",
    ".seau": "SeauSFX",
    ".ccg": "CCG Packer",
}

SUSPICIOUS_IMPORTS = {
    "process_injection": [
        "CreateRemoteThread", "NtCreateThreadEx", "RtlCreateUserThread",
        "VirtualAllocEx", "NtAllocateVirtualMemory", "WriteProcessMemory",
        "NtWriteVirtualMemory", "NtUnmapViewOfSection", "QueueUserAPC",
        "NtQueueApcThread", "SetThreadContext", "NtSetContextThread",
    ],
    "process_hollowing": [
        "NtUnmapViewOfSection", "ZwUnmapViewOfSection",
    ],
    "code_injection": [
        "SetWindowsHookEx", "SetWindowsHookExA", "SetWindowsHookExW",
    ],
    "dll_injection": [
        "LoadLibraryA", "LoadLibraryW", "LoadLibraryExA", "LoadLibraryExW",
        "LdrLoadDll",
    ],
    "memory_manipulation": [
        "VirtualAlloc", "VirtualAllocEx", "VirtualProtect", "VirtualProtectEx",
        "HeapCreate",
    ],
    "network": [
        "InternetOpenA", "InternetOpenW", "InternetOpenUrlA", "InternetOpenUrlW",
        "InternetConnectA", "InternetConnectW", "HttpOpenRequestA", "HttpSendRequestA",
        "URLDownloadToFileA", "URLDownloadToFileW", "URLDownloadToCacheFileA",
        "WinHttpOpen", "WinHttpConnect", "WinHttpSendRequest",
        "WSAStartup", "connect", "send", "recv", "WSASend", "WSARecv",
        "getaddrinfo", "gethostbyname",
    ],
    "file_operations": [
        "CreateFileA", "CreateFileW", "WriteFile", "ReadFile",
        "DeleteFileA", "DeleteFileW", "MoveFileA", "MoveFileW",
        "CopyFileA", "CopyFileW", "CreateDirectoryA",
    ],
    "registry": [
        "RegOpenKeyExA", "RegOpenKeyExW", "RegSetValueExA", "RegSetValueExW",
        "RegCreateKeyExA", "RegCreateKeyExW", "RegDeleteKeyA", "RegDeleteValueA",
    ],
    "crypto": [
        "CryptEncrypt", "CryptDecrypt", "CryptGenKey", "CryptImportKey",
        "CryptAcquireContextA", "CryptCreateHash", "CryptHashData",
        "BCryptEncrypt", "BCryptDecrypt", "BCryptGenerateSymmetricKey",
    ],
    "anti_debug": [
        "IsDebuggerPresent", "CheckRemoteDebuggerPresent",
        "NtQueryInformationProcess", "NtSetInformationThread",
        "OutputDebugStringA", "GetTickCount", "QueryPerformanceCounter",
    ],
    "execution": [
        "CreateProcessA", "CreateProcessW", "CreateProcessInternalW",
        "WinExec", "ShellExecuteA", "ShellExecuteW", "ShellExecuteExA",
        "system",
    ],
    "privilege_escalation": [
        "AdjustTokenPrivileges", "OpenProcessToken", "LookupPrivilegeValueA",
        "ImpersonateLoggedOnUser", "DuplicateTokenEx",
    ],
    "service_manipulation": [
        "CreateServiceA", "CreateServiceW", "StartServiceA",
        "OpenSCManagerA", "ChangeServiceConfigA",
    ],
    "screenshot_keylog": [
        "GetAsyncKeyState", "GetKeyState", "SetWindowsHookExA",
        "GetDC", "BitBlt", "GetForegroundWindow",
    ],
}


def analyze_pe(file_path: str) -> dict:
    """Perform comprehensive PE analysis."""
    if not HAS_PEFILE:
        return {"error": "pefile library not installed. Install with: pip install pefile"}

    try:
        pe = pefile.PE(file_path)
    except pefile.PEFormatError as e:
        return {"error": f"Invalid PE file: {e}"}

    result = {"format": "PE"}

    # --- Headers ---
    machine_map = {0x14c: "i386", 0x8664: "AMD64", 0x1c0: "ARM", 0xaa64: "ARM64"}
    subsystem_map = {
        1: "Native", 2: "Windows GUI", 3: "Windows CUI",
        7: "POSIX CUI", 9: "Windows CE", 14: "EFI Application",
    }

    timestamp = pe.FILE_HEADER.TimeDateStamp
    try:
        ts_str = datetime.fromtimestamp(timestamp, tz=timezone.utc).isoformat()
    except (OSError, ValueError, OverflowError):
        ts_str = f"invalid ({timestamp})"

    result["headers"] = {
        "machine": machine_map.get(pe.FILE_HEADER.Machine, hex(pe.FILE_HEADER.Machine)),
        "timestamp": ts_str,
        "timestamp_raw": timestamp,
        "characteristics": hex(pe.FILE_HEADER.Characteristics),
        "is_dll": bool(pe.FILE_HEADER.Characteristics & 0x2000),
        "is_exe": bool(pe.FILE_HEADER.Characteristics & 0x0002),
        "subsystem": subsystem_map.get(pe.OPTIONAL_HEADER.Subsystem, str(pe.OPTIONAL_HEADER.Subsystem)),
        "entry_point": hex(pe.OPTIONAL_HEADER.AddressOfEntryPoint),
        "image_base": hex(pe.OPTIONAL_HEADER.ImageBase),
        "num_sections": pe.FILE_HEADER.NumberOfSections,
        "size_of_image": pe.OPTIONAL_HEADER.SizeOfImage,
        "pe_type": "PE32+" if pe.OPTIONAL_HEADER.Magic == 0x20B else "PE32",
    }

    # Check for timestamp anomalies
    anomalies = []
    now_ts = int(datetime.now(tz=timezone.utc).timestamp())
    if timestamp == 0:
        anomalies.append("Compilation timestamp is zero (likely wiped)")
    elif timestamp > now_ts:
        anomalies.append("Compilation timestamp is in the future")
    elif timestamp < 946684800:  # Before 2000
        anomalies.append("Compilation timestamp before year 2000 (possibly fake)")

    # --- Sections ---
    sections = []
    with open(file_path, "rb") as f:
        file_data = f.read()

    for section in pe.sections:
        name = section.Name.decode("utf-8", errors="replace").rstrip("\x00")
        raw_data = section.get_data()
        entropy = calculate_entropy(raw_data)

        flags = []
        if section.Characteristics & 0x20000000:
            flags.append("x")  # executable
        if section.Characteristics & 0x40000000:
            flags.append("r")  # readable
        if section.Characteristics & 0x80000000:
            flags.append("w")  # writable

        sec_info = {
            "name": name,
            "virtual_address": hex(section.VirtualAddress),
            "virtual_size": section.Misc_VirtualSize,
            "raw_size": section.SizeOfRawData,
            "entropy": entropy,
            "entropy_class": entropy_classification(entropy),
            "flags": "".join(flags),
            "characteristics": hex(section.Characteristics),
        }

        # Check for anomalies
        if entropy > 7.0 and "x" in flags:
            anomalies.append(f"Section {name} is executable with high entropy ({entropy}) - likely packed")
        if "w" in flags and "x" in flags:
            anomalies.append(f"Section {name} is both writable and executable (WX)")
        if section.SizeOfRawData == 0 and section.Misc_VirtualSize > 0:
            anomalies.append(f"Section {name} has zero raw size but non-zero virtual size (unpacking target)")

        # Check for packer section names
        name_lower = name.lower().strip()
        if name_lower in PACKER_SECTION_NAMES:
            anomalies.append(f"Packer section detected: {name} ({PACKER_SECTION_NAMES[name_lower]})")

        sections.append(sec_info)

    result["sections"] = sections

    # --- Overall entropy ---
    result["entropy_overall"] = calculate_entropy(file_data)

    # --- Imports ---
    imports = {}
    suspicious_found = {}
    total_imports = 0

    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode("utf-8", errors="replace")
            funcs = []
            for imp in entry.imports:
                func_name = imp.name.decode("utf-8", errors="replace") if imp.name else f"ord({imp.ordinal})"
                funcs.append(func_name)
                total_imports += 1

                # Check against suspicious imports
                for category, api_list in SUSPICIOUS_IMPORTS.items():
                    if func_name in api_list:
                        suspicious_found.setdefault(category, []).append(f"{dll_name}:{func_name}")

            imports[dll_name] = funcs

    result["imports"] = imports
    result["total_import_count"] = total_imports
    result["suspicious_imports"] = suspicious_found

    if total_imports < 10 and total_imports > 0:
        anomalies.append(f"Very few imports ({total_imports}) - binary may be packed or use dynamic resolution")

    # --- Exports ---
    exports = []
    if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
        for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            name = exp.name.decode("utf-8", errors="replace") if exp.name else f"ord({exp.ordinal})"
            exports.append({
                "name": name,
                "ordinal": exp.ordinal,
                "address": hex(exp.address) if exp.address else None,
            })
    result["exports"] = exports

    # --- Import hash ---
    try:
        result["imphash"] = pe.get_imphash()
    except Exception:
        result["imphash"] = None

    # --- Resources ---
    resources = []
    if hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"):
        for res_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
            type_name = pefile.RESOURCE_TYPE.get(res_type.id, str(res_type.id))
            if hasattr(res_type, "directory"):
                for res_id in res_type.directory.entries:
                    if hasattr(res_id, "directory"):
                        for res_lang in res_id.directory.entries:
                            try:
                                data = pe.get_data(
                                    res_lang.data.struct.OffsetToData,
                                    res_lang.data.struct.Size,
                                )
                                ent = calculate_entropy(data)
                                magic_preview = data[:8].hex() if data else ""
                                resources.append({
                                    "type": type_name,
                                    "id": str(res_id.id) if res_id.id else (res_id.name.string.decode("utf-8", errors="replace") if res_id.name else "?"),
                                    "size": res_lang.data.struct.Size,
                                    "entropy": ent,
                                    "magic_bytes": magic_preview,
                                })
                                if ent > 7.0 and res_lang.data.struct.Size > 1024:
                                    anomalies.append(f"Resource {type_name}/{res_id.id} has high entropy ({ent}) - possible encrypted payload")
                                if magic_preview.startswith("4d5a"):
                                    anomalies.append(f"Resource {type_name}/{res_id.id} contains embedded PE (MZ header)")
                            except Exception:
                                pass
    result["resources"] = resources

    # --- Packer detection ---
    packer_indicators = []
    # Check sections
    for sec in sections:
        name_lower = sec["name"].lower().strip()
        if name_lower in PACKER_SECTION_NAMES:
            packer_indicators.append(f"Section name: {sec['name']} ({PACKER_SECTION_NAMES[name_lower]})")

    # Check high entropy
    high_entropy_sections = [s for s in sections if s["entropy"] > 7.0]
    if len(high_entropy_sections) == len(sections) and len(sections) > 1:
        packer_indicators.append("All sections have high entropy")
    elif len(high_entropy_sections) > 0:
        for s in high_entropy_sections:
            packer_indicators.append(f"High entropy section: {s['name']} ({s['entropy']})")

    # Check import count
    if 0 < total_imports < 10:
        packer_indicators.append(f"Very few imports ({total_imports})")

    # Check for known packer strings in overlay
    overlay_offset = pe.get_overlay_data_start_offset()
    if overlay_offset:
        overlay_data = file_data[overlay_offset:overlay_offset + 256]
        overlay_str = overlay_data.decode("ascii", errors="ignore")
        for packer_name in ["UPX!", "MPRESS", "ASPack", "PECompact"]:
            if packer_name in overlay_str:
                packer_indicators.append(f"Packer signature in overlay: {packer_name}")

    result["packer_detection"] = {
        "likely_packed": len(packer_indicators) >= 2,
        "indicators": packer_indicators,
    }

    # --- Digital signature ---
    has_security_dir = False
    if pe.OPTIONAL_HEADER.DATA_DIRECTORY[4].Size > 0:
        has_security_dir = True
    result["signature"] = {
        "has_authenticode": has_security_dir,
        "note": "Use osslsigncode or sigcheck for full verification" if has_security_dir else "No Authenticode signature present",
    }

    # --- Compiler detection ---
    compiler_hints = []

    # Check Rich header
    if hasattr(pe, "RICH_HEADER") and pe.RICH_HEADER:
        compiler_hints.append("MSVC (Rich header present)")

    # Check for .NET
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_lower = entry.dll.decode("utf-8", errors="replace").lower()
            if dll_lower == "mscoree.dll":
                compiler_hints.append(".NET (mscoree.dll import)")
                break

    # Check for common patterns in strings
    check_region = file_data[:min(len(file_data), 1024 * 1024)]
    check_str = check_region.decode("ascii", errors="ignore")
    if "AutoIt" in check_str:
        compiler_hints.append("AutoIt compiled script")
    if "PyInstaller" in check_str or "pyi-" in check_str:
        compiler_hints.append("PyInstaller (Python)")
    if "Py_InitModule" in check_str or "Py_SetProgramName" in check_str:
        compiler_hints.append("Python (embedded interpreter)")
    if "Go build" in check_str or "runtime.main" in check_str:
        compiler_hints.append("Go (Golang)")
    if "Borland" in check_str or "Embarcadero" in check_str:
        compiler_hints.append("Delphi/C++ Builder")
    if "Nim" in check_str and "nimbase" in check_str.lower():
        compiler_hints.append("Nim")

    for section in pe.sections:
        name = section.Name.decode("utf-8", errors="replace").rstrip("\x00")
        if name == "CODE" or name == ".idata":
            compiler_hints.append("Delphi/Borland (CODE/.idata section)")
            break

    result["compiler"] = {
        "hints": compiler_hints if compiler_hints else ["Unknown / could not determine"],
    }

    result["anomalies"] = anomalies

    pe.close()
    return result


# ---------------------------------------------------------------------------
# ELF analysis
# ---------------------------------------------------------------------------

def analyze_elf(file_path: str) -> dict:
    """Perform comprehensive ELF analysis."""
    if not HAS_ELFTOOLS:
        return {"error": "pyelftools library not installed. Install with: pip install pyelftools"}

    result = {"format": "ELF"}
    anomalies = []

    with open(file_path, "rb") as f:
        file_data = f.read()

    with open(file_path, "rb") as f:
        try:
            elf = ELFFile(f)
        except Exception as e:
            return {"error": f"Invalid ELF file: {e}"}

        # --- Headers ---
        type_map = {"ET_EXEC": "executable", "ET_DYN": "shared object", "ET_REL": "relocatable", "ET_CORE": "core"}
        result["headers"] = {
            "class": elf.elfclass,
            "data_encoding": elf.little_endian and "little-endian" or "big-endian",
            "os_abi": elf["e_ident"]["EI_OSABI"],
            "type": type_map.get(elf["e_type"], elf["e_type"]),
            "machine": elf["e_machine"],
            "entry_point": hex(elf["e_entry"]),
            "num_sections": elf.num_sections(),
            "num_segments": elf.num_segments(),
        }

        # --- Sections ---
        sections = []
        for section in elf.iter_sections():
            data = section.data()
            entropy = calculate_entropy(data)

            flags = ""
            sh_flags = section["sh_flags"]
            if sh_flags & 0x1:
                flags += "w"
            if sh_flags & 0x2:
                flags += "a"  # alloc
            if sh_flags & 0x4:
                flags += "x"

            sec_info = {
                "name": section.name,
                "type": section["sh_type"],
                "size": section["sh_size"],
                "entropy": entropy,
                "entropy_class": entropy_classification(entropy),
                "flags": flags,
                "address": hex(section["sh_addr"]),
            }
            sections.append(sec_info)

            if entropy > 7.0 and "x" in flags and section["sh_size"] > 256:
                anomalies.append(f"Section {section.name} is executable with high entropy ({entropy})")
            if "w" in flags and "x" in flags and section["sh_size"] > 0:
                anomalies.append(f"Section {section.name} is writable and executable (WX)")

        result["sections"] = sections
        result["entropy_overall"] = calculate_entropy(file_data)

        # --- Dynamic symbols (imports/exports) ---
        imports = []
        exports = []

        for section in elf.iter_sections():
            if isinstance(section, SymbolTableSection):
                for symbol in section.iter_symbols():
                    if symbol.name:
                        sym_info = {
                            "name": symbol.name,
                            "type": symbol["st_info"]["type"],
                            "bind": symbol["st_info"]["bind"],
                            "section_index": symbol["st_shndx"],
                        }
                        if symbol["st_shndx"] == "SHN_UNDEF":
                            imports.append(sym_info)
                        elif symbol["st_info"]["bind"] == "STB_GLOBAL":
                            exports.append(sym_info)

        result["imports"] = imports[:200]  # Limit output size
        result["exports"] = exports[:200]
        result["total_import_count"] = len(imports)
        result["total_export_count"] = len(exports)

        # --- Dynamic section (shared libraries) ---
        needed_libs = []
        rpath = None
        for section in elf.iter_sections():
            if isinstance(section, DynamicSection):
                for tag in section.iter_tags():
                    if tag.entry.d_tag == "DT_NEEDED":
                        needed_libs.append(tag.needed)
                    elif tag.entry.d_tag == "DT_RPATH":
                        rpath = tag.rpath
                    elif tag.entry.d_tag == "DT_RUNPATH":
                        rpath = tag.runpath

        result["needed_libraries"] = needed_libs
        if rpath:
            result["rpath"] = rpath
            anomalies.append(f"Binary has custom RPATH: {rpath}")

        # --- Packer detection ---
        packer_indicators = []
        if not needed_libs and result["headers"]["type"] == "executable":
            packer_indicators.append("Statically linked executable (no shared libraries)")

        stripped = True
        for section in elf.iter_sections():
            if section.name == ".symtab":
                stripped = False
                break
        if stripped:
            packer_indicators.append("Symbol table stripped")

        check_str = file_data[:65536].decode("ascii", errors="ignore")
        if "UPX!" in check_str:
            packer_indicators.append("UPX packer signature found")

        high_ent = [s for s in sections if s["entropy"] > 7.0 and s["size"] > 256]
        if high_ent:
            for s in high_ent:
                packer_indicators.append(f"High entropy section: {s['name']} ({s['entropy']})")

        result["packer_detection"] = {
            "likely_packed": len(packer_indicators) >= 2,
            "indicators": packer_indicators,
        }

        # --- Compiler detection ---
        compiler_hints = []
        comment_section = None
        for section in elf.iter_sections():
            if section.name == ".comment":
                comment_section = section.data().decode("ascii", errors="ignore")
                break

        if comment_section:
            if "GCC" in comment_section:
                compiler_hints.append(f"GCC ({comment_section.strip()})")
            elif "clang" in comment_section.lower():
                compiler_hints.append(f"Clang ({comment_section.strip()})")

        check_str_full = file_data.decode("ascii", errors="ignore")
        if "Go build" in check_str_full or "runtime.main" in check_str_full:
            compiler_hints.append("Go (Golang)")
        if "rust_begin_unwind" in check_str_full or "rust_panic" in check_str_full:
            compiler_hints.append("Rust")

        result["compiler"] = {
            "hints": compiler_hints if compiler_hints else ["Unknown / could not determine"],
        }

        result["anomalies"] = anomalies

    return result


# ---------------------------------------------------------------------------
# Format detection and dispatch
# ---------------------------------------------------------------------------

def detect_format(file_path: str) -> str:
    """Detect binary format from magic bytes."""
    with open(file_path, "rb") as f:
        magic = f.read(4)
    if magic[:2] == b"MZ":
        return "PE"
    elif magic == b"\x7fELF":
        return "ELF"
    elif magic in (b"\xfe\xed\xfa\xce", b"\xfe\xed\xfa\xcf",
                   b"\xce\xfa\xed\xfe", b"\xcf\xfa\xed\xfe",
                   b"\xca\xfe\xba\xbe"):
        return "MachO"
    return "unknown"


def analyze(file_path: str) -> dict:
    """Run analysis appropriate for the detected format."""
    fmt = detect_format(file_path)
    if fmt == "PE":
        return analyze_pe(file_path)
    elif fmt == "ELF":
        return analyze_elf(file_path)
    elif fmt == "MachO":
        return {"format": "Mach-O", "note": "Mach-O analysis requires macOS tools or lief library"}
    else:
        # Try PE then ELF as fallback
        result = analyze_pe(file_path)
        if "error" in result:
            result = analyze_elf(file_path)
        if "error" in result:
            return {"error": f"Unsupported or unrecognized binary format", "detected": fmt}
        return result


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Static analysis of PE and ELF binaries",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Examples:\n"
               "  %(prog)s --file sample.exe\n"
               "  %(prog)s --file sample.exe --imports --entropy\n"
               "  %(prog)s --file sample.elf --output report.json\n",
    )
    parser.add_argument("--input", "--file", "-f", dest="file", required=True, help="Path to binary file")
    parser.add_argument("--output", "-o", help="Write JSON report to file")
    parser.add_argument("--headers", action="store_true", help="Show header information only")
    parser.add_argument("--imports", action="store_true", help="Show imports only")
    parser.add_argument("--entropy", action="store_true", help="Show section entropy only")
    parser.add_argument("--packer-detect", action="store_true", help="Show packer detection only")
    parser.add_argument("--resources", action="store_true", help="Show resources only")
    parser.add_argument("--signature", action="store_true", help="Show signature info only")
    parser.add_argument("--compiler", action="store_true", help="Show compiler info only")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")

    args = parser.parse_args()

    if not os.path.isfile(args.file):
        print(f"Error: File not found: {args.file}", file=sys.stderr)
        sys.exit(1)

    report = analyze(args.file)

    # Filter output if specific section requested
    specific_sections = {
        "headers": args.headers,
        "imports": args.imports,
        "entropy": args.entropy,
        "packer_detect": args.packer_detect,
        "resources": args.resources,
        "signature": args.signature,
        "compiler": args.compiler,
    }

    if any(specific_sections.values()):
        filtered = {"file": args.file, "format": report.get("format", "unknown")}
        if args.headers:
            filtered["headers"] = report.get("headers", {})
        if args.imports:
            filtered["imports"] = report.get("imports", {})
            filtered["suspicious_imports"] = report.get("suspicious_imports", {})
            filtered["total_import_count"] = report.get("total_import_count", 0)
        if args.entropy:
            filtered["sections"] = [
                {"name": s["name"], "entropy": s["entropy"], "entropy_class": s["entropy_class"], "size": s.get("raw_size", s.get("size", 0))}
                for s in report.get("sections", [])
            ]
            filtered["entropy_overall"] = report.get("entropy_overall", 0)
        if args.packer_detect:
            filtered["packer_detection"] = report.get("packer_detection", {})
        if args.resources:
            filtered["resources"] = report.get("resources", [])
        if args.signature:
            filtered["signature"] = report.get("signature", {})
        if args.compiler:
            filtered["compiler"] = report.get("compiler", {})
        report = filtered

    result = json.dumps(report, indent=2, default=str)

    if args.output:
        Path(args.output).write_text(result)
        print(f"Report written to {args.output}")
    else:
        print(result)


if __name__ == "__main__":
    main()