diff options
Diffstat (limited to 'idatui/client.py')
| -rw-r--r-- | idatui/client.py | 636 |
1 files changed, 0 insertions, 636 deletions
diff --git a/idatui/client.py b/idatui/client.py deleted file mode 100644 index 2f74cbe..0000000 --- a/idatui/client.py +++ /dev/null @@ -1,636 +0,0 @@ -"""Persistent, thread-safe client for the ida-pro-mcp (idalib) MCP server. - -DEPRECATED. The default backend is now our own idalib worker (idatui.worker + -idatui.worker_client.WorkerClient), which talks a unix socket instead of HTTP and -is ~50-100x cheaper per call. This module (and the whole ida-pro-mcp -supervisor/HTTP path: server/patch_server.py, spawn.sh, launch.py's server -plumbing) is kept only for `--backend mcp` / --db attach and will be removed once -the worker is battle-tested. Prefer WorkerClient for new code. - -This is the foundation the whole TUI stands on. Unlike the throwaway CLI in the -ida-mcp skill (which re-does the MCP handshake and spins a fresh socket on every -call), this client: - - * Performs the MCP handshake exactly once and keeps the session warm - (measured: ~7ms/call warm vs ~60ms cold). - * Reuses TCP connections via a small keep-alive pool (no socket churn over a - long TUI session) while still allowing concurrent calls from worker threads. - * Has a precise, *grounded* error taxonomy (see below) so callers can tell - apart transport failures, protocol errors, hard tool errors, and soft - per-item "not found" results. - * Recovers automatically from an expired server session (404 -> re-handshake - -> retry once) and from transient transport hiccups (bounded retries). - * Auto-injects the mandatory ``database=<session_id>`` argument, resolving a - single open session automatically and refusing to guess when several are - open. - -Error taxonomy (verified against the live server, 2026-07): - - IDAConnectionError transport could not be established / was lost - IDATimeoutError a request exceeded its deadline - IDAProtocolError malformed / unexpected HTTP or JSON-RPC framing - IDARPCError JSON-RPC ``error`` object in the envelope - IDAToolError tool returned ``result.isError == true`` (hard failure: - bad params, unknown tool, missing database, ...) - IDASessionError no session / multiple sessions and none pinned - -Note the deliberate *non*-error: a tool that returns ``isError == false`` with an -``error`` field inside its payload (e.g. ``decompile`` on an unknown address, or -per-item failures in a batch tool) is treated as *data*, not an exception. The -caller inspects the payload. Raising on those would break every batch tool. - -The client is stdlib-only (``http.client`` / ``urllib.parse``). -""" - -from __future__ import annotations - -import http.client -import json -import socket -import threading -import time -from collections import deque -from dataclasses import dataclass -from typing import Any -from urllib.parse import urlsplit - -DEFAULT_URL = "http://127.0.0.1:8745/mcp" -PROTOCOL_VERSION = "2025-06-18" - -# Tools that must NOT receive an injected ``database`` argument. -SESSION_AGNOSTIC_TOOLS = frozenset({"idb_list", "idb_open", "int_convert"}) - -# Substrings (lowercased) that identify a stale/invalid IDB session in a tool -# error message. Verified against the live server: after a server restart the -# old session id is gone and every call fails with "Session not found: <id>". -_STALE_SESSION_MARKERS = ("session not found", "database is required") - - -# The error hierarchy and Session model now live in errors.py (transport- -# agnostic, shared with the idalib worker path); re-exported here so the -# deprecated mcp tooling and the stress tests keep importing them from client. -from .errors import ( # noqa: E402,F401 - IDAError, - IDAConnectionError, - IDATimeoutError, - IDAProtocolError, - IDARPCError, - IDAToolError, - IDASessionError, - Session, -) - - -# --------------------------------------------------------------------------- # -# Transport: a tiny keep-alive connection pool over http.client -# --------------------------------------------------------------------------- # -class _Transport: - """Keep-alive HTTP/1.1 pool. Thread-safe. One request == one pooled conn. - - Connections are checked out, used for exactly one request/response, then - returned to the pool if still healthy. On a dropped/broken connection the - request is retried on a fresh connection (bounded by ``max_retries``). - """ - - def __init__(self, url: str, *, max_retries: int = 2, pool_size: int = 8): - parts = urlsplit(url) - if parts.scheme not in ("http", "https"): - raise IDAConnectionError(f"unsupported scheme: {parts.scheme!r}") - self._https = parts.scheme == "https" - self._host = parts.hostname or "127.0.0.1" - self._port = parts.port or (443 if self._https else 80) - self._path = parts.path or "/" - self._max_retries = max_retries - self._pool_size = pool_size - self._idle: deque[http.client.HTTPConnection] = deque() - self._lock = threading.Lock() - - def _new_conn(self, timeout: float) -> http.client.HTTPConnection: - cls = http.client.HTTPSConnection if self._https else http.client.HTTPConnection - return cls(self._host, self._port, timeout=timeout) - - def _checkout(self, timeout: float) -> http.client.HTTPConnection: - with self._lock: - while self._idle: - conn = self._idle.popleft() - # Reuse only if the socket still looks alive. - if conn.sock is not None: - conn.timeout = timeout - try: - conn.sock.settimeout(timeout) - except OSError: - try: - conn.close() - except OSError: - pass - continue - return conn - return self._new_conn(timeout) - - def _checkin(self, conn: http.client.HTTPConnection) -> None: - with self._lock: - if len(self._idle) < self._pool_size: - self._idle.append(conn) - return - try: - conn.close() - except OSError: - pass - - def request( - self, body: bytes, headers: dict[str, str], timeout: float - ) -> tuple[int, dict[str, str], bytes]: - """POST ``body``; return (status, response_headers, response_body).""" - last_exc: Exception | None = None - for attempt in range(self._max_retries + 1): - conn = self._checkout(timeout) - try: - conn.request("POST", self._path, body=body, headers=headers) - resp = conn.getresponse() - status = resp.status - resp_headers = {k.lower(): v for k, v in resp.getheaders()} - data = resp.read() # must fully drain before reuse - except socket.timeout as e: - self._discard(conn) - raise IDATimeoutError(f"request timed out after {timeout}s") from e - except ( - http.client.RemoteDisconnected, - http.client.BadStatusLine, - ConnectionError, - OSError, - ) as e: - # A stale pooled connection, or the server closed on us. Drop it - # and retry on a fresh connection. - self._discard(conn) - last_exc = e - continue - else: - if resp.will_close: - self._discard(conn) - else: - self._checkin(conn) - return status, resp_headers, data - raise IDAConnectionError( - f"transport failed after {self._max_retries + 1} attempts: {last_exc}" - ) from last_exc - - def _discard(self, conn: http.client.HTTPConnection) -> None: - try: - conn.close() - except OSError: - pass - - def close(self) -> None: - with self._lock: - while self._idle: - self._discard(self._idle.popleft()) - - -# --------------------------------------------------------------------------- # -# The client -# --------------------------------------------------------------------------- # -class IDAClient: - """A warm, thread-safe handle to one MCP server (and one pinned IDB). - - Typical use:: - - with IDAClient(db="4f2223f9") as ida: - ida.health() - funcs = ida.call("list_funcs", queries=[{"count": 50}]) - code = ida.call("decompile", addr="main") - - Concurrency: safe to call from multiple threads (Textual workers). Requests - run on independent pooled connections; only the id counter and handshake - state are lock-guarded, so calls do not serialize on each other. - """ - - def __init__( - self, - url: str = DEFAULT_URL, - db: str | None = None, - *, - timeout: float = 30.0, - max_retries: int = 2, - pool_size: int = 8, - client_name: str = "idatui", - client_version: str = "0.0.1", - auto_recover_session: bool = True, - ): - self.url = url - self._db = db - self.timeout = timeout - self.auto_recover_session = auto_recover_session - self._transport = _Transport(url, max_retries=max_retries, pool_size=pool_size) - self._client_info = {"name": client_name, "version": client_version} - - self._rid = 0 - self._sid: str | None = None - self._ready = False - self._state_lock = threading.Lock() # guards _rid, _sid, _ready - self._handshake_lock = threading.Lock() # serializes (re)handshake - - # -- lifecycle --------------------------------------------------------- # - def __enter__(self) -> "IDAClient": - self.connect() - return self - - def __exit__(self, *exc) -> None: - self.close() - - def connect(self) -> "IDAClient": - """Ensure the MCP handshake has completed (idempotent, thread-safe).""" - if self._ready: - return self - self._handshake() - return self - - def close(self) -> None: - self._transport.close() - - # -- low-level plumbing ------------------------------------------------ # - def _next_id(self) -> int: - with self._state_lock: - self._rid += 1 - return self._rid - - def _headers(self) -> dict[str, str]: - h = { - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - } - sid = self._sid - if sid: - h["Mcp-Session-Id"] = sid - return h - - def _handshake(self) -> None: - with self._handshake_lock: - if self._ready: - return - with self._state_lock: - self._sid = None - rid = self._next_id() - status, headers, body = self._transport.request( - self._encode( - { - "jsonrpc": "2.0", - "id": rid, - "method": "initialize", - "params": { - "protocolVersion": PROTOCOL_VERSION, - "capabilities": {}, - "clientInfo": self._client_info, - }, - } - ), - { - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - }, - self.timeout, - ) - if status // 100 != 2: - raise IDAProtocolError( - f"initialize failed: HTTP {status}: {body[:200]!r}" - ) - sid = headers.get("mcp-session-id") - envelope = self._parse_body(body, rid) - self._raise_on_rpc_error(envelope) - with self._state_lock: - self._sid = sid - # notifications/initialized is a fire-and-forget notification. - self._transport.request( - self._encode({"jsonrpc": "2.0", "method": "notifications/initialized"}), - self._headers(), - self.timeout, - ) - with self._state_lock: - self._ready = True - - @staticmethod - def _encode(obj: dict) -> bytes: - return json.dumps(obj).encode() - - @staticmethod - def _parse_body(body: bytes, want_id: int) -> dict: - """Extract the JSON-RPC envelope for ``want_id`` from a (possibly SSE) body.""" - text = body.decode("utf-8", "replace") - found: dict | None = None - for line in text.splitlines(): - line = line.strip() - if line.startswith("data:"): - line = line[5:].strip() - if not line: - continue - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - if isinstance(obj, dict) and obj.get("id") == want_id: - found = obj - if found is None: - raise IDAProtocolError( - f"no JSON-RPC response with id={want_id} in body: {text[:200]!r}" - ) - return found - - @staticmethod - def _raise_on_rpc_error(envelope: dict) -> None: - err = envelope.get("error") - if err: - raise IDARPCError( - code=err.get("code", -1), - message=err.get("message", "unknown"), - data=err.get("data"), - ) - - def _rpc(self, method: str, params: dict, *, timeout: float | None = None) -> dict: - """Send a JSON-RPC request, recovering once from an expired session.""" - if not self._ready: - self._handshake() - to = self.timeout if timeout is None else timeout - rid = self._next_id() - payload = self._encode( - {"jsonrpc": "2.0", "id": rid, "method": method, "params": params} - ) - status, headers, body = self._transport.request(payload, self._headers(), to) - - if status == 404 and self._sid is not None: - # Server session expired: re-handshake and retry exactly once. - with self._state_lock: - self._ready = False - self._handshake() - rid = self._next_id() - payload = self._encode( - {"jsonrpc": "2.0", "id": rid, "method": method, "params": params} - ) - status, headers, body = self._transport.request( - payload, self._headers(), to - ) - - if status // 100 != 2: - raise IDAProtocolError(f"HTTP {status}: {body[:200]!r}") - - envelope = self._parse_body(body, rid) - self._raise_on_rpc_error(envelope) - return envelope - - # -- tool calls -------------------------------------------------------- # - @staticmethod - def _extract_payload(tool: str, result: dict) -> Any: - """Turn an MCP ``result`` object into a Python payload, or raise. - - Precedence: - 1. ``isError == true`` -> IDAToolError (hard failure) - 2. ``structuredContent`` present -> return it (already parsed) - 3. ``content[0].text`` is JSON -> return parsed JSON - 4. otherwise -> return the raw text string - """ - if result.get("isError"): - text = _first_text(result) or "(no message)" - raise IDAToolError(tool, text) - if "structuredContent" in result: - return result["structuredContent"] - text = _first_text(result) - if text is None: - return result - try: - return json.loads(text) - except (json.JSONDecodeError, TypeError): - return text - - def _prepare_args(self, tool: str, args: dict) -> tuple[dict, bool]: - """Return (arguments, db_was_injected). Never mutates the caller's dict.""" - prepared = dict(args) - if tool in SESSION_AGNOSTIC_TOOLS or "database" in prepared: - return prepared, False - prepared["database"] = self.resolve_db() - return prepared, True - - def call_envelope(self, tool: str, *, timeout: float | None = None, **args) -> dict: - """Call ``tool`` and return the full JSON-RPC envelope (for debugging).""" - prepared, _ = self._prepare_args(tool, args) - return self._rpc( - "tools/call", {"name": tool, "arguments": prepared}, timeout=timeout - ) - - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - """Call ``tool`` and return its payload (parsed JSON when possible). - - Raises IDAToolError on a hard tool failure. Soft/per-item errors (an - ``error`` field with ``isError == false``) are returned as data. - - If ``auto_recover_session`` is set and the db was auto-injected, a stale - "Session not found" error (e.g. after a server restart) triggers exactly - one transparent recovery: drop the stale pin, re-resolve the session, - and retry. A db the caller pinned explicitly is never silently switched. - """ - prepared, injected = self._prepare_args(tool, args) - envelope = self._rpc( - "tools/call", {"name": tool, "arguments": prepared}, timeout=timeout - ) - try: - return self._extract_payload(tool, envelope.get("result", {})) - except IDAToolError as e: - if not (injected and self.auto_recover_session and _is_stale_session(e)): - raise - # The pinned IDB session vanished (server restart). Re-resolve the - # sole session and retry once. resolve_db() raises IDASessionError - # if zero/many sessions exist, so we never guess. - self.set_db(None) - prepared2, _ = self._prepare_args(tool, args) - envelope2 = self._rpc( - "tools/call", {"name": tool, "arguments": prepared2}, timeout=timeout - ) - return self._extract_payload(tool, envelope2.get("result", {})) - - # -- session management ------------------------------------------------ # - def list_sessions(self) -> list[Session]: - result = self._rpc("tools/call", {"name": "idb_list", "arguments": {}}) - payload = self._extract_payload("idb_list", result.get("result", {})) - sessions = payload.get("sessions", []) if isinstance(payload, dict) else [] - return [Session.from_dict(s) for s in sessions] - - def resolve_db(self) -> str: - """Return the pinned session id, auto-resolving a lone open session. - - Raises IDASessionError if none is open, or if several are open and none - has been pinned via ``db=`` / :meth:`set_db`. - """ - if self._db: - return self._db - sessions = self.list_sessions() - usable = [s for s in sessions if s.session_id] - if len(usable) == 1: - self._db = usable[0].session_id - return self._db - if not usable: - raise IDASessionError( - "no open IDB session with a usable id; open one with idb_open" - ) - opts = ", ".join(f"{s.session_id} ({s.filename})" for s in usable) - raise IDASessionError( - f"multiple sessions open; pin one with db=... . Options: {opts}" - ) - - def set_db(self, db: str | None) -> None: - self._db = db - - @property - def db(self) -> str | None: - return self._db - - def _session_input_path(self, db: str) -> str: - for s in self.list_sessions(): - if s.session_id == db: - if not s.input_path: - raise IDASessionError(f"session {db} has no input_path to re-open") - return s.input_path - raise IDASessionError(f"session {db} not found") - - def bump_idle_ttl(self, idle_ttl_sec: int = 1_000_000_000, - path: str | None = None) -> None: - """Raise the worker's idle self-exit TTL so an interactive session never - gets reaped while the user is just reading. - - Headless idalib workers self-exit after ``idle_ttl_sec`` (default 600s) - with no requests. ``idb_open`` on the already-open path is idempotent - (returns the same session) and re-applies the TTL, so this simply opens - the pinned session's path with a huge TTL. The default (~31 years) is - effectively 'never'. - """ - p = path or self._session_input_path(self.resolve_db()) - self.call("idb_open", input_path=p, idle_ttl_sec=int(idle_ttl_sec)) - - def keepalive(self, interval: float = 120.0) -> "KeepAlive": - """Return a (not-yet-started) heartbeat that touches the worker's idle - watchdog every ``interval`` seconds. Belt-and-suspenders next to - :meth:`bump_idle_ttl`; also covers adopted sessions we didn't open. - """ - return KeepAlive(self, interval=interval) - - # -- convenience ------------------------------------------------------- # - def health(self) -> dict: - return self.call("server_health") - - def list_tools(self) -> list[tuple[str, str]]: - envelope = self._rpc("tools/list", {}) - out = [] - for t in envelope.get("result", {}).get("tools", []): - desc = (t.get("description") or "").splitlines() - out.append((t["name"], desc[0] if desc else "")) - return out - - def schema(self, tool: str) -> dict: - envelope = self._rpc("tools/list", {}) - for t in envelope.get("result", {}).get("tools", []): - if t["name"] == tool: - return t.get("inputSchema", {}) - raise IDAError(f"tool not found: {tool}") - - -class KeepAlive: - """Background heartbeat that keeps an idalib worker from idling out. - - Any forwarded request resets the worker's idle timer, so a periodic cheap - ``server_health`` is enough. Failures are swallowed (the next real call will - auto-recover); the heartbeat just keeps a chilling TUI's worker alive. - """ - - def __init__(self, client: "IDAClient", interval: float = 120.0): - if interval <= 0: - raise ValueError("interval must be > 0") - self._client = client - self.interval = interval - self._stop = threading.Event() - self._thread: threading.Thread | None = None - self.beats = 0 - self.failures = 0 - - def start(self) -> "KeepAlive": - if self._thread is not None: - return self - self._stop.clear() - self._thread = threading.Thread( - target=self._run, daemon=True, name="idatui-keepalive" - ) - self._thread.start() - return self - - def stop(self) -> None: - self._stop.set() - t = self._thread - self._thread = None - if t is not None: - t.join(timeout=2.0) - - def __enter__(self) -> "KeepAlive": - return self.start() - - def __exit__(self, *exc) -> None: - self.stop() - - def _run(self) -> None: - while not self._stop.wait(self.interval): - try: - self._client.health() - self.beats += 1 - except IDAError: - self.failures += 1 - - -def _is_stale_session(e: IDAToolError) -> bool: - msg = e.message.lower() - return any(marker in msg for marker in _STALE_SESSION_MARKERS) - - -def _first_text(result: dict) -> str | None: - content = result.get("content") - if isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - return item.get("text") - return None - - -# --------------------------------------------------------------------------- # -# Tiny self-check CLI: python -m idatui.client [--db ID] [--url URL] [health] -# --------------------------------------------------------------------------- # -def _main(argv: list[str]) -> int: - import os - - url = DEFAULT_URL - db = os.environ.get("IDA_MCP_DB") - rest: list[str] = [] - it = iter(argv) - for a in it: - if a == "--url": - url = next(it) - elif a == "--db": - db = next(it) - else: - rest.append(a) - - ida = IDAClient(url, db=db) - try: - ida.connect() - t0 = time.time() - sessions = ida.list_sessions() - print(f"sessions ({(time.time() - t0) * 1e3:.1f}ms):") - for s in sessions: - mark = "*" if s.is_active else " " - print(f" {mark} {s.session_id or '<none>':10} {s.filename}") - try: - h = ida.health() - print("health:", json.dumps(h, indent=1)[:400]) - except IDASessionError as e: - print(f"health: skipped ({e})") - finally: - ida.close() - return 0 - - -if __name__ == "__main__": - import sys - - raise SystemExit(_main(sys.argv[1:])) |