server.py

"""
MCP server for static analysis and patching of .NET binaries.
Uses dnlib for IL parsing/editing and ILSpy/dnSpy for C# decompilation.
"""
from __future__ import annotations

import logging
import shutil
import sys
import traceback
import urllib.error
import urllib.request
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Any

import click
import uvicorn
from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route

logging.basicConfig(
    level=logging.INFO,
    stream=sys.stderr,
    format='%(asctime)s [%(levelname)s] %(message)s'
)
logger = logging.getLogger(__name__)

# Suppress noisy FastMCP session cleanup errors
logging.getLogger("mcp.server.streamable_http").setLevel(logging.CRITICAL)

import clr
_DNLIB_DIR = Path(__file__).parent / "dll"
clr.AddReference("System")
clr.AddReference(str(_DNLIB_DIR / "dnlib.dll"))
clr.AddReference(str(_DNLIB_DIR / "dnSpy.Contracts.Logic.dll"))
clr.AddReference(str(_DNLIB_DIR / "ICSharpCode.Decompiler.dll"))
ilspy_assembly = clr.AddReference(str(_DNLIB_DIR / "dnSpy.Decompiler.ILSpy.Core.dll"))

import dnlib
import System
from dnSpy.Contracts.Decompiler import StringBuilderDecompilerOutput, DecompilationContext
import ICSharpCode.Decompiler       # noqa: F401 – side-effect: loads assembly into CLR
import dnSpy.Decompiler.ILSpy.Core  # noqa: F401 – side-effect: loads assembly into CLR

from dnlib.DotNet import *
from dnlib.DotNet.Emit import Instruction, OpCodes
from dnlib.DotNet import UTF8String


import os
SAMPLES_DIR = Path(os.environ.get("SAMPLES_DIR", "/data/samples"))
SAMPLES_DIR.mkdir(parents=True, exist_ok=True)

PYTHONNET_PROJECTS_DIR = Path(os.environ.get("PYTHONNET_PROJECTS_DIR", "/data/pythonnet_projects"))
PYTHONNET_PROJECTS_DIR.mkdir(parents=True, exist_ok=True)


# ---------------------------------------------------------------------------
# Module-level constants
# ---------------------------------------------------------------------------

_CALL_OPCODES: frozenset[str] = frozenset({
    str(OpCodes.Call), str(OpCodes.Callvirt),
})

_CALL_OPCODES_EXTENDED: frozenset[str] = frozenset({
    str(OpCodes.Call), str(OpCodes.Callvirt),
    str(OpCodes.Newobj), str(OpCodes.Ldftn),
})

_RUNTIME_ENTRY_POINTS: frozenset[str] = frozenset({
    ".ctor", ".cctor", "Main", "Finalize",
    "Dispose", "GetHashCode", "Equals", "ToString",
})

_CRYPTO_NAMESPACES: frozenset[str] = frozenset({
    "System.Security.Cryptography", "System.Security",
})

_CRYPTO_KEYWORDS: frozenset[str] = frozenset({
    "aes", "rsa", "sha", "md5", "des", "rc4", "base64", "cipher",
    "encrypt", "decrypt", "rijndael", "hmac", "pbkdf", "salsa", "chacha",
})

_CRYPTO_KEY_SIZES: frozenset[int] = frozenset({16, 24, 32})

_REFLECTION_MARKERS: frozenset[str] = frozenset({
    "system.reflection", "assembly.load", "getmethod", "gettype",
    "invoke", "activator.createinstance", "emit",
})

_BRANCH_INVERSIONS: dict[str, Any] = {
    str(OpCodes.Brtrue):    OpCodes.Brfalse,
    str(OpCodes.Brtrue_S):  OpCodes.Brfalse_S,
    str(OpCodes.Brfalse):   OpCodes.Brtrue,
    str(OpCodes.Brfalse_S): OpCodes.Brtrue_S,
    str(OpCodes.Beq):       OpCodes.Bne_Un,
    str(OpCodes.Beq_S):     OpCodes.Bne_Un_S,
    str(OpCodes.Bne_Un):    OpCodes.Beq,
    str(OpCodes.Bne_Un_S):  OpCodes.Beq_S,
    str(OpCodes.Blt):       OpCodes.Bge,
    str(OpCodes.Blt_S):     OpCodes.Bge_S,
    str(OpCodes.Bge):       OpCodes.Blt,
    str(OpCodes.Bge_S):     OpCodes.Blt_S,
    str(OpCodes.Bgt):       OpCodes.Ble,
    str(OpCodes.Bgt_S):     OpCodes.Ble_S,
    str(OpCodes.Ble):       OpCodes.Bgt,
    str(OpCodes.Ble_S):     OpCodes.Bgt_S,
}


# ---------------------------------------------------------------------------
# Context management
# ---------------------------------------------------------------------------

_pythonnet_contexts: dict[str, dict[str, Any]] = {}


class PythonnetContext:
    """Holds a loaded dnlib module and lazily-initialized decompiler."""

    binary_path: Path
    module: Any  # dnlib.DotNet.ModuleDefMD
    _type_index: dict[str, Any]
    _decompilation_context: Any | None
    _decompiler: Any | None

    def __init__(self, binary_path: Path) -> None:
        self.binary_path = binary_path
        logger.info(f"Loading .NET module: {binary_path}")
        self.module = dnlib.DotNet.ModuleDefMD.Load(str(binary_path))
        self._type_index = {str(t): t for t in self.module.GetTypes()}
        self._decompilation_context = None
        self._decompiler = None
        logger.info(f"PythonnetContext ready: {self.module} ({len(self._type_index)} types)")

    def _ensure_decompiler(self) -> None:
        """Lazily initialise the ILSpy C# decompiler (expensive, done once)."""
        if self._decompiler:
            return
        self._decompilation_context = System.Activator.CreateInstance(DecompilationContext)
        csharp_type = ilspy_assembly.GetType(
            "dnSpy.Decompiler.ILSpy.Core.CSharp.CSharpDecompiler"
        )
        settings_type = ilspy_assembly.GetType(
            "dnSpy.Decompiler.ILSpy.Core.Settings.CSharpVBDecompilerSettings"
        )
        settings = System.Activator.CreateInstance(settings_type, None)
        self._decompiler = System.Activator.CreateInstance(csharp_type, settings, None)

    def get_type(self, class_name: str) -> Any:
        """Return the dnlib TypeDef for *class_name* or raise KeyError."""
        t = self._type_index.get(class_name)
        if not t:
            raise KeyError(f"Class not found: {class_name}")
        return t

    def close(self) -> None:
        if self.module:
            self.module.Dispose()
            self.module = None
        logger.info("PythonnetContext closed")


def _get_context(analysis_id: str = "default") -> PythonnetContext:
    """Return the active context for *analysis_id* and refresh last_accessed."""
    if analysis_id not in _pythonnet_contexts:
        raise ValueError(
            f"No context loaded for analysis '{analysis_id}'. "
            "Call pythonnet_load_binary first."
        )
    entry = _pythonnet_contexts[analysis_id]
    entry["last_accessed"] = datetime.now()
    return entry["context"]


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _parse_il_offset(offset: str) -> int:
    """Parse an IL offset string to an integer.

    Accepts ``IL_0042``, ``0x42``, and plain decimal/hex strings.

    Raises:
        ValueError: if the string cannot be parsed.
    """
    s = offset.strip().upper().removeprefix("IL_").removeprefix("0X")
    try:
        return int(s, 16)
    except ValueError:
        raise ValueError(
            f"Invalid offset format: {offset!r} — use IL_XXXX, 0xN, or a decimal integer"
        )


