PCM

import os import sys import ast import json import traceback import shutil import http.client import signal import asyncio from typing import List, Optional, Dict, Any from fastmcp import FastMCP # The log_level is necessary for Cline to work: https://github.com/jlowin/fastmcp/issues/81 mcp = FastMCP("IDA Pro", log_level="ERROR") jsonrpc_request_id = 1 def handle_shutdown(sig, frame): print("[*] Shutting down MCP server gracefully...") if hasattr(handle_shutdown, "_is_shutting_down") and handle_shutdown._is_shutting_down: return handle_shutdown._is_shutting_down = True if hasattr(asyncio, "get_running_loop"): try: loop = asyncio.get_running_loop() for task in asyncio.all_tasks(loop): task.cancel() print("[*] All tasks cancelled, shutdown complete.") except RuntimeError: pass import os os._exit(0) def make_jsonrpc_request(method: str, *args): """Make a JSON-RPC request to the IDA plugin""" global jsonrpc_request_id conn = http.client.HTTPConnection("localhost", 13337) request = { "jsonrpc": "2.0", "method": method, "params": list(args), "id": jsonrpc_request_id, } jsonrpc_request_id += 1 try: conn.request("POST", "/mcp", json.dumps(request), {"Content-Type": "application/json"}) response = conn.getresponse() data = json.loads(response.read().decode()) if "error" in data: error = data["error"] error_message = f"JSON-RPC error {error['code']}: {error['message']}" if "data" in error and error["data"]: error_message += f"\nDetails: {error['data']}" raise Exception(error_message) return data.get("result") except Exception as e: traceback.print_exc() raise finally: conn.close() class MCPVisitor(ast.NodeVisitor): def __init__(self): self.types: dict[str, ast.ClassDef] = {} self.functions: dict[str, ast.FunctionDef] = {} self.descriptions: dict[str, str] = {} def visit_FunctionDef(self, node): for decorator in node.decorator_list: if isinstance(decorator, ast.Name): if decorator.id == "jsonrpc": for i, arg in enumerate(node.args.args): arg_name = arg.arg arg_type = arg.annotation if arg_type is None: raise Exception(f"Missing argument type for {node.name}.{arg_name}") if isinstance(arg_type, ast.Subscript): assert isinstance(arg_type.value, ast.Name) assert arg_type.value.id == "Annotated" assert isinstance(arg_type.slice, ast.Tuple) assert len(arg_type.slice.elts) == 2 annot_type = arg_type.slice.elts[0] annot_description = arg_type.slice.elts[1] assert isinstance(annot_description, ast.Constant) node.args.args[i].annotation = ast.Subscript( value=ast.Name(id="Annotated", ctx=ast.Load()), slice=ast.Tuple( elts=[annot_type, ast.Call(func=ast.Name(id="Field", ctx=ast.Load()), args=[], keywords=[ast.keyword(arg="description", value=annot_description)])], ctx=ast.Load() ), ctx=ast.Load(), ) elif isinstance(arg_type, ast.Name): pass else: raise Exception(f"Unexpected type annotation for {node.name}.{arg_name} -> {type(arg_type)}") body_comment = node.body[0] if isinstance(body_comment, ast.Expr) and isinstance(body_comment.value, ast.Constant): new_body = [body_comment] self.descriptions[node.name] = body_comment.value.value else: new_body = [] call_args = [ast.Constant(value=node.name)] for arg in node.args.args: call_args.append(ast.Name(id=arg.arg, ctx=ast.Load())) new_body.append(ast.Return(value=ast.Call(func=ast.Name(id="make_jsonrpc_request", ctx=ast.Load()), args=call_args, keywords=[]))) decorator_list = [ast.Call(func=ast.Attribute(value=ast.Name(id="mcp", ctx=ast.Load()), attr="tool", ctx=ast.Load()), args=[], keywords=[])] node_nobody = ast.FunctionDef(node.name, node.args, new_body, decorator_list, node.returns, node.type_comment, lineno=node.lineno, col_offset=node.col_offset) self.functions[node.name] = node_nobody def visit_ClassDef(self, node): for base in node.bases: if isinstance(base, ast.Name): if base.id == "TypedDict": self.types[node.name] = node SCRIPT_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'src') IDA_PLUGIN_PY = os.path.join(SCRIPT_DIR, 'ida', "mcp.py") GENERATED_PY = os.path.join(SCRIPT_DIR, "server_generated.py") # NOTE: This is in the global scope on purpose with open(IDA_PLUGIN_PY, "r") as f: code = f.read() module = ast.parse(code, IDA_PLUGIN_PY) visitor = MCPVisitor() visitor.visit(module) code = """# NOTE: This file has been automatically generated, do not modify! from typing import Annotated, Optional, TypedDict, List from pydantic import Field """ for type in visitor.types.values(): code += ast.unparse(type) code += "\n\n" for function in visitor.functions.values(): code += ast.unparse(function) code += "\n\n" with open(GENERATED_PY, "w") as f: f.write(code) exec(compile(code, GENERATED_PY, "exec")) def main(): import argparse parser = argparse.ArgumentParser(description="MCP server for IDA Pro") parser.add_argument("--generate-only", action="store_true", help="Generate the IDA plugin code and exit") parser.add_argument("--install-plugin", action="store_true", help="Install the IDA plugin") args = parser.parse_args() if args.generate_only: print(f"[*] Generating IDA plugin code...", file=sys.stderr) for function in visitor.functions.values(): signature = function.name + "(" for i, arg in enumerate(function.args.args): if i > 0: signature += ", " signature += arg.arg signature += ")" description = visitor.descriptions.get(function.name, "<no description>") if description[-1] != ".": description += "." print(f"- `{signature}`: {description}") sys.exit(0) elif args.install_plugin: print(f"[*] Installing IDA plugin...", file=sys.stderr) if sys.platform == "win32": ida_plugin_folder = os.path.join(os.getenv("APPDATA"), "Hex-Rays", "IDA Pro", "plugins") else: ida_plugin_folder = os.path.join(os.path.expanduser("~"), ".idapro", "plugins") plugin_destination = os.path.join(ida_plugin_folder, "pcm.py") if input(f"Installing IDA plugin to {plugin_destination}, proceed? [Y/n] ").lower() == "n": sys.exit(1) if not os.path.exists(ida_plugin_folder): os.makedirs(ida_plugin_folder) shutil.copy(IDA_PLUGIN_PY, plugin_destination) print(f"Installed plugin: {plugin_destination}") sys.exit(0) signal.signal(signal.SIGINT, handle_shutdown) signal.signal(signal.SIGTERM, handle_shutdown) print('[*] Starting MCP server...') try: mcp.run(transport="stdio") except KeyboardInterrupt: print("[*] Received keyboard interrupt, shutting down...") except Exception as e: print(f"[!] Error: {e}") traceback.print_exc() finally: print("[*] MCP server stopped.") if __name__ == "__main__": main()