#!/usr/bin/env python3
"""Packed malware analysis agent for UPX and generic packer detection and unpacking."""

import subprocess
import os
import sys
import hashlib
import math
from collections import Counter

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


def compute_hashes(filepath):
    """Compute file hashes."""
    md5 = hashlib.md5()
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            md5.update(chunk)
            sha256.update(chunk)
    return {"md5": md5.hexdigest(), "sha256": sha256.hexdigest()}


def calculate_entropy(data):
    """Calculate Shannon entropy of binary data."""
    if not data:
        return 0.0
    counter = Counter(data)
    length = len(data)
    return round(-sum((c / length) * math.log2(c / length) for c in counter.values()), 4)


def detect_upx(filepath):
    """Check for UPX packing signatures in the binary."""
    indicators = []
    with open(filepath, "rb") as f:
        data = f.read()

    if b"UPX!" in data:
        indicators.append("UPX! magic string found in binary")
    if b"UPX0" in data:
        indicators.append("UPX0 section name found")
    if b"UPX1" in data:
        indicators.append("UPX1 section name found")
    if b"UPX2" in data:
        indicators.append("UPX2 section name found")

    # Check for corrupted/modified UPX headers
    upx_pos = data.find(b"UPX!")
    if upx_pos != -1:
        # UPX version info follows the magic
        if upx_pos + 24 <= len(data):
            version_byte = data[upx_pos + 4]
            indicators.append(f"UPX version byte: 0x{version_byte:02X}")
    return indicators


def detect_generic_packing(filepath):
    """Detect generic packing indicators using PE section analysis."""
    if not HAS_PEFILE:
        return {"error": "pefile not installed: pip install pefile"}
    try:
        pe = pefile.PE(filepath)
    except pefile.PEFormatError:
        return {"error": "Not a valid PE file"}

    indicators = []
    sections = []
    high_entropy_count = 0

    for section in pe.sections:
        name = section.Name.rstrip(b"\x00").decode("utf-8", errors="replace")
        entropy = section.get_entropy()
        raw_size = section.SizeOfRawData
        virtual_size = section.Misc_VirtualSize
        sections.append({
            "name": name,
            "entropy": round(entropy, 4),
            "raw_size": raw_size,
            "virtual_size": virtual_size,
            "ratio": round(virtual_size / raw_size, 2) if raw_size > 0 else 0,
        })
        if entropy > 7.0:
            high_entropy_count += 1
            indicators.append(f"High entropy section: {name} ({entropy:.2f})")
        if virtual_size > raw_size * 5 and raw_size > 0:
            indicators.append(f"Suspicious size ratio in {name}: virtual/raw = {virtual_size/raw_size:.1f}x")

    imports = []
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode("utf-8", errors="replace")
            func_count = len(entry.imports)
            imports.append({"dll": dll_name, "functions": func_count})

    total_imports = sum(i["functions"] for i in imports)
    if total_imports < 10:
        indicators.append(f"Very few imports ({total_imports}) - typical of packed binaries")

    # Check for LoadLibrary/GetProcAddress (runtime import resolution)
    import_names = []
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            for imp in entry.imports:
                if imp.name:
                    import_names.append(imp.name.decode("utf-8", errors="replace"))
    if "LoadLibraryA" in import_names and "GetProcAddress" in import_names:
        indicators.append("LoadLibraryA + GetProcAddress present (runtime import resolution)")

    pe.close()
    return {
        "sections": sections,
        "imports": imports,
        "total_imports": total_imports,
        "high_entropy_sections": high_entropy_count,
        "indicators": indicators,
        "likely_packed": high_entropy_count > 0 or total_imports < 10,
    }


def unpack_upx(filepath, output_path=None):
    """Attempt to unpack a UPX-packed binary."""
    if output_path is None:
        output_path = filepath + ".unpacked"
    # First try standard UPX decompression
    cmd = ["upx", "-d", "-o", output_path, filepath]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode == 0:
        return True, "Standard UPX unpack succeeded", output_path

    # If standard fails, try fixing UPX headers
    return False, result.stderr.strip(), None