def _find_method(t: Any, method_name: str) -> Any | None:
    """Return the first method on *t* whose Name equals *method_name*, or None."""
    return next((m for m in t.Methods if m.Name == method_name), None)


def _resolve_method(t: Any, method_name: str, method_index: int = -1) -> Any:
    """Resolve a method by name on *t*, handling overloads and full signatures.

    *method_name* may be a bare name (``"Foo"``) or a full signature as
    returned by dnlib / ``list_all_methods`` (``"Foo(System.String)"`` or
    ``"System.Void Ns.Class::Foo(System.String)"``). The signature portion is
    stripped before matching so both forms work transparently.

    - Exactly one match → returns the method directly.
    - Multiple matches and ``method_index == -1`` → returns a dict with key
      ``"overloads"`` listing all candidates (index + full signature) so the
      caller can retry with a specific ``method_index``.
    - ``method_index >= 0`` → returns the method at that position.

    The return value is either a live method object or a ``dict`` (error or
    overload disambiguation). Callers should check ``isinstance(result, dict)``
    and return it early on that branch.
    """
    # Strip parameter list: "Foo(System.String)" -> "Foo"
    bare_name = method_name.split("(")[0].strip()
    # Strip return type + class prefix: "System.Void Ns.Class::Foo" -> "Foo"
    if "::" in bare_name:
        bare_name = bare_name.split("::")[-1].strip()

    matches = [m for m in t.Methods if m.Name == bare_name]
    if not matches:
        return {"error": f"Method not found: {t.Name}.{method_name}"}
    if len(matches) == 1:
        return matches[0]
    if method_index >= 0:
        if method_index >= len(matches):
            return {"error": f"method_index {method_index} out of range (0–{len(matches) - 1})"}
        return matches[method_index]
    return {
        "overloads": [
            {"index": i, "signature": str(m)}
            for i, m in enumerate(matches)
        ],
        "hint": (
            f"{len(matches)} overloads found for '{method_name}'. "
            "Re-call with method_index=N to select one."
        ),
    }


# ---------------------------------------------------------------------------
# MCP tools
# ---------------------------------------------------------------------------

def pythonnet_load_binary(binary_path: str, analysis_id: str = "default") -> str:
    """Load a .NET binary for analysis.

    The binary is copied into /data/pythonnet_projects/{analysis_id}/ so that
    any subsequent save operations land in a dedicated project directory rather
    than overwriting the original sample.

    Args:
        binary_path: Absolute path to the .NET binary (e.g. /data/samples/malware.exe).
        analysis_id: Session identifier — use different IDs to analyse multiple
                     binaries simultaneously.

    Returns:
        A human-readable status string with the module name, type count, and
        the project directory path.
    """
    global _pythonnet_contexts

    if analysis_id in _pythonnet_contexts:
        logger.info(f"Closing previous context for analysis '{analysis_id}'…")
        _pythonnet_contexts[analysis_id]["context"].close()
        del _pythonnet_contexts[analysis_id]

    src = Path(binary_path)
    if not src.is_absolute():
        # Bare filename: search /data/samples/ then /data/pythonnet_projects/_uploads/
        for search_dir in (SAMPLES_DIR, PYTHONNET_PROJECTS_DIR / "_uploads"):
            candidate = search_dir / src.name
            if candidate.exists():
                src = candidate
                break
        else:
            return f"Error: File not found: {binary_path} (searched {SAMPLES_DIR} and {PYTHONNET_PROJECTS_DIR / '_uploads'})"
    elif not src.exists():
        return f"Error: File not found: {binary_path}"

    try:
        # Create a per-analysis project directory and copy the binary there.
        project_dir = PYTHONNET_PROJECTS_DIR / analysis_id
        project_dir.mkdir(parents=True, exist_ok=True)
        working_copy = project_dir / src.name
        shutil.copy2(src, working_copy)
        logger.info(f"Copied {src} -> {working_copy}")

        ctx = PythonnetContext(binary_path=working_copy)
        _pythonnet_contexts[analysis_id] = {
            "context": ctx,
            "project_dir": project_dir,
            "original_name": src.name,
            "last_accessed": datetime.now(),
            "loaded_at": datetime.now(),
        }
        return (
            f"Loaded {binary_path}\n"
            f"Working copy : {working_copy}\n"
            f"Module : {ctx.module}\n"
            f"Types  : {len(ctx._type_index)}\n"
            f"Project dir  : {project_dir}\n"
            f"Ready for analysis (use pythonnet_* tools)."
        )
    except Exception as e:
        logger.error(f"Error loading binary: {e}", exc_info=True)
        return f"Error loading binary: {e}"


