#!/usr/bin/env python3
""".NET malware reverse engineering agent using subprocess wrappers for dnSpy/de4dot."""

import subprocess
import os
import sys
import re
import hashlib
import struct


def compute_hashes(filepath):
    """Compute hashes for sample identification."""
    with open(filepath, "rb") as f:
        data = f.read()
    return {
        "md5": hashlib.md5(data).hexdigest(),
        "sha256": hashlib.sha256(data).hexdigest(),
        "size": len(data),
    }


def detect_dotnet_assembly(filepath):
    """Check if file is a .NET assembly by looking for CLI header."""
    with open(filepath, "rb") as f:
        data = f.read(512)
    if data[:2] != b"MZ":
        return {"is_dotnet": False, "reason": "Not a PE file"}
    try:
        pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
        if pe_offset + 4 > len(data):
            return {"is_dotnet": False, "reason": "Invalid PE header"}
        if data[pe_offset:pe_offset + 4] != b"PE\x00\x00":
            return {"is_dotnet": False, "reason": "Invalid PE signature"}
        opt_offset = pe_offset + 24
        magic = struct.unpack_from("<H", data, opt_offset)[0]
        if magic == 0x10B:
            clr_offset = opt_offset + 208
        elif magic == 0x20B:
            clr_offset = opt_offset + 224
        else:
            return {"is_dotnet": False, "reason": "Unknown PE format"}
        if clr_offset + 8 <= len(data):
            clr_rva, clr_size = struct.unpack_from("<II", data, clr_offset)
            if clr_rva > 0 and clr_size > 0:
                return {"is_dotnet": True, "clr_header_rva": clr_rva, "clr_size": clr_size}
        return {"is_dotnet": False, "reason": "No CLR header"}
    except (struct.error, IndexError):
        return {"is_dotnet": False, "reason": "Parse error"}


def detect_obfuscator(filepath):
    """Detect .NET obfuscator using Detect It Easy."""
    try:
        result = subprocess.run(
            ["diec", filepath], capture_output=True, text=True, timeout=30
        )
        output = result.stdout
        obfuscators = {
            "ConfuserEx": "confuser" in output.lower(),
            "SmartAssembly": "smartassembly" in output.lower(),
            ".NET Reactor": "reactor" in output.lower(),
            "Dotfuscator": "dotfuscator" in output.lower(),
            "Babel": "babel" in output.lower(),
            "Eazfuscator": "eazfuscator" in output.lower(),
            "Crypto Obfuscator": "crypto" in output.lower() and "obfuscator" in output.lower(),
        }
        detected = [name for name, found in obfuscators.items() if found]
        return {"detected": detected, "raw_output": output.strip()}
    except FileNotFoundError:
        return {"detected": [], "raw_output": "diec not installed"}


def deobfuscate_with_de4dot(filepath, output_path):
    """Run de4dot to deobfuscate .NET assembly."""
    try:
        result = subprocess.run(
            ["de4dot", filepath, "-o", output_path],
            capture_output=True, text=True, timeout=120
        )
        return {
            "success": result.returncode == 0,
            "output_path": output_path,
            "stdout": result.stdout[-500:] if result.stdout else "",
        }
    except FileNotFoundError:
        return {"success": False, "error": "de4dot not installed"}


