#!/usr/bin/env python3
"""Linux ELF malware static analysis agent using pyelftools and binary inspection."""

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

try:
    from elftools.elf.elffile import ELFFile
    HAS_ELFTOOLS = True
except ImportError:
    HAS_ELFTOOLS = False


def compute_hashes(filepath):
    """Compute MD5, SHA1, and SHA256 hashes of a file."""
    md5 = hashlib.md5()
    sha1 = hashlib.sha1()
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            md5.update(chunk)
            sha1.update(chunk)
            sha256.update(chunk)
    return {"md5": md5.hexdigest(), "sha1": sha1.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 -sum((c / length) * math.log2(c / length) for c in counter.values())


def analyze_elf_header(filepath):
    """Parse ELF header and extract key properties."""
    if not HAS_ELFTOOLS:
        return {"error": "pyelftools not installed: pip install pyelftools"}
    with open(filepath, "rb") as f:
        elf = ELFFile(f)
        symtab = elf.get_section_by_name(".symtab")
        info = {
            "class": f"{elf.elfclass}-bit",
            "endian": "Little" if elf.little_endian else "Big",
            "machine": elf.header.e_machine,
            "type": elf.header.e_type,
            "entry_point": f"0x{elf.header.e_entry:X}",
            "stripped": symtab is None,
            "num_sections": elf.num_sections(),
            "num_segments": elf.num_segments(),
        }
    return info


def analyze_sections(filepath):
    """Analyze ELF sections for entropy and suspicious characteristics."""
    if not HAS_ELFTOOLS:
        return []
    sections = []
    with open(filepath, "rb") as f:
        elf = ELFFile(f)
        for section in elf.iter_sections():
            data = section.data()
            if len(data) == 0:
                continue
            entropy = calculate_entropy(data)
            sections.append({
                "name": section.name,
                "type": section["sh_type"],
                "size": len(data),
                "entropy": round(entropy, 4),
                "high_entropy": entropy > 7.0,
                "flags": section["sh_flags"],
            })
    return sections


def extract_strings(filepath, min_length=6):
    """Extract ASCII strings from the binary and categorize by type."""
    stdout, _, rc = subprocess.run(
        ["strings", "-n", str(min_length), filepath],
        capture_output=True, text=True, timeout=120
    ).stdout, "", 0
    if not stdout:
        return {}
    all_strings = stdout.strip().splitlines()
    categorized = {
        "urls": [], "ips": [], "domains": [], "shell_commands": [],
        "crypto_mining": [], "persistence": [], "ssh_related": [],
        "total": len(all_strings),
    }
    for s in all_strings:
        s_lower = s.lower()
        if any(proto in s_lower for proto in ["http://", "https://", "ftp://"]):
            categorized["urls"].append(s)
        if any(p in s_lower for p in ["stratum", "xmr", "monero", "pool.", "mining"]):
            categorized["crypto_mining"].append(s)
        if any(p in s_lower for p in ["crontab", "systemd", "init.d", "rc.local",
                                       "ld.so.preload", "systemctl"]):
            categorized["persistence"].append(s)
        if any(p in s_lower for p in ["ssh", "authorized_keys", "id_rsa", "shadow", "passwd"]):
            categorized["ssh_related"].append(s)
        if any(p in s_lower for p in ["bash", "wget", "curl", "chmod", "/tmp/", "/dev/"]):
            categorized["shell_commands"].append(s)
        import re
        if re.match(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", s):
            categorized["ips"].append(s)
        if re.match(r"[a-zA-Z0-9.-]+\.(com|net|org|io|ru|cn|xyz)", s):
            categorized["domains"].append(s)
    return categorized


def check_packing(filepath):
    """Check if the binary is packed with UPX or other packers."""
    with open(filepath, "rb") as f:
        data = f.read(4096)
    indicators = []
    if b"UPX!" in data:
        indicators.append("UPX packer detected (UPX! magic)")
    if b"UPX0" in data or b"UPX1" in data:
        indicators.append("UPX section names found")
    stdout, _, _ = subprocess.run(["upx", "-t", filepath],
                                   capture_output=True, text=True,
                                   stderr=subprocess.STDOUT, timeout=120).stdout, "", 0
    if stdout and "packed" in stdout.lower():
        indicators.append("UPX verification confirms packing")
    return indicators


def analyze_dynamic_linking(filepath):
    """Analyze dynamic linking information and imported functions."""
    stdout, _, rc = subprocess.run(["readelf", "-d", filepath],
                                    capture_output=True, text=True, timeout=120).stdout, "", 0
    dynamic_info = {"libraries": [], "rpath": None}
    if stdout:
        for line in stdout.splitlines():
            if "NEEDED" in line:
                lib = line.split("[")[-1].rstrip("]") if "[" in line else ""
                dynamic_info["libraries"].append(lib)
            if "RPATH" in line or "RUNPATH" in line:
                dynamic_info["rpath"] = line.split("[")[-1].rstrip("]")

    readelf_proc = subprocess.run(
        ["readelf", "-r", filepath],
        capture_output=True, text=True,
        timeout=120,
    )
    import re as _re
    suspicious_funcs = _re.compile(r'socket|connect|exec|fork|open|write|bind|listen|send|recv')
    stdout2 = "\n".join(
        line for line in (readelf_proc.stdout or "").splitlines()
        if suspicious_funcs.search(line)
    )
    dynamic_info["suspicious_imports"] = [
        line.strip() for line in (stdout2 or "").splitlines() if line.strip()
    ]
    return dynamic_info


def detect_malware_type(strings_data):
    """Classify malware type based on extracted strings."""
    classifications = []
    if strings_data.get("crypto_mining"):
        classifications.append("Cryptominer")
    if any("flood" in s.lower() or "ddos" in s.lower()
           for s in strings_data.get("shell_commands", [])):
        classifications.append("DDoS Botnet")
    if strings_data.get("ssh_related") and strings_data.get("persistence"):
        classifications.append("Backdoor/Trojan")
    if any("insmod" in s or "modprobe" in s or "init_module" in s
           for s in strings_data.get("shell_commands", [])):
        classifications.append("Rootkit")
    if any("ransom" in s.lower() or "encrypt" in s.lower() or "bitcoin" in s.lower()
           for cat in strings_data.values() if isinstance(cat, list) for s in cat):
        classifications.append("Ransomware")
    return classifications or ["Unknown"]


if __name__ == "__main__":
    print("=" * 60)
    print("Linux ELF Malware Analysis Agent")
    print("Static analysis with pyelftools, strings, readelf")
    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}")
        print(f"[*] Size: {os.path.getsize(target)} bytes")

        hashes = compute_hashes(target)
        print(f"[*] MD5:    {hashes['md5']}")
        print(f"[*] SHA256: {hashes['sha256']}")

        elf_info = analyze_elf_header(target)
        print(f"\n--- ELF Header ---")
        for k, v in elf_info.items():
            print(f"  {k}: {v}")

        packing = check_packing(target)
        if packing:
            for p in packing:
                print(f"[!] {p}")

        sections = analyze_sections(target)
        high_ent = [s for s in sections if s.get("high_entropy")]
        if high_ent:
            print(f"\n[!] High entropy sections (possible packing/encryption):")
            for s in high_ent:
                print(f"  {s['name']}: entropy={s['entropy']}, size={s['size']}")

        strings_data = extract_strings(target)
        print(f"\n--- Strings Analysis ({strings_data.get('total', 0)} total) ---")
        for category in ["urls", "ips", "domains", "crypto_mining", "persistence", "ssh_related"]:
            items = strings_data.get(category, [])
            if items:
                print(f"  {category}: {len(items)}")
                for item in items[:5]:
                    print(f"    - {item}")

        classification = detect_malware_type(strings_data)
        print(f"\n[*] Classification: {', '.join(classification)}")
    else:
        print(f"\n[DEMO] Usage: python agent.py <elf_binary>")
        print("[*] Provide a Linux ELF binary for analysis.")