1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
|
"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor.
Opens ONE database in-process (on the main thread, as idalib requires) and
serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed
pickle. Same tool implementations as the MCP path (we call
``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are
byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no
supervisor / HTTP / JSON / 50KB-truncation machinery.
python -m idatui.worker <sock_path> <binary_path>
The socket only appears once the database is open + analyzed, so a client can
poll ``connect()`` to know when the worker is ready. Requests are served
serially on the main thread (idalib is single-threaded; every tool runs inline
through its own execute_sync, which is a no-op on the main thread).
Protocol (both directions length-prefixed: 4-byte big-endian len + pickle):
request = (tool_name: str, kwargs: dict)
response = (ok: bool, result_or_error)
tool_name == "__shutdown__" ends the worker.
"""
from __future__ import annotations
import os
import pickle
import socket
import struct
import sys
import uuid
# --------------------------------------------------------------------------- #
# framing
# --------------------------------------------------------------------------- #
def _recvn(sock: socket.socket, n: int) -> bytes | None:
buf = bytearray()
while len(buf) < n:
chunk = sock.recv(n - len(buf))
if not chunk:
return None
buf += chunk
return bytes(buf)
def send(sock: socket.socket, obj) -> None:
data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
sock.sendall(struct.pack(">I", len(data)) + data)
def recv(sock: socket.socket):
hdr = _recvn(sock, 4)
if hdr is None:
return None
(n,) = struct.unpack(">I", hdr)
body = _recvn(sock, n)
return None if body is None else pickle.loads(body)
# --------------------------------------------------------------------------- #
# worker
# --------------------------------------------------------------------------- #
def _ensure_tools_injected() -> None:
"""Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...)
into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient
(nothing else has to inject these tools first). Must run BEFORE
ida_pro_mcp.ida_mcp is imported (the injected code lives in api_types.py)."""
import importlib.util
repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
patch = os.path.join(repo, "server", "patch_server.py")
if not os.path.exists(patch):
return
try:
spec = importlib.util.spec_from_file_location("_idatui_patch", patch)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types
mod.main()
except Exception as e: # noqa: BLE001 -- tools may already be present
sys.stderr.write(f"idatui: tool injection skipped: {e}\n")
def _open_and_register(binpath: str):
"""Open the DB (main thread) then import ida-pro-mcp so every @tool registers
against this live database. Returns (tools_dict, module_name, save_fn)."""
_ensure_tools_injected() # before any ida_pro_mcp import
import idapro
idapro.enable_console_messages(False)
if idapro.open_database(binpath, run_auto_analysis=True): # nonzero == failure
raise RuntimeError(
f"failed to open {binpath}: the .i64 is likely held by a running "
f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash "
f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)")
import ida_auto
ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp)
# importing the package registers all api_*/patched tools against MCP_SERVER
from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433
import ida_nalt
module = os.path.basename(ida_nalt.get_root_filename() or binpath)
def save():
import idc
try:
idc.save_database(idc.get_idb_path(), 0)
except Exception: # noqa: BLE001
import ida_loader, ida_pro # noqa: WPS433
ida_loader.save_database(idc.get_idb_path(), 0)
return MCP_SERVER.tools.methods, module, save
def serve(sockpath: str, binpath: str) -> None:
tools, module, save = _open_and_register(binpath)
sid = uuid.uuid4().hex[:8]
def dispatch(name: str, args: dict):
args = dict(args)
args.pop("database", None) # single-DB worker: no session routing
# session-management shims (were the supervisor's job):
if name in ("idb_open",):
return {"success": True,
"session": {"session_id": sid, "module": module,
"input_path": binpath}}
if name in ("idb_save", "save"):
save()
return {"success": True}
if name in ("server_health", "ping", "health", "state"):
return {"module": module, "ok": True, "session_id": sid}
if name in ("idb_list",):
return {"sessions": [{"session_id": sid, "module": module,
"input_path": binpath}]}
fn = tools.get(name)
if fn is None:
raise KeyError(f"unknown tool: {name!r}")
result = fn(**args)
# Match the MCP server's structuredContent: a dict passes through, any
# other return (list/scalar) is wrapped as {"result": ...}. domain.py
# parses that exact shape (e.g. lookup_funcs -> payload["result"]).
return result if isinstance(result, dict) else {"result": result}
try:
os.unlink(sockpath)
except OSError:
pass
srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
srv.bind(sockpath)
srv.listen(8)
try:
while True:
conn, _ = srv.accept()
try:
while True:
req = recv(conn)
if req is None:
break
name, args = req
if name == "__shutdown__":
return
try:
send(conn, (True, dispatch(name, args)))
except Exception as e: # noqa: BLE001 -- report, keep serving
send(conn, (False, f"{type(e).__name__}: {e}"))
except (ConnectionError, OSError):
pass
finally:
conn.close()
finally:
try:
import idapro
idapro.close_database(save=False)
except Exception: # noqa: BLE001
pass
try:
os.unlink(sockpath)
except OSError:
pass
def main(argv=None) -> None:
argv = argv if argv is not None else sys.argv[1:]
if len(argv) < 2:
sys.stderr.write("usage: python -m idatui.worker <sock> <binary>\n")
raise SystemExit(2)
try:
serve(argv[0], argv[1])
except SystemExit:
raise
except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1
import traceback
sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n")
traceback.print_exc()
sys.stderr.flush()
raise SystemExit(1)
if __name__ == "__main__":
main()
|