def fix_upx_headers(filepath, output_path):
    """Attempt to fix corrupted UPX magic bytes for unpacking."""
    with open(filepath, "rb") as f:
        data = bytearray(f.read())

    # Look for known UPX section names that might be renamed
    modified = False
    # Common modifications: UPX0/UPX1 renamed to something else
    for i in range(len(data) - 3):
        # Look for section header pattern near typical PE section table location
        if data[i:i+3] in [b"UP0", b"UP1", b"UX0", b"UX1"]:
            # Might be modified UPX section name
            pass

    # Fix UPX! magic if corrupted
    for i in range(len(data) - 4):
        if data[i:i+2] == b"UX" and data[i+2:i+4] == b"!\x00":
            data[i:i+3] = b"UPX"
            modified = True

    if modified:
        with open(output_path, "wb") as f:
            f.write(data)
        return True
    return False


def compare_packed_unpacked(packed_path, unpacked_path):
    """Compare packed vs unpacked binary properties."""
    if not HAS_PEFILE:
        return {}
    comparison = {}
    for label, path in [("packed", packed_path), ("unpacked", unpacked_path)]:
        try:
            pe = pefile.PE(path)
            imports = 0
            if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
                for entry in pe.DIRECTORY_ENTRY_IMPORT:
                    imports += len(entry.imports)
            sections = len(pe.sections)
            pe.close()
            comparison[label] = {
                "size": os.path.getsize(path),
                "sections": sections,
                "imports": imports,
                "sha256": compute_hashes(path)["sha256"],
            }
        except Exception as e:
            comparison[label] = {"error": str(e)}
    return comparison


if __name__ == "__main__":
    print("=" * 60)
    print("Packed Malware Analysis Agent")
    print("UPX detection, packer identification, automated unpacking")
    print("=" * 60)

    target = sys.argv[1] if len(sys.argv) > 1 else None

    if target and os.path.exists(target):
        print(f"\n[*] Analyzing: {target}")
        hashes = compute_hashes(target)
        print(f"[*] SHA-256: {hashes['sha256']}")
        print(f"[*] Size: {os.path.getsize(target)} bytes")

        print("\n--- UPX Signature Check ---")
        upx_indicators = detect_upx(target)
        for ind in upx_indicators:
            print(f"  [!] {ind}")

        print("\n--- Generic Packing Analysis ---")
        packing = detect_generic_packing(target)
        if "error" not in packing:
            print(f"  Likely packed: {packing['likely_packed']}")
            print(f"  Total imports: {packing['total_imports']}")
            print(f"  High entropy sections: {packing['high_entropy_sections']}")
            for ind in packing.get("indicators", []):
                print(f"  [!] {ind}")
            print("\n  Sections:")
            for s in packing.get("sections", []):
                flag = " [HIGH]" if s["entropy"] > 7.0 else ""
                print(f"    {s['name']:10s} entropy={s['entropy']:.2f} "
                      f"raw={s['raw_size']} virt={s['virtual_size']}{flag}")

        if upx_indicators:
            print("\n--- UPX Unpacking ---")
            success, msg, output = unpack_upx(target)
            if success:
                print(f"  [OK] {msg}")
                print(f"  [*] Unpacked file: {output}")
                print("\n--- Comparison ---")
                comp = compare_packed_unpacked(target, output)
                for label, data in comp.items():
                    if "error" not in data:
                        print(f"  {label}: size={data['size']}, "
                              f"sections={data['sections']}, imports={data['imports']}")
            else:
                print(f"  [FAIL] {msg}")
                print("  [*] Try fixing UPX headers or use dynamic unpacking with a debugger")
    else:
        print(f"\n[DEMO] Usage: python agent.py <packed_binary.exe>")