import sys
import signal
import inspect
import logging
import argparse
import importlib
from pathlib import Path
from typing import Annotated, Optional, Any
import typing_inspection.introspection as intro
from mcp.server.fastmcp import FastMCP
# ida/idapro must go first to initialize idalib
# IDA 9.0+ uses 'ida' module, IDA 8.x uses 'idapro'
try:
import ida as idapro # IDA 9.0+
except ImportError:
import idapro # IDA 8.x
import ida_auto
import ida_hexrays
logger = logging.getLogger(__name__)
mcp = FastMCP("github.com/icryo/ida-pro-mcp#idalib")
# Global state for database management
_db_state = {
"loaded": False,
"path": None,
"tools_registered": False,
"unsafe_mode": False,
}
def _wrap_with_db_check(func):
"""Wrap a function to check if database is loaded before calling"""
import functools
@functools.wraps(func)
def wrapper(*args, **kwargs):
if not _db_state["loaded"]:
raise RuntimeError("No database is currently loaded. Use load_database tool first.")
return func(*args, **kwargs)
return wrapper
def _register_ida_tools():
"""Register IDA tools from mcp-plugin.py - must be called at startup"""
if _db_state["tools_registered"]:
logger.debug("IDA tools already registered")
return
plugin = importlib.import_module("ida_pro_mcp.mcp-plugin")
logger.debug("adding IDA tools...")
for name, callable in plugin.rpc_registry.methods.items():
if _db_state["unsafe_mode"] or name not in plugin.rpc_registry.unsafe:
logger.debug("adding tool: %s: %s", name, callable)
# Wrap the callable to check database state
wrapped = _wrap_with_db_check(callable)
mcp.add_tool(wrapped, name)
_db_state["tools_registered"] = True
fixup_tool_argument_descriptions(mcp)
logger.info("IDA tools registered successfully")
@mcp.tool()
def load_database(
path: Annotated[str, "Path to the binary file or IDB to analyze"],
run_auto_analysis: Annotated[bool, "Run automatic analysis after loading"] = True
) -> dict[str, Any]:
"""Load a new binary file or IDB database for analysis. Closes any currently open database first."""
global _db_state
input_path = Path(path)
if not input_path.exists():
return {"success": False, "error": f"File not found: {path}"}
# Close existing database if one is open
if _db_state["loaded"]:
logger.info("Closing current database: %s", _db_state["path"])
idapro.close_database()
_db_state["loaded"] = False
_db_state["path"] = None
# Open new database
logger.info("Opening database: %s", input_path)
if idapro.open_database(str(input_path), run_auto_analysis=run_auto_analysis):
return {"success": False, "error": f"Failed to open database: {path}"}
if run_auto_analysis:
logger.debug("Waiting for auto-analysis to complete...")
ida_auto.auto_wait()
# Initialize Hex-Rays if not already done
if not ida_hexrays.init_hexrays_plugin():
logger.warning("Failed to initialize Hex-Rays decompiler - decompilation will not be available")
_db_state["loaded"] = True
_db_state["path"] = str(input_path.absolute())
return {
"success": True,
"path": _db_state["path"],
"message": f"Successfully loaded {input_path.name}"
}
@mcp.tool()
def close_database() -> dict[str, Any]:
"""Close the currently open database."""
global _db_state
if not _db_state["loaded"]:
return {"success": False, "error": "No database is currently open"}
logger.info("Closing database: %s", _db_state["path"])
idapro.close_database()
old_path = _db_state["path"]
_db_state["loaded"] = False
_db_state["path"] = None
return {
"success": True,
"message": f"Successfully closed database: {old_path}"
}
@mcp.tool()
def save_database(
path: Annotated[Optional[str], "Optional path to save the database to. If not provided, saves to the current IDB path."] = None
) -> dict[str, Any]:
"""
Save the current database to disk. This preserves all analysis including:
- Function names and boundaries
- Type definitions and assignments
- Comments (all types)
- Structures and enums
- Cross-references
- Everything else in the IDB
If no path is provided, saves to the original location (creating .idb/.i64 file).
If a path is provided, saves to that location (useful for creating backups or snapshots).
"""
global _db_state
if not _db_state["loaded"]:
return {"success": False, "error": "No database is currently open"}
import idc
if path:
save_path = Path(path)
# Ensure it has proper extension
if save_path.suffix not in [".idb", ".i64"]:
# Add appropriate extension based on bitness
import ida_ida
if ida_ida.inf_is_64bit():
save_path = save_path.with_suffix(".i64")
else:
save_path = save_path.with_suffix(".idb")
logger.info("Saving database to: %s", save_path)
# Use idc.save_database with explicit path
success = idc.save_database(str(save_path), idc.DBFL_COMP)
saved_to = str(save_path.absolute())
else:
logger.info("Saving database to current location")
# Save to current location
success = idc.save_database("", idc.DBFL_COMP)
saved_to = _db_state["path"]
if success:
return {
"success": True,
"path": saved_to,
"message": f"Successfully saved database to {saved_to}"
}
else:
return {
"success": False,
"error": "Failed to save database"
}
@mcp.tool()
def get_database_status() -> dict[str, Any]:
"""Get the current database status - whether a database is loaded and its path."""
return {
"loaded": _db_state["loaded"],
"path": _db_state["path"],
"tools_available": _db_state["tools_registered"]
}
def fixup_tool_argument_descriptions(mcp: FastMCP):
# In our tool definitions within `mcp-plugin.py`, we use `typing.Annotated` on function parameters
# to attach documentation. For example:
#
# def get_function_by_name(
# name: Annotated[str, "Name of the function to get"]
# ) -> Function:
# """Get a function by its name"""
# ...
#
# However, the interpretation of Annotated is left up to static analyzers and other tools.
# FastMCP doesn't have any special handling for these comments, so we splice them into the
# tool metadata ourselves here.
#
# Example, before:
#
# tool.parameter={
# properties: {
# name: {
# title: "Name",
# type: "string"
# }
# },
# required: ["name"],
# title: "get_function_by_nameArguments",
# type: "object"
# }
#
# Example, after:
#
# tool.parameter={
# properties: {
# name: {
# title: "Name",
# type: "string"
# description: "Name of the function to get"
# }
# },
# required: ["name"],
# title: "get_function_by_nameArguments",
# type: "object"
# }
#
# References:
# - https://docs.python.org/3/library/typing.html#typing.Annotated
# - https://fastapi.tiangolo.com/python-types/#type-hints-with-metadata-annotations
# unfortunately, FastMCP.list_tools() is async, so we break with best practices and reach into `._tool_manager`
# rather than spinning up an asyncio runtime just to fetch the (non-async) list of tools.
for tool in mcp._tool_manager.list_tools():
sig = inspect.signature(tool.fn)
for name, parameter in sig.parameters.items():
# this instance is a raw `typing._AnnotatedAlias` that we can't do anything with directly.
# it renders like:
#
# typing.Annotated[str, 'Name of the function to get']
if not parameter.annotation:
continue
# this instance will look something like:
#
# InspectedAnnotation(type=<class 'str'>, qualifiers=set(), metadata=['Name of the function to get'])
#
annotation = intro.inspect_annotation(
parameter.annotation,
annotation_source=intro.AnnotationSource.ANY
)
# for our use case, where we attach a single string annotation that is meant as documentation,
# we extract that string and assign it to "description" in the tool metadata.
if annotation.type is not str:
continue
if len(annotation.metadata) != 1:
continue
description = annotation.metadata[0]
if not isinstance(description, str):
continue
logger.debug("adding parameter documentation %s(%s='%s')", tool.name, name, description)
tool.parameters["properties"][name]["description"] = description
def main():
global _db_state
parser = argparse.ArgumentParser(description="MCP server for IDA Pro via idalib")
parser.add_argument("--verbose", "-v", action="store_true", help="Show debug messages")
parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to listen on, default: 127.0.0.1")
parser.add_argument("--port", type=int, default=8745, help="Port to listen on, default: 8745")
parser.add_argument("--unsafe", action="store_true", help="Enable unsafe functions (DANGEROUS)")
parser.add_argument("input_path", type=Path, nargs="?", default=None, help="Optional path to the input file to analyze. If not provided, use load_database tool to load files.")
args = parser.parse_args()
if args.verbose:
log_level = logging.DEBUG
if hasattr(idapro, 'enable_console_messages'):
idapro.enable_console_messages(True)
else:
log_level = logging.INFO
if hasattr(idapro, 'enable_console_messages'):
idapro.enable_console_messages(False)
mcp.settings.log_level = logging.getLevelName(log_level)
mcp.settings.host = args.host
mcp.settings.port = args.port
logging.basicConfig(level=log_level)
# reset logging levels that might be initialized in idapythonrc.py
# which is evaluated during import of idalib.
logging.getLogger().setLevel(log_level)
# Store unsafe mode for later tool registration
_db_state["unsafe_mode"] = args.unsafe
# If input_path provided, load it immediately (backwards compatible)
if args.input_path is not None:
if not args.input_path.exists():
raise FileNotFoundError(f"Input file not found: {args.input_path}")
logger.info("opening database: %s", args.input_path)
if idapro.open_database(str(args.input_path), run_auto_analysis=True):
raise RuntimeError("failed to analyze input file")
logger.debug("idalib: waiting for analysis...")
ida_auto.auto_wait()
if not ida_hexrays.init_hexrays_plugin():
raise RuntimeError("failed to initialize Hex-Rays decompiler")
_db_state["loaded"] = True
_db_state["path"] = str(args.input_path.absolute())
else:
logger.info("Starting without a database. Use load_database tool to load a binary.")
# Register IDA tools at startup (MCP doesn't support dynamic tool registration)
# Tools will check if database is loaded before executing
_register_ida_tools()
# Setup signal handlers to ensure IDA database is properly closed on shutdown.
#
# PROBLEM: When uvicorn (used by FastMCP's SSE transport) receives SIGINT/SIGTERM:
# 1. It captures the signal and performs graceful shutdown
# 2. After shutdown, it re-raises the signal with the default handler
# 3. The default handler immediately terminates the process at the OS level
# 4. This bypasses all remaining Python code (try/except/finally blocks)
#
# SOLUTION: Register our signal handlers BEFORE calling mcp.run(). When a signal
# arrives, our handlers execute first, allowing us to close the IDA database
# cleanly before the process terminates.
def cleanup_and_exit(signum, frame):
if _db_state["loaded"]:
logger.info("Closing IDA database...")
idapro.close_database()
logger.info("IDA database closed.")
sys.exit(0)
signal.signal(signal.SIGINT, cleanup_and_exit)
signal.signal(signal.SIGTERM, cleanup_and_exit)
# NOTE: npx @modelcontextprotocol/inspector for debugging
logger.info("MCP Server available at: http://%s:%d/sse", mcp.settings.host, mcp.settings.port)
mcp.run(transport="sse")
if __name__ == "__main__":
main()
