diff options
| author | blasty <blasty@local> | 2026-07-09 11:14:23 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-07-09 11:14:23 +0200 |
| commit | eec630887ce72a7dafae4c33f5e85ad7c746b038 (patch) | |
| tree | 7f5c736249e6ba32317cb4c17a25a4652661593b | |
| download | ida-tui-eec630887ce72a7dafae4c33f5e85ad7c746b038.tar.gz ida-tui-eec630887ce72a7dafae4c33f5e85ad7c746b038.tar.xz ida-tui-eec630887ce72a7dafae4c33f5e85ad7c746b038.zip | |
persistent MCP client: warm handshake, keep-alive pool, error taxonomy, thread-safe
- IDAClient: single handshake, ~7ms warm calls (vs ~60ms cold CLI)
- keep-alive connection pool over http.client; concurrent-safe (40 threads OK)
- grounded error taxonomy: IDAToolError on isError, soft per-item errors as data
- session-expiry recovery (404 -> re-handshake -> retry), bounded transport retries
- auto db-injection + resolve_db (ignores stale empty-id sessions)
- live smoke test: 13/13 pass
| -rw-r--r-- | .gitignore | 12 | ||||
| -rw-r--r-- | idatui/__init__.py | 25 | ||||
| -rw-r--r-- | idatui/__main__.py | 6 | ||||
| -rw-r--r-- | idatui/client.py | 563 | ||||
| -rw-r--r-- | pyproject.toml | 20 | ||||
| -rwxr-xr-x | spawn.sh | 2 | ||||
| -rw-r--r-- | tests/smoke_client.py | 152 |
7 files changed, 780 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2b1bd85 --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +__pycache__/ +*.pyc +.venv/ +dist/ +build/ +*.egg-info/ +*.i64 +*.id0 +*.id1 +*.id2 +*.nam +*.til diff --git a/idatui/__init__.py b/idatui/__init__.py new file mode 100644 index 0000000..450d64f --- /dev/null +++ b/idatui/__init__.py @@ -0,0 +1,25 @@ +"""idatui — a minimal keyboard-first TUI for IDA Pro over ida-pro-mcp (idalib).""" + +from .client import ( + IDAClient, + IDAError, + IDAConnectionError, + IDATimeoutError, + IDAProtocolError, + IDARPCError, + IDAToolError, + IDASessionError, + Session, +) + +__all__ = [ + "IDAClient", + "IDAError", + "IDAConnectionError", + "IDATimeoutError", + "IDAProtocolError", + "IDARPCError", + "IDAToolError", + "IDASessionError", + "Session", +] diff --git a/idatui/__main__.py b/idatui/__main__.py new file mode 100644 index 0000000..29f5cf4 --- /dev/null +++ b/idatui/__main__.py @@ -0,0 +1,6 @@ +"""``python -m idatui`` -> client self-check (until the TUI app lands).""" +import sys + +from .client import _main + +raise SystemExit(_main(sys.argv[1:])) diff --git a/idatui/client.py b/idatui/client.py new file mode 100644 index 0000000..2c8cc68 --- /dev/null +++ b/idatui/client.py @@ -0,0 +1,563 @@ +"""Persistent, thread-safe client for the ida-pro-mcp (idalib) MCP server. + +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"}) + + +# --------------------------------------------------------------------------- # +# Exceptions +# --------------------------------------------------------------------------- # +class IDAError(Exception): + """Base class for all client errors.""" + + +class IDAConnectionError(IDAError): + """The transport could not be established or was lost.""" + + +class IDATimeoutError(IDAError): + """A request exceeded its deadline.""" + + +class IDAProtocolError(IDAError): + """Malformed or unexpected HTTP / JSON-RPC framing.""" + + +class IDARPCError(IDAError): + """The JSON-RPC envelope carried an ``error`` object.""" + + def __init__(self, code: int, message: str, data: Any = None): + super().__init__(f"JSON-RPC error {code}: {message}") + self.code = code + self.message = message + self.data = data + + +class IDAToolError(IDAError): + """A tool call returned ``result.isError == true`` (a hard failure).""" + + def __init__(self, tool: str, message: str): + super().__init__(f"tool {tool!r} failed: {message}") + self.tool = tool + self.message = message + + +class IDASessionError(IDAError): + """No IDB session is open, or several are and none was pinned.""" + + +# --------------------------------------------------------------------------- # +# Session model +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Session: + session_id: str + filename: str + input_path: str + is_active: bool = False + is_analyzing: bool = False + + @classmethod + def from_dict(cls, d: dict) -> "Session": + return cls( + session_id=d.get("session_id", ""), + filename=d.get("filename", ""), + input_path=d.get("input_path", ""), + is_active=bool(d.get("is_active", False)), + is_analyzing=bool(d.get("is_analyzing", False)), + ) + + +# --------------------------------------------------------------------------- # +# 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", + ): + self.url = url + self._db = db + self.timeout = timeout + 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 call_envelope(self, tool: str, *, timeout: float | None = None, **args) -> dict: + """Call ``tool`` and return the full JSON-RPC envelope (for debugging).""" + self._inject_db(tool, args) + return self._rpc("tools/call", {"name": tool, "arguments": args}, 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. + """ + envelope = self.call_envelope(tool, timeout=timeout, **args) + return self._extract_payload(tool, envelope.get("result", {})) + + def _inject_db(self, tool: str, args: dict) -> None: + if tool in SESSION_AGNOSTIC_TOOLS: + return + if "database" in args: + return + args["database"] = self.resolve_db() + + # -- 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 + + # -- 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}") + + +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:])) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e491542 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,20 @@ +[project] +name = "idatui" +version = "0.0.1" +description = "A minimal keyboard-first TUI frontend for IDA Pro over the ida-pro-mcp (idalib) server." +requires-python = ">=3.11" +# The client layer is intentionally stdlib-only (urllib/http.client), matching the +# ida-mcp skill philosophy: no install needed to talk to the server. +dependencies = [] + +[project.optional-dependencies] +# The TUI layer (added in a later phase) pulls in Textual. +tui = ["textual>=0.60"] +dev = ["pytest>=8"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["idatui"] diff --git a/spawn.sh b/spawn.sh new file mode 100755 index 0000000..8bc0b6f --- /dev/null +++ b/spawn.sh @@ -0,0 +1,2 @@ +#!/bin/sh +uv run idalib-mcp --host 127.0.0.1 --port 8745 $(pwd)/bin/ls diff --git a/tests/smoke_client.py b/tests/smoke_client.py new file mode 100644 index 0000000..647b589 --- /dev/null +++ b/tests/smoke_client.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +"""Live smoke test for idatui.client against a running ida-pro-mcp server. + +Run with the venv Python while a server is up (see spawn.sh): + + IDA_MCP_DB=<session_id> python3 tests/smoke_client.py + # or: python3 tests/smoke_client.py --db <session_id> + +Exercises: handshake+warm latency, happy-path payloads, the full error taxonomy +(hard tool errors vs soft per-item errors), session resolution, connection-pool +reuse, and concurrent calls from threads. +""" +import concurrent.futures +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from idatui.client import ( # noqa: E402 + IDAClient, + IDASessionError, + IDAToolError, +) + +PASS, FAIL = 0, 0 + + +def check(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f" ok {name}") + else: + FAIL += 1 + print(f" FAIL {name} {detail}") + + +def main(argv): + url = "http://127.0.0.1:8745/mcp" + db = os.environ.get("IDA_MCP_DB") + it = iter(argv) + for a in it: + if a == "--url": + url = next(it) + elif a == "--db": + db = next(it) + + ida = IDAClient(url, db=db) + ida.connect() + + print("[sessions]") + sessions = ida.list_sessions() + check("list_sessions returns >=1", len(sessions) >= 1, str(sessions)) + if db is None: + usable = [s for s in sessions if s.session_id] + if len(usable) == 1: + db = usable[0].session_id + else: + print(f" (multiple/zero sessions; pin one with --db) -> {usable}") + print(" Options:", ", ".join(f"{s.session_id}={s.filename}" for s in usable)) + return 2 + ida.set_db(db) + print(f" using db={db}") + + print("[handshake / latency]") + # Warm calls; measure. + for _ in range(3): + ida.health() + times = [] + for _ in range(20): + t = time.time() + ida.call("list_funcs", queries=[{"count": 50}]) + times.append((time.time() - t) * 1e3) + times.sort() + med = times[len(times) // 2] + print(f" list_funcs x20 warm: min={times[0]:.1f} med={med:.1f} max={times[-1]:.1f} ms") + check("warm call under 50ms median", med < 50, f"med={med:.1f}ms") + + print("[happy path]") + h = ida.health() + check("health is dict", isinstance(h, dict), type(h).__name__) + funcs = ida.call("list_funcs", queries=[{"count": 5}]) + check("list_funcs shape", isinstance(funcs, (list, dict)), type(funcs).__name__) + dis = ida.call("disasm", addr="main", max_instructions=10) + lines = dis.get("asm", {}).get("lines") if isinstance(dis, dict) else None + check("disasm main has lines", bool(lines), str(dis)[:120]) + check("disasm line carries addr", bool(lines and "addr" in lines[0]), + str(lines[0]) if lines else "no lines") + + print("[error taxonomy]") + # Hard tool error: wrong params -> isError true -> IDAToolError + try: + ida.call("xrefs_to", targets=["main"]) + check("bad params raises IDAToolError", False, "no exception") + except IDAToolError as e: + check("bad params raises IDAToolError", True) + check(" ...message surfaced", "addrs" in e.message or "param" in e.message.lower(), + e.message) + # Hard tool error: unknown tool -> isError true + try: + ida.call("definitely_not_a_tool") + check("unknown tool raises IDAToolError", False, "no exception") + except IDAToolError: + check("unknown tool raises IDAToolError", True) + # Soft/per-item error: bad addr to decompile -> isError false -> DATA, not raise + try: + payload = ida.call("decompile", addr="zzz_nope_addr") + soft = isinstance(payload, dict) and payload.get("error") + check("soft error returned as data (not raised)", bool(soft), str(payload)[:160]) + except IDAToolError as e: + check("soft error returned as data (not raised)", False, f"raised: {e}") + + print("[session resolution]") + tmp = IDAClient(url, db=None) + tmp.connect() + all_sessions = tmp.list_sessions() + usable = [s for s in all_sessions if s.session_id] + if len(usable) > 1: + try: + tmp.resolve_db() + check("multi-session resolve raises", False, "no exception") + except IDASessionError: + check("multi-session resolve raises", True) + else: + check("single-session auto-resolves", tmp.resolve_db() == usable[0].session_id) + tmp.close() + + print("[connection reuse]") + # Fire many calls; pool should keep idle conns bounded and all succeed. + ok = 0 + for _ in range(30): + if ida.call("list_funcs", queries=[{"count": 1}]) is not None: + ok += 1 + check("30 sequential calls all succeed (pooled)", ok == 30, f"{ok}/30") + + print("[concurrency]") + def worker(i): + return ida.call("disasm", addr="main", max_instructions=5) + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex: + results = list(ex.map(worker, range(40))) + good = sum(1 for r in results if isinstance(r, dict) and r.get("asm")) + check("40 concurrent calls across 8 threads", good == 40, f"{good}/40") + + ida.close() + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) |
