import json import os import secrets import time import httpx from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import JSONResponse, Response, HTMLResponse from starlette.routing import Route SSE_PORT = int(os.environ.get("MCP_PORT", 8888)) NGROK_URL = os.environ.get("NGROK_URL", "") AUTH_TOKEN = os.environ.get("IDA_MCP_TOKEN") INSTANCE_TTL = int(os.environ.get("IDA_MCP_INSTANCE_TTL", 30)) http = httpx.AsyncClient(timeout=120.0) INSTANCES = {} AUTH_CODES = {} AUTH_CODE_TTL = 120 def _live_instances(): now = time.time() return {k: v for k, v in INSTANCES.items() if now - v["last_seen"] <= INSTANCE_TTL} def _check_token(request: Request) -> bool: if not AUTH_TOKEN: return True return request.headers.get("Authorization", "") == f"Bearer {AUTH_TOKEN}" class ApiError(Exception): def __init__(self, message: str, code: str = "error"): super().__init__(message) self.code = code class ResolveError(ApiError): pass def resolve_instance(instance_id): live = _live_instances() if instance_id: if instance_id not in live: available = ", ".join(live.keys()) or "(none)" raise ResolveError( f"Instance '{instance_id}' not found or not responding. Available: {available}. " f"Call ida_list_instances to see what's currently open.", code="instance_not_found", ) v = live[instance_id] return v["host"], v["port"], instance_id, False if not live: raise ResolveError( "No IDA instances registered. Is the plugin running, and IDA_MCP_REGISTRY_URL " f"pointing at this server (http://127.0.0.1:{SSE_PORT})?", code="no_instances", ) if len(live) == 1: (only_id, v), = live.items() return v["host"], v["port"], only_id, False chosen_id, chosen = max(live.items(), key=lambda kv: kv[1]["last_seen"]) return chosen["host"], chosen["port"], chosen_id, True def _base(host, port): return f"http://{host}:{port}" def _unwrap(r: httpx.Response): if r.status_code >= 400: code = "internal_error" if r.status_code >= 500 else "bad_request" try: body = r.json() msg = body.get("error", r.text) code = body.get("code", code) except Exception: msg = r.text raise ApiError(f"IDA plugin error ({r.status_code}): {msg}", code=code) return r.json() async def ida_get(host, port, path): headers = {"Authorization": f"Bearer {AUTH_TOKEN}"} if AUTH_TOKEN else {} r = await http.get(f"{_base(host, port)}{path}", headers=headers) return _unwrap(r) async def ida_post(host, port, path, data): headers = {"Authorization": f"Bearer {AUTH_TOKEN}"} if AUTH_TOKEN else {} r = await http.post(f"{_base(host, port)}{path}", json=data, headers=headers) return _unwrap(r) _INSTANCE_PARAM = { "instance": { "type": "string", "description": "IDA instance id (see ida_list_instances). Omit if only one IDA is open.", } } def _schema(props=None, required=None): p = dict(_INSTANCE_PARAM) p.update(props or {}) return {"type": "object", "properties": p, "required": required or []} TOOLS = [ { "name": "ida_list_instances", "description": ( "List currently open IDA instances (auto-registered by the plugin). Call this first " "when unsure which binary/instance to target, or when a tool call errors about multiple " "instances being open." ), "inputSchema": {"type": "object", "properties": {}}, }, { "name": "ida_ping", "description": "Check connectivity to an IDA instance; reports whether Hex-Rays decompiler is available.", "inputSchema": _schema(), }, { "name": "ida_list_functions", "description": "List functions with pagination. ALWAYS pass filter to avoid dumping huge lists into context.", "inputSchema": _schema( { "filter": {"type": "string", "description": "Substring filter on function name — set this!"}, "limit": {"type": "integer", "default": 100}, "offset": {"type": "integer", "default": 0}, } ), }, { "name": "ida_search_functions", "description": "Search functions by substring in name. Prefer this over ida_list_functions for lookups.", "inputSchema": _schema({"q": {"type": "string"}, "limit": {"type": "integer", "default": 50}}, ["q"]), }, { "name": "ida_search_strings", "description": "Search decoded strings by substring.", "inputSchema": _schema({"q": {"type": "string"}, "limit": {"type": "integer", "default": 30}}, ["q"]), }, { "name": "ida_get_name", "description": "Resolve a symbol name (or 0xADDR) to its value — string contents or raw dword.", "inputSchema": _schema({"name": {"type": "string"}}, ["name"]), }, { "name": "ida_get_imports", "description": "List the import table, optionally filtered by substring.", "inputSchema": _schema({"filter": {"type": "string"}}), }, { "name": "ida_get_segments", "description": "List segments (name, address range, size, r/w/x permissions) — useful for orienting in an unfamiliar binary.", "inputSchema": _schema(), }, { "name": "ida_get_entry_points", "description": "List entry points of the binary.", "inputSchema": _schema(), }, { "name": "ida_decompile", "description": "Decompile a single function (Hex-Rays) by name or address (0x...).", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_decompile_many", "description": ( "Decompile several functions in one call (single round-trip instead of N) — use this instead " "of calling ida_decompile in a loop when you already know the list of targets, e.g. after xrefs." ), "inputSchema": _schema( {"targets": {"type": "array", "items": {"type": "string"}, "description": "Names or 0x addresses"}}, ["targets"], ), }, { "name": "ida_get_disasm", "description": ( "Raw disassembly listing for a function — use when Hex-Rays fails or asm-level detail is needed. " "Response includes total_instructions/returned/truncated — if truncated is true, call again with a " "higher count or a nonzero offset to page through the rest instead of assuming you've seen it all." ), "inputSchema": _schema( { "target": {"type": "string"}, "count": {"type": "integer", "default": 30}, "offset": {"type": "integer", "default": 0, "description": "Skip this many instructions before returning — for paging past a truncated result"}, }, ["target"], ), }, { "name": "ida_xrefs_to", "description": "Find all locations that reference an address/function (callers, data refs).", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_xrefs_from", "description": "Find all addresses a function references (callees).", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_get_prototype", "description": "Get the current type/prototype string of a function.", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_set_prototype", "description": 'Set a function\'s prototype, e.g. "int __fastcall foo(int a, char *b);".', "inputSchema": _schema({"target": {"type": "string"}, "prototype": {"type": "string"}}, ["target", "prototype"]), }, { "name": "ida_set_lvar_type", "description": ( 'Set the type of a local variable inside a decompiled function, e.g. type="int" or "MyStruct *". ' "Experimental — Hex-Rays API varies slightly across IDA versions; check the error message if it fails." ), "inputSchema": _schema( {"function": {"type": "string"}, "lvar_name": {"type": "string"}, "type": {"type": "string"}}, ["function", "lvar_name", "type"], ), }, { "name": "ida_rename_function", "description": "Rename a function.", "inputSchema": _schema({"target": {"type": "string"}, "new_name": {"type": "string"}}, ["target", "new_name"]), }, { "name": "ida_rename_local", "description": "Rename a local variable inside a decompiled function.", "inputSchema": _schema( {"function": {"type": "string"}, "old_name": {"type": "string"}, "new_name": {"type": "string"}}, ["function", "old_name", "new_name"], ), }, { "name": "ida_set_comment", "description": "Set a function-level comment.", "inputSchema": _schema({"target": {"type": "string"}, "comment": {"type": "string"}}, ["target", "comment"]), }, { "name": "ida_get_notes", "description": ( "Read the persistent AI scratchpad note stored at an address (survives across sessions/agents, " "stored inside the idb). Use this to check if a function was already analyzed before re-doing work." ), "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_set_notes", "description": ( "Write a persistent AI scratchpad note at an address — e.g. your analysis conclusion for a " "function, so future calls (or other agents) don't have to re-derive it. Empty string deletes it." ), "inputSchema": _schema({"target": {"type": "string"}, "note": {"type": "string"}}, ["target", "note"]), }, { "name": "ida_rebuild_name_cache", "description": "Force-rebuild the plugin's internal name cache (rarely needed; use if get_name misses a symbol that clearly exists).", "inputSchema": _schema(), }, { "name": "ida_get_bytes", "description": "Read raw bytes at an address as hex — for shellcode, keys, magic constants, or anything not worth disassembling.", "inputSchema": _schema({"target": {"type": "string"}, "size": {"type": "integer", "default": 64}}, ["target"]), }, { "name": "ida_find_pattern", "description": ( 'Search the whole binary for an IDA-style hex byte pattern, e.g. "48 8B ?? ??" (?? = wildcard byte). ' "Experimental — underlying IDA search API varies across versions." ), "inputSchema": _schema({"pattern": {"type": "string"}, "limit": {"type": "integer", "default": 50}}, ["pattern"]), }, { "name": "ida_patch_bytes", "description": "Overwrite raw bytes at an address in the binary. Destructive — confirm with the user before calling.", "inputSchema": _schema({"target": {"type": "string"}, "hex_bytes": {"type": "string", "description": "e.g. '90 90 90' for 3 NOPs"}}, ["target", "hex_bytes"]), }, { "name": "ida_make_data", "description": "Define a data item (byte/word/dword/qword/struct) at an address. Experimental; struct requires struct_name from ida_list_local_types.", "inputSchema": _schema( { "target": {"type": "string"}, "data_type": {"type": "string", "enum": ["byte", "word", "dword", "qword", "struct"]}, "struct_name": {"type": "string"}, }, ["target", "data_type"], ), }, { "name": "ida_get_line_comment", "description": "Read the comment at a specific address (disassembly-line level, distinct from the whole-function comment).", "inputSchema": _schema({"target": {"type": "string"}, "repeatable": {"type": "boolean", "default": False}}, ["target"]), }, { "name": "ida_set_line_comment", "description": "Set a comment at a specific address (disassembly-line level, distinct from ida_set_comment which is whole-function).", "inputSchema": _schema( {"target": {"type": "string"}, "comment": {"type": "string"}, "repeatable": {"type": "boolean", "default": False}}, ["target", "comment"], ), }, { "name": "ida_top_functions", "description": ( "Triage tool for an unfamiliar/stripped binary: list the most cross-referenced or largest functions " "first, instead of guessing where to start." ), "inputSchema": _schema({"by": {"type": "string", "enum": ["xrefs", "size"], "default": "xrefs"}, "limit": {"type": "integer", "default": 20}}), }, { "name": "ida_call_graph", "description": "Build a callers or callees graph up to N levels deep from a function, in one call instead of walking xrefs manually.", "inputSchema": _schema( { "target": {"type": "string"}, "depth": {"type": "integer", "default": 2}, "direction": {"type": "string", "enum": ["callees", "callers"], "default": "callees"}, }, ["target"], ), }, { "name": "ida_make_function", "description": "Create a function at an address that auto-analysis missed.", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_undefine", "description": "Remove function/data definition at an address (undo a bad auto-analysis guess before redefining it).", "inputSchema": _schema({"target": {"type": "string"}}, ["target"]), }, { "name": "ida_list_local_types", "description": "List structs/enums currently defined in Local Types.", "inputSchema": _schema({"limit": {"type": "integer", "default": 200}}), }, { "name": "ida_create_struct", "description": 'Define a new struct in Local Types from a field list, e.g. fields=[{"name":"flag","type":"int"},{"name":"ptr","type":"void *"}].', "inputSchema": _schema( { "name": {"type": "string"}, "fields": {"type": "array", "items": {"type": "object", "properties": {"name": {"type": "string"}, "type": {"type": "string"}}}}, }, ["name", "fields"], ), }, { "name": "ida_get_objc_classes", "description": ( "List Objective-C classes and methods recovered from the binary (by _OBJC_CLASS_$_ symbols and " "+[Class sel]/-[Class sel] function naming). Use this first on any Mach-O/iOS/macOS target with ObjC " "— it's often more useful than decompiling method-by-method." ), "inputSchema": _schema({"filter": {"type": "string"}, "limit": {"type": "integer", "default": 200}}), }, { "name": "ida_find_immediate", "description": ( "Search all instructions for an immediate operand matching a given value (decimal or 0x-hex) — " "e.g. find where a magic constant, size, or flag literal is used, unlike ida_find_pattern which " "matches raw bytes regardless of instruction semantics." ), "inputSchema": _schema({"value": {"type": "string", "description": "e.g. '1337' or '0x539'"}, "limit": {"type": "integer", "default": 100}}, ["value"]), }, { "name": "ida_find_vtables", "description": ( "Heuristic search for virtual-function-table-like arrays: runs of consecutive pointers in " "non-executable segments where every pointer is the start of a real function. Useful for finding " "C++ vtables. Experimental — heuristic, can have false positives/negatives; narrow with 'segment' " "(e.g. '__const', '__data') on large binaries for speed." ), "inputSchema": _schema( { "segment": {"type": "string", "description": "Optional segment name substring to narrow the scan"}, "min_entries": {"type": "integer", "default": 2}, "limit": {"type": "integer", "default": 50}, } ), }, { "name": "ida_create_local_type", "description": ( 'Define a raw C declaration (struct/enum/union/typedef) in Local Types, e.g. decl="enum Color { RED, GREEN, BLUE };". ' "More general than ida_create_struct — use this for enums/unions/typedefs, or structs with bitfields/nesting." ), "inputSchema": _schema({"decl": {"type": "string"}}, ["decl"]), }, { "name": "ida_xrefs_to_many", "description": "Get xrefs_to for several targets in one call — e.g. right after ida_search_functions, instead of calling ida_xrefs_to in a loop.", "inputSchema": _schema({"targets": {"type": "array", "items": {"type": "string"}}}, ["targets"]), }, { "name": "ida_reload_plugin", "description": ( "Re-read the IDA plugin's Python code from disk without restarting IDA — use after the plugin file " "itself was edited, so bug fixes/new endpoints take effect without a full IDA restart. Note: the " "background heartbeat/registration thread keeps its old code until an actual IDA restart; only " "request-handling logic hot-reloads." ), "inputSchema": _schema(), }, { "name": "ida_compare", "description": ( "Decompile a target in two different IDA instances and return both side by side — e.g. to diff " "a patched binary against the original, or a symboled build against a stripped one." ), "inputSchema": { "type": "object", "properties": { "instance_a": {"type": "string", "description": "First instance id (see ida_list_instances)"}, "target_a": {"type": "string"}, "instance_b": {"type": "string", "description": "Second instance id"}, "target_b": {"type": "string"}, }, "required": ["instance_a", "target_a", "instance_b", "target_b"], }, }, ] async def execute_tool(name: str, arguments: dict): try: if name == "ida_list_instances": live = _live_instances() return { "count": len(live), "instances": [ {"id": k, "host": v["host"], "port": v["port"], "file": v["file"]} for k, v in live.items() ], } if name == "ida_compare": ha, pa, _, _ = resolve_instance(arguments["instance_a"]) hb, pb, _, _ = resolve_instance(arguments["instance_b"]) a = await ida_get(ha, pa, f"/decompile/{arguments['target_a']}") b = await ida_get(hb, pb, f"/decompile/{arguments['target_b']}") return {"instance_a": arguments["instance_a"], "instance_b": arguments["instance_b"], "a": a, "b": b} host, port, inst_id, auto_picked = resolve_instance(arguments.get("instance")) result = await _dispatch(name, arguments, host, port) if isinstance(result, dict) and auto_picked: result.setdefault( "_instance_used", f"{inst_id} (auto-picked, multiple instances were open — pass 'instance' to target a specific one)" ) return result except ApiError as e: return {"error": str(e), "code": e.code} except httpx.ConnectError: return {"error": "Cannot connect to the IDA instance. Is the plugin still running?", "code": "instance_unreachable"} except KeyError as e: return {"error": f"Missing required argument: {e}", "code": "bad_request"} except Exception as e: return {"error": str(e), "code": "internal_error"} async def _dispatch(name: str, arguments: dict, host, port): if name == "ida_ping": return await ida_get(host, port, "/ping") elif name == "ida_list_functions": f = arguments.get("filter", "") limit = arguments.get("limit", 100) offset = arguments.get("offset", 0) return await ida_get(host, port, f"/functions?filter={f}&limit={limit}&offset={offset}") elif name == "ida_search_functions": return await ida_get(host, port, f"/functions/search?q={arguments['q']}&limit={arguments.get('limit', 50)}") elif name == "ida_search_strings": return await ida_get(host, port, f"/strings/search?q={arguments['q']}&limit={arguments.get('limit', 30)}") elif name == "ida_get_name": return await ida_get(host, port, f"/name/{arguments['name']}") elif name == "ida_get_imports": return await ida_get(host, port, f"/imports?filter={arguments.get('filter', '')}") elif name == "ida_get_segments": return await ida_get(host, port, "/segments") elif name == "ida_get_entry_points": return await ida_get(host, port, "/entry_points") elif name == "ida_decompile": return await ida_get(host, port, f"/decompile/{arguments['target']}") elif name == "ida_decompile_many": return await ida_post(host, port, "/decompile_batch", {"targets": arguments["targets"]}) elif name == "ida_get_disasm": return await ida_get( host, port, f"/disasm/{arguments['target']}?count={arguments.get('count', 30)}&offset={arguments.get('offset', 0)}" ) elif name == "ida_xrefs_to": return await ida_get(host, port, f"/xrefs_to/{arguments['target']}") elif name == "ida_xrefs_from": return await ida_get(host, port, f"/xrefs_from/{arguments['target']}") elif name == "ida_get_prototype": return await ida_get(host, port, f"/prototype/{arguments['target']}") elif name == "ida_set_prototype": return await ida_post(host, port, "/set_prototype", {"target": arguments["target"], "prototype": arguments["prototype"]}) elif name == "ida_set_lvar_type": return await ida_post( host, port, "/set_lvar_type", {"function": arguments["function"], "lvar_name": arguments["lvar_name"], "type": arguments["type"]}, ) elif name == "ida_rename_function": return await ida_post(host, port, "/rename_function", {"target": arguments["target"], "new_name": arguments["new_name"]}) elif name == "ida_rename_local": return await ida_post( host, port, "/rename_local", {"function": arguments["function"], "old_name": arguments["old_name"], "new_name": arguments["new_name"]}, ) elif name == "ida_set_comment": return await ida_post(host, port, "/set_comment", {"target": arguments["target"], "comment": arguments["comment"]}) elif name == "ida_get_notes": return await ida_get(host, port, f"/notes/{arguments['target']}") elif name == "ida_set_notes": return await ida_post(host, port, "/notes", {"target": arguments["target"], "note": arguments["note"]}) elif name == "ida_rebuild_name_cache": return await ida_get(host, port, "/name_cache/rebuild") elif name == "ida_get_bytes": return await ida_get(host, port, f"/bytes/{arguments['target']}?size={arguments.get('size', 64)}") elif name == "ida_find_pattern": return await ida_get(host, port, f"/find?pattern={arguments['pattern']}&limit={arguments.get('limit', 50)}") elif name == "ida_patch_bytes": return await ida_post(host, port, "/patch_bytes", {"target": arguments["target"], "hex_bytes": arguments["hex_bytes"]}) elif name == "ida_make_data": return await ida_post( host, port, "/make_data", {"target": arguments["target"], "data_type": arguments["data_type"], "struct_name": arguments.get("struct_name")}, ) elif name == "ida_get_line_comment": rep = "1" if arguments.get("repeatable") else "0" return await ida_get(host, port, f"/line_comment/{arguments['target']}?repeatable={rep}") elif name == "ida_set_line_comment": return await ida_post( host, port, "/line_comment", {"target": arguments["target"], "comment": arguments["comment"], "repeatable": bool(arguments.get("repeatable", False))}, ) elif name == "ida_top_functions": return await ida_get(host, port, f"/top_functions?by={arguments.get('by', 'xrefs')}&limit={arguments.get('limit', 20)}") elif name == "ida_call_graph": return await ida_get( host, port, f"/call_graph/{arguments['target']}?depth={arguments.get('depth', 2)}&direction={arguments.get('direction', 'callees')}", ) elif name == "ida_make_function": return await ida_post(host, port, "/make_function", {"target": arguments["target"]}) elif name == "ida_undefine": return await ida_post(host, port, "/undefine", {"target": arguments["target"]}) elif name == "ida_list_local_types": return await ida_get(host, port, f"/local_types?limit={arguments.get('limit', 200)}") elif name == "ida_create_struct": return await ida_post(host, port, "/create_struct", {"name": arguments["name"], "fields": arguments["fields"]}) elif name == "ida_get_objc_classes": f = arguments.get("filter", "") return await ida_get(host, port, f"/objc_classes?filter={f}&limit={arguments.get('limit', 200)}") elif name == "ida_find_immediate": return await ida_get(host, port, f"/find_immediate?value={arguments['value']}&limit={arguments.get('limit', 100)}") elif name == "ida_find_vtables": seg = arguments.get("segment", "") or "" return await ida_get( host, port, f"/find_vtables?segment={seg}&min_entries={arguments.get('min_entries', 2)}&limit={arguments.get('limit', 50)}", ) elif name == "ida_create_local_type": return await ida_post(host, port, "/create_type", {"decl": arguments["decl"]}) elif name == "ida_xrefs_to_many": return await ida_post(host, port, "/xrefs_to_batch", {"targets": arguments["targets"]}) elif name == "ida_reload_plugin": return await ida_get(host, port, "/reload") return {"error": f"Unknown tool: {name}", "code": "unknown_tool"} async def handle_register_instance(request: Request): if not _check_token(request): return JSONResponse({"error": "unauthorized"}, status_code=401) try: body = await request.json() except Exception: return JSONResponse({"error": "invalid json"}, status_code=400) inst_id = body.get("id") host = body.get("host") port = body.get("port") if not inst_id or not host or not port: return JSONResponse({"error": "id, host, port are required"}, status_code=400) INSTANCES[inst_id] = { "host": host, "port": port, "file": body.get("file", ""), "pid": body.get("pid"), "last_seen": time.time(), } return JSONResponse({"ok": True}) async def handle_unregister_instance(request: Request): if not _check_token(request): return JSONResponse({"error": "unauthorized"}, status_code=401) try: body = await request.json() except Exception: body = {} INSTANCES.pop(body.get("id"), None) return JSONResponse({"ok": True}) async def handle_list_instances(request: Request): if not _check_token(request): return JSONResponse({"error": "unauthorized"}, status_code=401) live = _live_instances() return JSONResponse({"count": len(live), "instances": [{"id": k, **v} for k, v in live.items()]}) def make_response(req_id, result): return {"jsonrpc": "2.0", "id": req_id, "result": result} def make_error(req_id, code, message): return {"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}} async def handle_jsonrpc(msg: dict): method = msg.get("method") params = msg.get("params", {}) req_id = msg.get("id") if method == "initialize": return make_response( req_id, { "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "ida-mcp", "version": "4.0.0"}, }, ) elif method in ("notifications/initialized", "initialized"): return None elif method == "tools/list": return make_response(req_id, {"tools": TOOLS}) elif method == "tools/call": result = await execute_tool(params.get("name"), params.get("arguments", {})) is_error = isinstance(result, dict) and "error" in result return make_response( req_id, { "content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False, indent=2)}], "isError": is_error, }, ) elif method == "ping": return make_response(req_id, {}) else: return make_error(req_id, -32601, f"Method not found: {method}") async def handle_mcp_post(request: Request): if request.method == "HEAD": return Response(status_code=200) if request.method == "GET": return JSONResponse({"error": "use POST for JSON-RPC"}, status_code=400) if not _check_token(request): return JSONResponse(make_error(None, -32000, "unauthorized"), status_code=401) try: body = await request.json() except Exception: return JSONResponse(make_error(None, -32700, "Parse error"), status_code=400) if isinstance(body, list): responses = [r for r in [await handle_jsonrpc(m) for m in body] if r is not None] return JSONResponse(responses) result = await handle_jsonrpc(body) if result is None: return Response(status_code=202) return JSONResponse(result) async def handle_health(request: Request): return JSONResponse( { "status": "ok", "server": "ida-mcp-dynamic", "version": "4.0.0", "instances": len(_live_instances()), "auth_required": bool(AUTH_TOKEN), } ) async def handle_oauth_protected_resource(request: Request): return JSONResponse( { "resource": NGROK_URL, "authorization_servers": [], "bearer_methods_supported": ["header"], "scopes_supported": [], } ) async def handle_oauth_authorization_server(request: Request): return JSONResponse( { "issuer": NGROK_URL, "authorization_endpoint": f"{NGROK_URL}/authorize", "token_endpoint": f"{NGROK_URL}/token", "registration_endpoint": f"{NGROK_URL}/register", "grant_types_supported": ["authorization_code"], "response_types_supported": ["code"], "code_challenge_methods_supported": ["S256", "plain"], "token_endpoint_auth_methods_supported": ["none"], "scopes_supported": [], } ) def _authorize_form_html(redirect_uri: str, state: str, error: str = "") -> str: error_html = f'