def pythonnet_list_all_methods(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """List all methods with a body in the loaded binary.

    Args:
        limit: Maximum number of methods to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "methods": [str, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    methods = [
        str(m)
        for t in ctx._type_index.values()
        for m in t.Methods
        if m.HasBody
    ]
    if limit:
        methods = methods[:limit]

    return {"count": len(methods), "methods": methods}


def pythonnet_list_all_methods_inside_class(
    class_name: str,
    analysis_id: str = "default",
) -> dict:
    """List all methods (with or without body) inside a specific class.

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "count": int, "methods": [str, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    methods = [str(m) for m in t.Methods]
    return {"class": class_name, "count": len(methods), "methods": methods}


def pythonnet_list_all_classes(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """List all types defined in the loaded binary.

    Args:
        limit: Maximum number of classes to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "classes": [str, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    classes = list(ctx._type_index.keys())
    if limit:
        classes = classes[:limit]

    return {"count": len(classes), "classes": classes}


def pythonnet_decompile_class(class_name: str, analysis_id: str = "default") -> dict:
    """Decompile an entire class to C# source code.

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "decompiled_code": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    ctx._ensure_decompiler()
    out = System.Activator.CreateInstance(StringBuilderDecompilerOutput)
    ctx._decompiler.Decompile(t, out, ctx._decompilation_context)
    return {"class": class_name, "decompiled_code": out.GetText()}


def pythonnet_decompile_method(
    class_name: str,
    method_name: str,
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Decompile a single method to C# source code.

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        method_name: Method name to decompile.
        method_index: 0-based index to select among overloads. Omit (or pass
                      ``-1``) for auto-selection; if multiple overloads exist
                      the response will list them with their indices.
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "decompiled_code": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method

    ctx._ensure_decompiler()
    out = System.Activator.CreateInstance(StringBuilderDecompilerOutput)
    ctx._decompiler.Decompile(method, out, ctx._decompilation_context)
    logger.info(f"Decompiled {class_name}.{method_name}")
    return {"class": class_name, "method": method_name, "decompiled_code": out.GetText()}


def pythonnet_decompile_all_classes(
    skip_compiler_generated: bool = True,
    analysis_id: str = "default",
) -> dict:
    """Decompile every class in the binary to C# source in one call.

    Expensive but useful for handing the full codebase to an LLM for global
    analysis. Classes that fail are collected in ``errors`` rather than
    aborting the run.

    Args:
        skip_compiler_generated: Skip types whose name starts with ``<`` or
                                 ``$`` (compiler-generated noise). Default ``True``.
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "error_count": int, "classes": {name: src}, "errors": {name: msg}}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    ctx._ensure_decompiler()
    classes: dict[str, str] = {}
    errors: dict[str, str] = {}

    for class_name, t in ctx._type_index.items():
        if skip_compiler_generated:
            name = str(t.Name)
            if name.startswith("<") or name.startswith("$"):
                continue
        try:
            out = System.Activator.CreateInstance(StringBuilderDecompilerOutput)
            ctx._decompiler.Decompile(t, out, ctx._decompilation_context)
            classes[class_name] = out.GetText()
        except Exception as e:
            errors[class_name] = str(e)

    return {
        "count": len(classes),
        "error_count": len(errors),
        "classes": classes,
        "errors": errors,
    }


def pythonnet_rename_class(
    class_name: str,
    new_name: str,
    new_namespace: str = "",
    analysis_id: str = "default",
) -> dict:
    """Rename a class and optionally its namespace (useful for deobfuscation).

    Args:
        class_name: Full class path to rename (e.g. ``OldNs.OldClass``).
        new_name: New class name (without namespace).
        new_namespace: If provided, also update the namespace.
        analysis_id: Analysis session ID.

    Returns:
        ``{"old_class_name": str, "new_class_name": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    t.Name = UTF8String(new_name)
    if new_namespace:
        t.Namespace = UTF8String(new_namespace)

    # Recompute full name the same way the index was built
    new_full_name = str(t)
    ctx._type_index[new_full_name] = ctx._type_index.pop(class_name)
    logger.info(f"Renamed class {class_name} -> {new_full_name}")
    return {"old_class_name": class_name, "new_class_name": new_full_name}


def pythonnet_rename_method(
    class_name: str,
    method_name: str,
    new_method_name: str,
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Rename a method (useful for deobfuscation).

    Args:
        class_name: Full class path containing the method.
        method_name: Current method name.
        new_method_name: New method name.
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "old_method_name": str, "new_method_name": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method

    method.Name = UTF8String(new_method_name)
    logger.info(f"Renamed method {class_name}.{method_name} -> {new_method_name}")
    return {"class": class_name, "old_method_name": method_name, "new_method_name": new_method_name}


def pythonnet_rename_field(
    class_name: str,
    field_name: str,
    new_field_name: str,
    analysis_id: str = "default",
) -> dict:
    """Rename a field (useful for deobfuscation).

    Args:
        class_name: Full class path containing the field.
        field_name: Current field name.
        new_field_name: New field name.
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "old_field_name": str, "new_field_name": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    field = next((f for f in t.Fields if f.Name == field_name), None)
    if field is None:
        return {"error": f"Field not found: {class_name}.{field_name}"}

    field.Name = UTF8String(new_field_name)
    logger.info(f"Renamed field {class_name}.{field_name} -> {new_field_name}")
    return {"class": class_name, "old_field_name": field_name, "new_field_name": new_field_name}


def pythonnet_save_binary(analysis_id: str = "default") -> dict:
    """Save the (possibly modified) module to disk.

    Must be called after any rename/patch operation to persist changes.
    Always saves to the project directory as ``{original_stem}_patched{ext}``
    (e.g. ``/data/pythonnet_projects/{analysis_id}/malware_patched.dll``).

    Args:
        analysis_id: Analysis session ID.

    Returns:
        ``{"output_path": str, "project_dir": str, "status": "saved", "download_url": str}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    entry = _pythonnet_contexts[analysis_id]
    project_dir: Path = entry.get("project_dir", PYTHONNET_PROJECTS_DIR / analysis_id)
    project_dir.mkdir(parents=True, exist_ok=True)

    original_name = entry.get("original_name") or Path(str(ctx.binary_path)).name
    stem = Path(original_name).stem
    suffix = Path(original_name).suffix
    output_path = str(project_dir / f"{stem}_patched{suffix}")

    try:
        ctx.module.Write(output_path)
        logger.info(f"Saved binary to {output_path}")
        return {
            "output_path": output_path,
            "project_dir": str(project_dir),
            "status": "saved",
            "download_url": f"/download?path={output_path}",
        }
    except Exception as e:
        logger.error(f"Error saving binary: {e}", exc_info=True)
        return {"error": str(e)}


def pythonnet_get_strings(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """List all unique string literals (``ldstr``) across all method bodies.

    Args:
        limit: Maximum number of strings to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "strings": [str, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    seen: dict[str, None] = {}
    for t in ctx._type_index.values():
        for method in t.Methods:
            if not method.HasBody:
                continue
            for instr in method.Body.Instructions:
                if instr.OpCode == OpCodes.Ldstr and instr.Operand is not None:
                    seen[str(instr.Operand)] = None

    strings = list(seen)
    if limit:
        strings = strings[:limit]

    return {"count": len(strings), "strings": strings}


def pythonnet_search_string(
    query: str,
    case_sensitive: bool = False,
    analysis_id: str = "default",
) -> dict:
    """Find all methods that reference a given string literal.

    Inverse of ``pythonnet_get_strings``: start from a suspicious string and
    trace back to the code that uses it.

    Args:
        query: String (or partial string) to search for.
        case_sensitive: Case-sensitive match (default ``False``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"query": str, "count": int, "matches": [{"class", "method", "string"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    needle = query if case_sensitive else query.lower()
    matches = []
    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            if not method.HasBody:
                continue
            for instr in method.Body.Instructions:
                if instr.OpCode != OpCodes.Ldstr or instr.Operand is None:
                    continue
                value = str(instr.Operand)
                haystack = value if case_sensitive else value.lower()
                if needle in haystack:
                    matches.append({
                        "class": class_name,
                        "method": str(method.Name),
                        "string": value,
                    })
                    break  # one hit per method is enough

    return {"query": query, "count": len(matches), "matches": matches}


def pythonnet_search_method_by_name(
    query: str,
    case_sensitive: bool = False,
    analysis_id: str = "default",
) -> dict:
    """Search for methods by partial name across all classes.

    Useful when code is obfuscated but partial names remain
    (e.g. searching ``"decrypt"`` or ``"http"``).

    Args:
        query: Partial method name to search for.
        case_sensitive: Case-sensitive match (default ``False``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"query": str, "count": int, "matches": [{"class", "method", "has_body"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    needle = query if case_sensitive else query.lower()
    matches = []
    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            name = str(method.Name)
            haystack = name if case_sensitive else name.lower()
            if needle in haystack:
                matches.append({
                    "class": class_name,
                    "method": name,
                    "has_body": bool(method.HasBody),
                })

    return {"query": query, "count": len(matches), "matches": matches}


def pythonnet_get_imports(analysis_id: str = "default") -> dict:
    """List all external assemblies and types referenced by the binary.

    Useful for quick orientation: ``System.Net.WebClient`` → network activity,
    ``System.Security.Cryptography`` → crypto, etc.

    Args:
        analysis_id: Analysis session ID.

    Returns:
        ``{"assembly_refs": [str], "imports_by_assembly": {assembly: [type, …]}}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    assembly_refs = [str(a) for a in ctx.module.GetAssemblyRefs()]

    imports_by_assembly: dict[str, list[str]] = {}
    for type_ref in ctx.module.GetTypeRefs():
        scope = str(type_ref.Scope) if type_ref.Scope is not None else "<unknown>"
        full_name = str(type_ref)
        bucket = imports_by_assembly.setdefault(scope, [])
        if full_name not in bucket:
            bucket.append(full_name)

    return {"assembly_refs": assembly_refs, "imports_by_assembly": imports_by_assembly}


def pythonnet_get_fields(class_name: str, analysis_id: str = "default") -> dict:
    """List all fields of a class with their types and visibility.

    Essential for understanding the layout of obfuscated classes.

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "count": int, "fields": [{"name", "type", "static", "visibility"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    fields = []
    for f in t.Fields:
        fields.append({
            "name": str(f.Name),
            "type": str(f.FieldType),
            "static": bool(f.IsStatic),
            "visibility": (
                "public"    if f.IsPublic    else
                "private"   if f.IsPrivate   else
                "protected" if f.IsFamily    else
                "internal"  if f.IsAssembly  else
                "other"
            ),
        })

    return {"class": class_name, "count": len(fields), "fields": fields}


def pythonnet_get_custom_attributes(
    class_name: str,
    method_name: str = "",
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Read custom attributes on a class or one of its methods.

    Revealing in .NET malware: C2 frameworks (Covenant ``[Task]``),
    obfuscators, and packers all leave characteristic attributes
    (``[Obfuscation]``, ``[UnmanagedFunctionPointer]``, etc.).

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        method_name: If provided, read attributes on this method instead of the class.
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"target": str, "count": int, "attributes": [{"attribute", "constructor_args", "named_args"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    if method_name:
        target_obj = _resolve_method(t, method_name, method_index)
        if isinstance(target_obj, dict):
            return target_obj
        target_label = f"{class_name}.{method_name}"
        attr_collection = target_obj.CustomAttributes
    else:
        target_label = class_name
        attr_collection = t.CustomAttributes

    attributes = []
    for attr in attr_collection:
        attributes.append({
            "attribute": str(attr.AttributeType),
            "constructor_args": [
                {"type": str(a.Type), "value": str(a.Value)}
                for a in attr.ConstructorArguments
            ],
            "named_args": [
                {"name": str(a.Name), "type": str(a.Type), "value": str(a.Value)}
                for a in attr.NamedArguments
            ],
        })

    return {"target": target_label, "count": len(attributes), "attributes": attributes}


def pythonnet_get_opcodes(
    class_name: str,
    method_name: str,
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Dump the raw IL instructions of a method.

    Useful when C# decompilation fails or produces misleading output — IL is
    always the ground truth.

    Args:
        class_name: Full class path (e.g. ``Namespace.ClassName``).
        method_name: Method to inspect.
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "count": int, "instructions": [{"offset", "opcode", "operand"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    instructions = [
        {
            "offset": f"IL_{instr.Offset:04X}",
            "opcode": str(instr.OpCode.Name),
            "operand": str(instr.Operand) if instr.Operand is not None else None,
        }
        for instr in method.Body.Instructions
    ]

    return {
        "class": class_name,
        "method": method_name,
        "count": len(instructions),
        "instructions": instructions,
    }


def pythonnet_get_method_calls(
    class_name: str,
    method_name: str,
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """List all outgoing calls from a method (``call``/``callvirt`` instructions).

    Allows building a call graph without decompiling. Each entry includes the
    target's full dnlib name and whether the dispatch is virtual.

    Args:
        class_name: Full class path containing the method.
        method_name: Method to inspect.
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "count": int, "calls": [{"callee", "virtual"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    callvirt = str(OpCodes.Callvirt)
    calls = [
        {
            "callee": str(instr.Operand),
            "virtual": str(instr.OpCode) == callvirt,
        }
        for instr in method.Body.Instructions
        if str(instr.OpCode) in _CALL_OPCODES and instr.Operand is not None
    ]

    return {"class": class_name, "method": method_name, "count": len(calls), "calls": calls}


def pythonnet_get_callers(
    class_name: str,
    method_name: str,
    analysis_id: str = "default",
) -> dict:
    """Find all methods that call a given method (reverse call graph).

    Scans every method body for ``call``/``callvirt`` instructions targeting
    the specified method.

    Args:
        class_name: Full class path of the target method.
        method_name: Target method name.
        analysis_id: Analysis session ID.

    Returns:
        ``{"target": str, "count": int, "callers": [{"class", "method"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    if not any(m.Name == method_name for m in t.Methods):
        return {"error": f"Method not found: {class_name}.{method_name}"}

    callers = []
    for caller_class, caller_type in ctx._type_index.items():
        for caller_method in caller_type.Methods:
            if not caller_method.HasBody:
                continue
            for instr in caller_method.Body.Instructions:
                if str(instr.OpCode) not in _CALL_OPCODES or instr.Operand is None:
                    continue
                operand = str(instr.Operand)
                if class_name in operand and method_name in operand:
                    callers.append({"class": caller_class, "method": str(caller_method.Name)})
                    break  # one hit per caller method is enough

    return {"target": f"{class_name}.{method_name}", "count": len(callers), "callers": callers}


def pythonnet_patch_ldstr(
    class_name: str,
    method_name: str,
    offset: str,
    new_value: str,
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Replace the string operand of an ``ldstr`` instruction at a given IL offset.

    Use ``pythonnet_get_opcodes`` first to find the target offset, then
    ``pythonnet_nop_instructions`` to remove the decrypt call.

    Args:
        class_name: Full class path containing the method.
        method_name: Method to patch.
        offset: IL offset of the ``ldstr`` (e.g. ``"IL_0004"``, ``"0x4"``, or ``"4"``).
        new_value: Replacement string (e.g. the decrypted plaintext).
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "offset": str, "old_value": str, "new_value": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    try:
        target_offset = _parse_il_offset(offset)
    except ValueError as e:
        return {"error": str(e)}

    instr = next((i for i in method.Body.Instructions if i.Offset == target_offset), None)
    if instr is None:
        return {"error": f"No instruction at offset {offset} in {class_name}.{method_name}"}
    if instr.OpCode != OpCodes.Ldstr:
        return {"error": f"Instruction at {offset} is {instr.OpCode.Name!r}, not ldstr"}

    old_value = str(instr.Operand)
    instr.Operand = new_value
    logger.info(f"Patched ldstr at {offset} in {class_name}.{method_name}: {old_value!r} -> {new_value!r}")
    return {
        "class": class_name,
        "method": method_name,
        "offset": offset,
        "old_value": old_value,
        "new_value": new_value,
    }


def pythonnet_nop_instructions(
    class_name: str,
    method_name: str,
    offsets: list[str],
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Replace one or more instructions with ``nop`` at the given IL offsets.

    Useful to erase a decrypt call after the surrounding ``ldstr`` has been
    patched to the plaintext value.

    Args:
        class_name: Full class path containing the method.
        method_name: Method to patch.
        offsets: List of IL offsets to NOP out (e.g. ``["IL_0006", "IL_0008"]``).
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "nopped": [{"offset", "original_opcode"}], "errors": […]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    instr_map = {i.Offset: i for i in method.Body.Instructions}
    nopped = []
    errors = []

    for offset in offsets:
        try:
            target_offset = _parse_il_offset(offset)
        except ValueError as e:
            errors.append({"offset": offset, "error": str(e)})
            continue

        instr = instr_map.get(target_offset)
        if instr is None:
            errors.append({"offset": offset, "error": "no instruction at this offset"})
            continue

        original_opcode = str(instr.OpCode.Name)
        instr.OpCode = OpCodes.Nop
        instr.Operand = None
        nopped.append({"offset": offset, "original_opcode": original_opcode})
        logger.info(f"NOPed {original_opcode} at {offset} in {class_name}.{method_name}")

    return {"class": class_name, "method": method_name, "nopped": nopped, "errors": errors}


def pythonnet_patch_return(
    class_name: str,
    method_name: str,
    return_value: str = "true",
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Replace a method body with a minimal return sequence (bypass a check).

    Useful for bypassing licence checks, auth gates, or anti-debug routines.
    Clears the existing body and replaces it with the shortest valid return.

    Args:
        class_name: Full class path containing the method.
        method_name: Method to patch.
        return_value: What to return — ``"true"``/``"1"``, ``"false"``/``"0"``,
                      ``"void"``, or an integer string for ``ldc.i4``. Default ``"true"``.
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "injected": [str, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    method.Body.Instructions.Clear()
    method.Body.ExceptionHandlers.Clear()
    method.Body.Variables.Clear()

    rv = return_value.strip().lower()
    if rv == "void":
        injected = ["ret"]
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret))
    elif rv in ("true", "1"):
        injected = ["ldc.i4.1", "ret"]
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_1))
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret))
    elif rv in ("false", "0"):
        injected = ["ldc.i4.0", "ret"]
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_0))
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret))
    else:
        try:
            int_val = int(return_value)
        except ValueError:
            return {
                "error": (
                    f"Unsupported return_value: {return_value!r} — "
                    "use 'true', 'false', 'void', or an integer"
                )
            }
        injected = [f"ldc.i4 {int_val}", "ret"]
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4, int_val))
        method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret))

    method.Body.UpdateInstructionOffsets()
    logger.info(f"Patched return of {class_name}.{method_name} -> {return_value}")
    return {"class": class_name, "method": method_name, "injected": injected}


def pythonnet_patch_branch(
    class_name: str,
    method_name: str,
    offset: str,
    action: str = "invert",
    method_index: int = -1,
    analysis_id: str = "default",
) -> dict:
    """Flip or force a conditional branch instruction at a given IL offset.

    Actions:

    * ``"invert"`` — ``brtrue`` ↔ ``brfalse``, ``beq`` ↔ ``bne``, ``blt`` ↔ ``bge``, etc.
    * ``"force"``  — replace with an unconditional ``br`` (always taken).
    * ``"nop"``    — erase the branch entirely (never taken).

    Args:
        class_name: Full class path containing the method.
        method_name: Method to patch.
        offset: IL offset of the branch instruction (e.g. ``"IL_001A"``).
        action: ``"invert"``, ``"force"``, or ``"nop"`` (default ``"invert"``).
        method_index: 0-based index to select among overloads (default ``-1`` = auto).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "offset": str, "old_opcode": str, "new_opcode": str, "action": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    method = _resolve_method(t, method_name, method_index)
    if isinstance(method, dict):
        return method
    if not method.HasBody:
        return {"error": f"Method has no body: {class_name}.{method_name}"}

    try:
        target_offset = _parse_il_offset(offset)
    except ValueError as e:
        return {"error": str(e)}

    instr = next((i for i in method.Body.Instructions if i.Offset == target_offset), None)
    if instr is None:
        return {"error": f"No instruction at offset {offset}"}

    old_opcode = str(instr.OpCode.Name)
    action = action.strip().lower()

    if action == "invert":
        new_opcode = _BRANCH_INVERSIONS.get(str(instr.OpCode))
        if new_opcode is None:
            return {"error": f"Cannot invert opcode {old_opcode!r} — not a known conditional branch"}
        instr.OpCode = new_opcode
    elif action == "force":
        instr.OpCode = OpCodes.Br_S if "_S" in str(instr.OpCode) else OpCodes.Br
    elif action == "nop":
        instr.OpCode = OpCodes.Nop
        instr.Operand = None
    else:
        return {"error": f"Unknown action {action!r} — use 'invert', 'force', or 'nop'"}

    logger.info(
        f"Patched branch at {offset} in {class_name}.{method_name}: "
        f"{old_opcode} -> {instr.OpCode.Name} ({action})"
    )
    return {
        "class": class_name,
        "method": method_name,
        "offset": offset,
        "old_opcode": old_opcode,
        "new_opcode": str(instr.OpCode.Name),
        "action": action,
    }


def pythonnet_set_field_constant(
    class_name: str,
    field_name: str,
    new_value: str,
    analysis_id: str = "default",
) -> dict:
    """Patch the initialization value of a static field in ``.cctor``.

    Finds the ``stsfld`` for the field in the static constructor and rewrites
    the preceding ``ldc.i4*`` instruction.

    Args:
        class_name: Full class path containing the field.
        field_name: Static field name to patch.
        new_value: New integer value as a string (e.g. ``"0"``, ``"1"``, ``"42"``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "field": str, "offset": str, "old_value": str, "new_value": str}``
    """
    try:
        ctx = _get_context(analysis_id)
        t = ctx.get_type(class_name)
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    cctor = _find_method(t, ".cctor")
    if cctor is None:
        return {"error": f"No static constructor (.cctor) in {class_name}"}
    if not cctor.HasBody:
        return {"error": f".cctor has no body in {class_name}"}

    try:
        int_value = int(new_value)
    except ValueError:
        return {"error": f"new_value must be an integer string, got {new_value!r}"}

    instrs = list(cctor.Body.Instructions)
    for i, instr in enumerate(instrs):
        if str(instr.OpCode.Name) != "stsfld" or instr.Operand is None:
            continue
        if field_name not in str(instr.Operand):
            continue
        if i == 0:
            return {"error": "stsfld is the first instruction — no preceding ldc to patch"}
        prev = instrs[i - 1]
        if not str(prev.OpCode.Name).startswith("ldc.i4"):
            return {"error": f"Instruction before stsfld is {prev.OpCode.Name!r}, expected ldc.i4*"}

        old_value = str(prev.Operand)
        prev.OpCode = OpCodes.Ldc_I4
        prev.Operand = int_value
        logger.info(f"Patched {class_name}.{field_name} init: {old_value} -> {new_value}")
        return {
            "class": class_name,
            "field": field_name,
            "offset": f"IL_{prev.Offset:04X}",
            "old_value": old_value,
            "new_value": new_value,
        }

    return {"error": f"stsfld for field {field_name!r} not found in .cctor of {class_name}"}


def pythonnet_get_entry_point(analysis_id: str = "default") -> dict:
    """Return the managed entry point of the module (``Main`` or ``EntryPoint`` metadata).

    First reflex on an unknown binary to understand where execution starts.

    Args:
        analysis_id: Analysis session ID.

    Returns:
        ``{"class": str, "method": str, "full_name": str}`` or
        ``{"entry_point": null, "note": str}`` for DLLs.
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    ep = ctx.module.EntryPoint
    if ep is None:
        return {
            "entry_point": None,
            "note": "No managed entry point (DLL or native entry point)",
        }

    return {
        "class": str(ep.DeclaringType),
        "method": str(ep.Name),
        "full_name": str(ep),
    }


def pythonnet_get_module_info(analysis_id: str = "default") -> dict:
    """Return global metadata about the .NET module and assembly.

    Covers architecture, runtime version, assembly version, culture, public
    key token, and assembly-level custom attributes (obfuscator watermarks
    often appear here).

    Args:
        analysis_id: Analysis session ID.

    Returns:
        ``{"module": str, "machine": str, "cor_flags": str, "runtime_version": str, "assembly": {…}}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    mod = ctx.module
    asm = mod.Assembly
    asm_info = None
    if asm is not None:
        asm_info = {
            "name": str(asm.Name),
            "version": str(asm.Version),
            "culture": str(asm.Culture) if asm.Culture else None,
            "public_key_token": str(asm.PublicKeyToken) if asm.PublicKeyToken else None,
            "custom_attributes": [str(a.AttributeType) for a in asm.CustomAttributes],
        }

    return {
        "module": str(mod),
        "machine": str(mod.Machine),
        "cor_flags": str(mod.Cor20HeaderFlags),
        "runtime_version": str(mod.RuntimeVersion),
        "assembly": asm_info,
    }


def pythonnet_get_resources(analysis_id: str = "default") -> dict:
    """List all embedded resources in the binary.

    Resources are commonly used to hide payloads, encrypted configs, or
    dropped assemblies.

    Args:
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "resources": [{"name", "type", "size"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    resources = []
    for r in ctx.module.Resources:
        entry: dict[str, Any] = {"name": str(r.Name), "type": str(r.ResourceType)}
        try:
            entry["size"] = r.CreateReader().Length
        except Exception:
            entry["size"] = None
        resources.append(entry)

    return {"count": len(resources), "resources": resources}


def pythonnet_extract_resource(
    resource_name: str,
    analysis_id: str = "default",
) -> dict:
    """Extract a raw embedded resource to a file for separate analysis.

    Useful when a resource contains a dropped payload, encrypted config, or
    a secondary assembly.  The file is always saved inside the project
    directory (``/data/pythonnet_projects/{analysis_id}/{resource_name}``).

    Args:
        resource_name: Exact resource name (from ``pythonnet_get_resources``).
        analysis_id: Analysis session ID.

    Returns:
        ``{"resource": str, "output_path": str, "bytes_written": int, "download_url": str}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    resource = next(
        (r for r in ctx.module.Resources if str(r.Name) == resource_name), None
    )
    if resource is None:
        return {"error": f"Resource not found: {resource_name!r}"}

    entry = _pythonnet_contexts[analysis_id]
    project_dir: Path = entry.get("project_dir", PYTHONNET_PROJECTS_DIR / analysis_id)
    project_dir.mkdir(parents=True, exist_ok=True)
    # Sanitise resource name to avoid path traversal
    safe_name = Path(resource_name).name or "resource.bin"
    output_path = str(project_dir / safe_name)

    try:
        data = bytes(resource.CreateReader().ReadRemainingBytes())
        Path(output_path).write_bytes(data)
        logger.info(f"Extracted resource {resource_name!r} -> {output_path} ({len(data)} bytes)")
        return {
            "resource": resource_name,
            "output_path": output_path,
            "bytes_written": len(data),
            "download_url": f"/download?path={output_path}",
        }
    except Exception as e:
        logger.error(f"Error extracting resource: {e}", exc_info=True)
        return {"error": str(e)}


def pythonnet_find_pinvoke(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """List all P/Invoke declarations (``DllImport``) in the binary.

    Reveals native API usage directly. ``VirtualAlloc``, ``WriteProcessMemory``,
    ``CreateRemoteThread`` → shellcode injection; WinINet/WinHTTP → C2 comms.

    Args:
        limit: Maximum number of entries to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "pinvokes": [{"class", "managed_name", "native_name", "dll"}], "by_dll": {dll: [str]}}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    pinvokes = []
    by_dll: dict[str, list[str]] = {}

    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            if not method.IsPinvokeImpl or method.ImplMap is None:
                continue
            impl = method.ImplMap
            dll = str(impl.Module.Name) if impl.Module is not None else "<unknown>"
            native_name = str(impl.Name) if impl.Name else str(method.Name)
            pinvokes.append({
                "class": class_name,
                "managed_name": str(method.Name),
                "native_name": native_name,
                "dll": dll,
            })
            by_dll.setdefault(dll, []).append(
                f"{class_name}.{method.Name} -> {native_name}"
            )

    if limit:
        pinvokes = pinvokes[:limit]

    return {"count": len(pinvokes), "pinvokes": pinvokes, "by_dll": by_dll}


def pythonnet_find_reflection_calls(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """List all uses of ``System.Reflection`` in the binary.

    Reflection is a classic sign of payload droppers, in-memory loaders, and
    obfuscators (``Assembly.Load``, ``MethodInfo.Invoke``, ``Type.GetMethod``, …).

    Args:
        limit: Maximum number of matches to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "matches": [{"class", "method", "offset", "callee"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    matches = []
    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            if not method.HasBody:
                continue
            for instr in method.Body.Instructions:
                if str(instr.OpCode) not in _CALL_OPCODES or instr.Operand is None:
                    continue
                operand = str(instr.Operand).lower()
                if any(marker in operand for marker in _REFLECTION_MARKERS):
                    matches.append({
                        "class": class_name,
                        "method": str(method.Name),
                        "offset": f"IL_{instr.Offset:04X}",
                        "callee": str(instr.Operand),
                    })

    if limit:
        matches = matches[:limit]

    return {"count": len(matches), "matches": matches}


def pythonnet_find_large_byte_arrays(
    min_size: int = 256,
    analysis_id: str = "default",
) -> dict:
    """Find methods that allocate large byte arrays (potential payloads or shellcode buffers).

    Args:
        min_size: Minimum array size to report in bytes (default 256).
        analysis_id: Analysis session ID.

    Returns:
        ``{"min_size": int, "count": int, "matches": [{"class", "method", "offset", "size", "element_type"}, …]}``
        Results are sorted by size descending.
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    matches = []
    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            if not method.HasBody:
                continue
            instrs = list(method.Body.Instructions)
            for i, instr in enumerate(instrs):
                if str(instr.OpCode.Name) != "newarr":
                    continue
                operand = str(instr.Operand) if instr.Operand is not None else ""
                if "byte" not in operand.lower():
                    continue
                if i > 0 and str(instrs[i - 1].OpCode.Name).startswith("ldc.i4"):
                    try:
                        size = int(str(instrs[i - 1].Operand)) if instrs[i - 1].Operand is not None else 0
                    except (ValueError, TypeError):
                        size = 0
                    if size >= min_size:
                        matches.append({
                            "class": class_name,
                            "method": str(method.Name),
                            "offset": f"IL_{instr.Offset:04X}",
                            "size": size,
                            "element_type": operand,
                        })

    matches.sort(key=lambda x: x["size"], reverse=True)
    return {"min_size": min_size, "count": len(matches), "matches": matches}


def pythonnet_get_call_graph(
    class_name: str,
    method_name: str,
    max_depth: int = 5,
    analysis_id: str = "default",
) -> dict:
    """Build a call graph by BFS from a root method, following internal calls only.

    External framework calls are listed as edges but not expanded further.
    Returns edges as ``(caller, callee)`` pairs.

    Args:
        class_name: Root class to start from.
        method_name: Root method to start from.
        max_depth: Maximum BFS depth (default 5).
        analysis_id: Analysis session ID.

    Returns:
        ``{"root": str, "visited_count": int, "edge_count": int, "edges": [{"caller", "callee", "depth", "external"?}]}``
    """
    try:
        ctx = _get_context(analysis_id)
        ctx.get_type(class_name)  # validate
    except (ValueError, KeyError) as e:
        return {"error": str(e)}

    # Primary index: "ClassName::MethodName" -> (class_name, method)
    internal_methods: dict[str, tuple[str, Any]] = {}
    # Inverted index: method_name -> [(class_name, method)] for O(1) operand resolution
    by_method_name: dict[str, list[tuple[str, Any]]] = {}
    for cname, t in ctx._type_index.items():
        for m in t.Methods:
            mname = str(m.Name)
            internal_methods[f"{cname}::{mname}"] = (cname, m)
            by_method_name.setdefault(mname, []).append((cname, m))

    def resolve_internal(operand_str: str) -> tuple[str, Any] | tuple[None, None]:
        # For each method name that appears in the operand string, check its class
        for mname, candidates in by_method_name.items():
            if mname not in operand_str:
                continue
            for cname, m in candidates:
                if cname in operand_str:
                    return cname, m
        return None, None

    root_key = f"{class_name}::{method_name}"
    if root_key not in internal_methods:
        return {"error": f"Method not found: {class_name}.{method_name}"}

    edges: list[dict] = []
    visited: set[str] = set()
    queue = [(class_name, method_name, 0)]

    while queue:
        cur_class, cur_method, depth = queue.pop(0)
        node_key = f"{cur_class}::{cur_method}"
        if node_key in visited or depth > max_depth:
            continue
        visited.add(node_key)

        pair = internal_methods.get(node_key)
        if pair is None:
            continue
        _, m = pair
        if not m.HasBody:
            continue

        for instr in m.Body.Instructions:
            if str(instr.OpCode) not in _CALL_OPCODES or instr.Operand is None:
                continue
            operand_str = str(instr.Operand)
            callee_class, callee_method = resolve_internal(operand_str)

            if callee_class:
                callee_name = str(callee_method.Name)
                edges.append({
                    "caller": f"{cur_class}.{cur_method}",
                    "callee": f"{callee_class}.{callee_name}",
                    "depth": depth,
                })
                queue.append((callee_class, callee_name, depth + 1))
            else:
                edges.append({
                    "caller": f"{cur_class}.{cur_method}",
                    "callee": operand_str,
                    "depth": depth,
                    "external": True,
                })

    return {
        "root": f"{class_name}.{method_name}",
        "visited_count": len(visited),
        "edge_count": len(edges),
        "edges": edges,
    }


def pythonnet_find_dead_code(
    limit: int = 0,
    analysis_id: str = "default",
) -> dict:
    """Find methods that are never called from within the binary.

    Unreferenced methods are often artefacts of obfuscation, leftover debug
    code, or decoys. Well-known runtime entry points (``.ctor``, ``.cctor``,
    ``Main``, ``Finalize``, etc.) are excluded.

    Args:
        limit: Maximum number of results to return (0 = unlimited).
        analysis_id: Analysis session ID.

    Returns:
        ``{"count": int, "dead_methods": [{"class", "method", "full"}, …]}``
    """
    try:
        ctx = _get_context(analysis_id)
    except ValueError as e:
        return {"error": str(e)}

    referenced: set[str] = set()
    for t in ctx._type_index.values():
        for method in t.Methods:
            if not method.HasBody:
                continue
            for instr in method.Body.Instructions:
                if str(instr.OpCode) in _CALL_OPCODES_EXTENDED and instr.Operand is not None:
                    referenced.add(str(instr.Operand))

    dead = []
    for class_name, t in ctx._type_index.items():
        for method in t.Methods:
            if str(method.Name) in _RUNTIME_ENTRY_POINTS or not method.HasBody:
                continue
            full_sig = str(method)
            if not any(full_sig in ref or str(method.Name) in ref for ref in referenced):
                dead.append({
                    "class": class_name,
                    "method": str(method.Name),
                    "full": full_sig,
                })

    if limit:
        dead = dead[:limit]

    return {"count": len(dead), "dead_methods": dead}


def pythonnet_load_from_url(
    url: str,
    filename: str = "",
    analysis_id: str = "default",
) -> str:
    """Download a .NET binary from a URL and load it for analysis.

    The container fetches the file itself — no manual path needed. Useful when
    the binary is reachable from inside the Docker network (internal HTTP
    server, object storage, etc.).

    The downloaded file is staged in a temporary directory, then copied into
    ``/data/pythonnet_projects/{analysis_id}/`` as the working copy — same
    behaviour as ``pythonnet_load_binary``.

    Args:
        url: HTTP/HTTPS URL to download from.
        filename: Override the saved filename (default: derived from the URL).
        analysis_id: Session identifier.

    Returns:
        Same status string as ``pythonnet_load_binary``.
    """
    dest_name = filename.strip() or Path(url.split("?")[0]).name or "binary.bin"
    # Sanitise: no path traversal
    dest_name = Path(dest_name).name
    # /data/samples is mounted read-only in this container; use a writable staging area
    staging = PYTHONNET_PROJECTS_DIR / "_downloads"
    staging.mkdir(parents=True, exist_ok=True)
    dest = staging / dest_name

    try:
        logger.info(f"Downloading {url} -> {dest}")
        urllib.request.urlretrieve(url, dest)  # noqa: S310 – URL is caller-supplied
    except urllib.error.URLError as e:
        return f"Error downloading {url}: {e}"
    except Exception as e:
        return f"Error: {e}"

    return pythonnet_load_binary(str(dest), analysis_id)


def pythonnet_list_uploaded_files(extension: str = "") -> dict:
    """List .NET binaries available for analysis.

    Searches two directories:
    - ``/data/samples/`` — files uploaded via the main SARA upload endpoint
    - ``/data/pythonnet_projects/_uploads/`` — files uploaded directly to this server

    Args:
        extension: Filter by extension, e.g. ``".exe"`` or ``".dll"`` (default: all).

    Returns:
        ``{"count": int, "files": [{"name", "path", "size_bytes", "source"}]}``
    """
    ext = extension.strip().lower()
    files = []
    sources = [
        (SAMPLES_DIR, "samples"),
        (PYTHONNET_PROJECTS_DIR / "_uploads", "uploads"),
    ]
    for search_dir, source_label in sources:
        if not search_dir.exists():
            continue
        for p in sorted(search_dir.iterdir()):
            if not p.is_file():
                continue
            if ext and p.suffix.lower() != ext:
                continue
            files.append({
                "name": p.name,
                "path": str(p),
                "size_bytes": p.stat().st_size,
                "source": source_label,
            })

    return {"count": len(files), "files": files}


_STAGING_DIRS: frozenset[str] = frozenset({"_uploads", "_downloads"})


def pythonnet_list_active_sessions() -> dict:
    """List all analysis sessions, past and present.

    Combines in-memory sessions (currently loaded) with project directories
    found on disk under ``/data/pythonnet_projects/``. Staging directories
    (``_uploads``, ``_downloads``) are excluded.

    Returns:
        ``{"count": int, "sessions": [...]}``

        Active sessions: ``{"analysis_id", "status": "active", "binary", "types", "loaded_at", "last_accessed"}``

        On-disk sessions: ``{"analysis_id", "status": "on_disk", "binary", "files"}``
    """
    sessions = []

    # In-memory sessions (active)
    active_ids = set()
    for analysis_id, entry in _pythonnet_contexts.items():
        ctx: PythonnetContext = entry["context"]
        sessions.append({
            "analysis_id": analysis_id,
            "status": "active",
            "binary": str(ctx.binary_path),
            "types": len(ctx._type_index),
            "loaded_at": entry["loaded_at"].isoformat(),
            "last_accessed": entry["last_accessed"].isoformat(),
        })
        active_ids.add(analysis_id)

    # Past sessions on disk
    if PYTHONNET_PROJECTS_DIR.exists():
        for project_dir in sorted(PYTHONNET_PROJECTS_DIR.iterdir()):
            if not project_dir.is_dir():
                continue
            if project_dir.name in _STAGING_DIRS:
                continue
            if project_dir.name in active_ids:
                continue
            files = [p for p in project_dir.iterdir() if p.is_file()]
            sessions.append({
                "analysis_id": project_dir.name,
                "status": "on_disk",
                "binary": files[0].name if files else None,
                "files": len(files),
            })

    return {"count": len(sessions), "sessions": sessions}


def pythonnet_unload_binary(analysis_id: str = "default") -> dict:
    """Unload a binary and free its memory.

    Args:
        analysis_id: Session to close.

    Returns:
        ``{"analysis_id": str, "status": "unloaded"}``
    """
    if analysis_id not in _pythonnet_contexts:
        return {"error": f"No session '{analysis_id}' found"}
    _pythonnet_contexts[analysis_id]["context"].close()
    del _pythonnet_contexts[analysis_id]
    logger.info(f"Unloaded session '{analysis_id}'")
    return {"analysis_id": analysis_id, "status": "unloaded"}


# ---------------------------------------------------------------------------
# HTTP upload endpoint (mounted alongside the MCP app)
# ---------------------------------------------------------------------------

async def _download_endpoint(request: Request) -> Response:
    """GET /download?path=/absolute/path/to/file — download a file from the container.

    Intended for retrieving patched binaries after a ``pythonnet_save_binary``
    call. Only files whose resolved path starts with an allowed prefix are
    served (no arbitrary filesystem reads).

    Example::

        curl -OJ "http://docker-host:8001/download?path=/data/pythonnet_projects/my-analysis/patched.dll"
    """
    from starlette.responses import FileResponse, Response as StarletteResponse

    raw_path = request.query_params.get("path", "").strip()
    if not raw_path:
        return JSONResponse({"error": "Missing 'path' query parameter"}, status_code=400)

    target = Path(raw_path).resolve()

    # Restrict to project output directory only
    if not str(target).startswith(str(PYTHONNET_PROJECTS_DIR.resolve())):
        return JSONResponse(
            {"error": f"Path not in an allowed directory: {target}"},
            status_code=403,
        )

    if not target.exists():
        return JSONResponse({"error": f"File not found: {target}"}, status_code=404)
    if not target.is_file():
        return JSONResponse({"error": f"Not a file: {target}"}, status_code=400)

    logger.info(f"Serving download: {target}")
    return FileResponse(
        path=str(target),
        filename=target.name,
        media_type="application/octet-stream",
    )


async def _upload_endpoint(request: Request) -> JSONResponse:
    """POST /upload — upload a binary file into the container.

    Accepts ``multipart/form-data`` with a ``file`` field.
    Requires ``python-multipart`` to be installed.

    Example::

        curl -F file=@malware.exe http://docker-host:8001/upload

    Returns ``{"path": "/data/samples/malware.exe", "size_bytes": N}``.
    """
    try:
        form = await request.form()
    except Exception as e:
        return JSONResponse({"error": f"Failed to parse form: {e}"}, status_code=400)

    upload = form.get("file")
    if upload is None:
        return JSONResponse({"error": "Missing 'file' field in form data"}, status_code=400)

    # Sanitise filename — no path traversal
    raw_name = getattr(upload, "filename", None) or "upload.bin"
    safe_name = Path(raw_name).name or "upload.bin"
    # /data/samples is mounted read-only in this container; save to a writable staging area
    staging = PYTHONNET_PROJECTS_DIR / "_uploads"
    staging.mkdir(parents=True, exist_ok=True)
    dest = staging / safe_name

    try:
        content = await upload.read()
        dest.write_bytes(content)
        logger.info(f"Uploaded {safe_name} ({len(content)} bytes) -> {dest}")
        return JSONResponse({"path": str(dest), "size_bytes": len(content)})
    except Exception as e:
        logger.error(f"Upload error: {e}", exc_info=True)
        return JSONResponse({"error": str(e)}, status_code=500)


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

@click.command()
@click.option('--binary', type=click.Path(exists=True), required=False,
              help='Binary to analyse (optional — can also load via pythonnet_load_binary)')
@click.option('--host', default='0.0.0.0', help='Bind host')
@click.option('--port', default=8001, type=int, help='Bind port')
def main(
    binary: str | None,
    host: str,
    port: int,
) -> None:
    """Headless Pythonnet MCP server for .NET binary analysis."""
    try:
        logger.info("Pythonnet initialized")
        logger.info(f"Creating MCP server — host={host}, port={port}")

        mcp = FastMCP(
            "pythonnet-mcp-headless",
            host=host,
            port=port,
            stateless_http=True,
            json_response=True,
        )

        tools = [
            pythonnet_load_binary,
            pythonnet_load_from_url,
            pythonnet_list_uploaded_files,
            pythonnet_list_all_methods,
            pythonnet_list_all_methods_inside_class,
            pythonnet_list_all_classes,
            pythonnet_decompile_class,
            pythonnet_decompile_method,
            pythonnet_decompile_all_classes,
            pythonnet_rename_class,
            pythonnet_rename_method,
            pythonnet_rename_field,
            pythonnet_save_binary,
            pythonnet_get_strings,
            pythonnet_search_string,
            pythonnet_search_method_by_name,
            pythonnet_get_imports,
            pythonnet_get_fields,
            pythonnet_get_custom_attributes,
            pythonnet_get_opcodes,
            pythonnet_get_method_calls,
            pythonnet_get_callers,
            pythonnet_get_call_graph,
            pythonnet_get_entry_point,
            pythonnet_get_module_info,
            pythonnet_get_resources,
            pythonnet_extract_resource,
            pythonnet_find_pinvoke,
            pythonnet_find_reflection_calls,
            pythonnet_find_large_byte_arrays,
            pythonnet_find_dead_code,
            pythonnet_patch_ldstr,
            pythonnet_nop_instructions,
            pythonnet_patch_return,
            pythonnet_patch_branch,
            pythonnet_set_field_constant,
            pythonnet_list_active_sessions,
            pythonnet_unload_binary,
        ]
        for tool in tools:
            mcp.tool()(tool)

        # Combine MCP app with the /upload HTTP endpoint on the same port.
        # The upload route is listed first so it takes priority over the MCP
        # catch-all mount.
        mcp_asgi = mcp.streamable_http_app()

        @asynccontextmanager
        async def lifespan(app):
            logger.info("Starting MCP lifespan (task group init)…")
            try:
                async with mcp_asgi.router.lifespan_context(app):
                    logger.info("MCP lifespan ready")
                    yield
            except Exception as e:
                logger.error(f"MCP lifespan error: {e}\n{traceback.format_exc()}")
                raise
            finally:
                logger.info("MCP lifespan ended")

        async def _asgi_exception_handler(request: Request, exc: Exception) -> JSONResponse:
            logger.error(
                f"Unhandled ASGI exception on {request.method} {request.url.path}: "
                f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}"
            )
            return JSONResponse({"error": str(exc)}, status_code=500)

        app = Starlette(
            lifespan=lifespan,
            routes=[
                Route("/upload", _upload_endpoint, methods=["POST"]),
                Route("/download", _download_endpoint, methods=["GET"]),
                Mount("/", mcp_asgi),
            ],
            exception_handlers={Exception: _asgi_exception_handler},
        )

        logger.info(
            f"Starting server on {host}:{port} "
            f"({len(tools)} MCP tools · POST /upload available)"
            f"({len(tools)} MCP tools · GET /download available)"
        )
        uvicorn.run(app, host=host, port=port, log_level="warning")

    except KeyboardInterrupt:
        logger.info("Shutting down…")
    except Exception as e:
        logger.error(f"Error: {e}", exc_info=True)
        print(e)
        raise SystemExit(1)


if __name__ == "__main__":
    main()