def extract_strings(filepath, min_length=8):
    """Extract strings and classify for IOCs."""
    with open(filepath, "rb") as f:
        data = f.read()
    unicode_strings = re.findall(
        rb"(?:[\x20-\x7e]\x00){%d,}" % min_length, data
    )
    ascii_strings = re.findall(
        rb"[\x20-\x7e]{%d,}" % min_length, data
    )
    all_strings = set()
    for s in ascii_strings:
        all_strings.add(s.decode("ascii", errors="ignore"))
    for s in unicode_strings:
        all_strings.add(s.decode("utf-16-le", errors="ignore"))
    indicators = {
        "urls": [], "ips": [], "emails": [],
        "registry_keys": [], "file_paths": [],
        "base64_strings": [], "suspicious_strings": [],
    }
    suspicious_keywords = [
        "keylog", "screenshot", "clipboard", "password", "credential",
        "smtp", "telegram", "discord", "webhook", "stealer",
        "inject", "hook", "persist", "startup",
    ]
    for s in all_strings:
        if re.search(r"https?://", s):
            indicators["urls"].append(s)
        if re.search(r"\b(\d{1,3}\.){3}\d{1,3}\b", s):
            indicators["ips"].append(s)
        if re.search(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", s):
            indicators["emails"].append(s)
        if re.search(r"HKLM|HKCU|SOFTWARE\\", s, re.IGNORECASE):
            indicators["registry_keys"].append(s)
        if re.search(r"[A-Za-z0-9+/]{40,}={0,2}$", s):
            indicators["base64_strings"].append(s[:100])
        for kw in suspicious_keywords:
            if kw in s.lower():
                indicators["suspicious_strings"].append(s[:100])
                break
    for key in indicators:
        indicators[key] = list(set(indicators[key]))[:20]
    return indicators


def analyze_dotnet_metadata(filepath):
    """Extract .NET metadata using monodis or ilspy CLI if available."""
    metadata = {}
    try:
        result = subprocess.run(
            ["monodis", "--assembly", filepath],
            capture_output=True, text=True, timeout=15
        )
        if result.returncode == 0:
            metadata["assembly_info"] = result.stdout.strip()
    except FileNotFoundError:
        pass
    try:
        result = subprocess.run(
            ["monodis", "--typedef", filepath],
            capture_output=True, text=True, timeout=15
        )
        if result.returncode == 0:
            types = re.findall(r"(\S+)\s+flags", result.stdout)
            metadata["type_count"] = len(types)
            metadata["types"] = types[:30]
    except FileNotFoundError:
        pass
    try:
        result = subprocess.run(
            ["monodis", "--method", filepath],
            capture_output=True, text=True, timeout=15
        )
        if result.returncode == 0:
            methods = re.findall(r"(\S+)\s+\(", result.stdout)
            metadata["method_count"] = len(methods)
    except FileNotFoundError:
        pass
    return metadata


def analyze_dotnet_malware(filepath, output_dir="/tmp/dotnet_analysis"):
    """Full .NET malware analysis pipeline."""
    os.makedirs(output_dir, exist_ok=True)
    report = {"file": filepath}
    report["hashes"] = compute_hashes(filepath)
    report["dotnet_check"] = detect_dotnet_assembly(filepath)
    if not report["dotnet_check"].get("is_dotnet"):
        report["error"] = "Not a .NET assembly"
        return report
    report["obfuscator"] = detect_obfuscator(filepath)
    deobf_path = os.path.join(output_dir, "deobfuscated.exe")
    report["deobfuscation"] = deobfuscate_with_de4dot(filepath, deobf_path)
    analysis_target = deobf_path if report["deobfuscation"].get("success") else filepath
    report["strings"] = extract_strings(analysis_target)
    report["metadata"] = analyze_dotnet_metadata(analysis_target)
    return report


def print_report(report):
    print(".NET Malware Analysis Report")
    print("=" * 50)
    print(f"File: {report['file']}")
    print(f"SHA-256: {report['hashes']['sha256']}")
    print(f".NET Assembly: {report['dotnet_check'].get('is_dotnet', False)}")
    obf = report.get("obfuscator", {})
    if obf.get("detected"):
        print(f"Obfuscator: {', '.join(obf['detected'])}")
    deobf = report.get("deobfuscation", {})
    print(f"Deobfuscation: {'Success' if deobf.get('success') else 'Failed/Skipped'}")
    meta = report.get("metadata", {})
    if meta:
        print(f"Types: {meta.get('type_count', 'N/A')}, Methods: {meta.get('method_count', 'N/A')}")
    strings = report.get("strings", {})
    if strings:
        print("\nExtracted Indicators:")
        for cat, values in strings.items():
            if values:
                print(f"  {cat}: {len(values)}")
                for v in values[:3]:
                    print(f"    - {v[:80]}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python agent.py <dotnet_executable>")
        sys.exit(1)
    result = analyze_dotnet_malware(sys.argv[1])
    print_report(result)