{error}

' if error else "" return f"""

IDA MCP — enter access token

{error_html}
""" def _issue_code() -> str: code = secrets.token_urlsafe(24) AUTH_CODES[code] = time.time() + AUTH_CODE_TTL return code def _redirect_with_code(redirect_uri: str, state: str) -> HTMLResponse: code = _issue_code() sep = "&" if "?" in redirect_uri else "?" location = f"{redirect_uri}{sep}code={code}" if state: location += f"&state={state}" return HTMLResponse(f'') async def handle_authorize(request: Request): params = request.query_params redirect_uri = params.get("redirect_uri") state = params.get("state", "") if not redirect_uri: return JSONResponse({"error": "missing redirect_uri"}, status_code=400) if not AUTH_TOKEN: return _redirect_with_code(redirect_uri, state) submitted = params.get("token") if submitted is None: return HTMLResponse(_authorize_form_html(redirect_uri, state)) if submitted != AUTH_TOKEN: return HTMLResponse(_authorize_form_html(redirect_uri, state, error="Wrong token"), status_code=401) return _redirect_with_code(redirect_uri, state) async def handle_register(request: Request): try: body = await request.json() except Exception: body = {} redirect_uris = body.get("redirect_uris", []) return JSONResponse( { "client_id": "claude-ai-client", "client_secret": "not-used", "client_id_issued_at": int(time.time()), "client_secret_expires_at": 0, "redirect_uris": redirect_uris, "grant_types": body.get("grant_types", ["authorization_code"]), "response_types": body.get("response_types", ["code"]), "token_endpoint_auth_method": "none", "application_type": "web", }, status_code=201, ) async def handle_token(request: Request): if not AUTH_TOKEN: return JSONResponse({"access_token": "ida-mcp-open-access", "token_type": "bearer", "expires_in": 86400}) try: form = await request.form() code = form.get("code") except Exception: code = None if not code: try: body = await request.json() code = body.get("code") except Exception: code = None expiry = AUTH_CODES.pop(code, None) if code else None if not expiry or expiry < time.time(): return JSONResponse({"error": "invalid_grant", "error_description": "missing/expired/used code"}, status_code=400) return JSONResponse({"access_token": AUTH_TOKEN, "token_type": "bearer", "expires_in": 86400}) def create_app(): return Starlette( routes=[ Route("/health", handle_health), Route("/instances", handle_list_instances), Route("/instances/register", handle_register_instance, methods=["POST"]), Route("/instances/unregister", handle_unregister_instance, methods=["POST"]), Route("/.well-known/oauth-protected-resource", handle_oauth_protected_resource), Route("/.well-known/oauth-protected-resource/sse", handle_oauth_protected_resource), Route("/.well-known/oauth-authorization-server", handle_oauth_authorization_server), Route("/register", handle_register, methods=["POST"]), Route("/authorize", handle_authorize, methods=["GET"]), Route("/token", handle_token, methods=["POST"]), Route("/sse", handle_mcp_post, methods=["GET", "POST", "HEAD"]), ] ) if __name__ == "__main__": import uvicorn print(f"[IDA MCP] Dynamic multi-instance mode. MCP endpoint: http://0.0.0.0:{SSE_PORT}/sse") print(f"[IDA MCP] Waiting for IDA plugin(s) to self-register on /instances/register ...") if not AUTH_TOKEN: print("[IDA MCP] WARNING: IDA_MCP_TOKEN is not set — server is open to anyone who can reach it.") print("[IDA MCP] Set IDA_MCP_TOKEN before exposing this via ngrok/the internet.") if NGROK_URL: print(f"[IDA MCP] NGROK_URL: {NGROK_URL}") uvicorn.run(create_app(), host="0.0.0.0", port=SSE_PORT)