875 lines
36 KiB
Python
875 lines
36 KiB
Python
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'<p style="color:#c00">{error}</p>' if error else ""
|
|
return f"""<!doctype html><html><body style="font-family:sans-serif;max-width:420px;margin:80px auto">
|
|
<h3>IDA MCP — enter access token</h3>
|
|
{error_html}
|
|
<form method="GET" action="/authorize">
|
|
<input type="hidden" name="redirect_uri" value="{redirect_uri}">
|
|
<input type="hidden" name="state" value="{state}">
|
|
<input type="password" name="token" placeholder="IDA_MCP_TOKEN" style="width:100%;padding:8px;margin:8px 0" autofocus>
|
|
<button type="submit" style="padding:8px 16px">Authorize</button>
|
|
</form>
|
|
</body></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'<script>window.location="{location}"</script>')
|
|
|
|
|
|
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)
|