diff options
| author | blasty <blasty@local> | 2026-07-24 14:53:56 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-07-24 14:53:56 +0200 |
| commit | 22ffea91210610d1def042e230ffe8d223022b30 (patch) | |
| tree | a70c84470a0f4f46901cb221f2d0dd3859895b41 | |
| parent | mcp: collapse app + launcher to worker-only (diff) | |
| download | ida-tui-22ffea91210610d1def042e230ffe8d223022b30.tar.gz ida-tui-22ffea91210610d1def042e230ffe8d223022b30.tar.xz ida-tui-22ffea91210610d1def042e230ffe8d223022b30.zip | |
mcp: delete the ida-pro-mcp transport, supervisor, and mcp-only tests
The idalib worker is the only backend now, so remove the dead HTTP/supervisor
surface entirely (~2200 lines):
* deleted idatui/client.py (the IDAClient HTTP/JSON-RPC transport + session
manager), idatui/tui.py (the old mcp TUI entry, superseded by launch.py),
spawn.sh, and systemd/ (the supervisor unit).
* deleted the mcp-only tests (stress_client, smoke_client, test_keepalive,
stress_paging, rpc_smoke, serverctl.sh, pane_smoke, test_domain) -- the worker
pilot (tests/test_scenarios.py) supersedes them.
* migrated the tmux RPC harness (idatui/pane.py) to the worker: it spawns
`idatui.launch <binary> --rpc <sock>` instead of the mcp `idatui.tui`, drops
the supervisor auto-start/ensure machinery, and reaps our own worker
(idatui/worker.py) instead of ida_pro_mcp.idalib_server. --db/--url/--no-
ensure-server are gone; --open is required.
* __init__ / __main__ / domain no longer import client (exceptions come from
errors.py, the domain client hint is WorkerClient); pyproject points both
console scripts at idatui.launch; README + ida-tui header describe the
worker-only flow.
What stays (by design): the ida_pro_mcp *package* (the worker reuses its @tool
functions in-process) and server/patch_server.py (the worker injects its custom
tools on startup). Verified: whole package imports + IdaTui constructs + pilot
lists 31 scenarios. The worker pilot (134 pass / 2 known flakes) is the E2E gate.
| -rw-r--r-- | README.md | 67 | ||||
| -rwxr-xr-x | ida-tui | 11 | ||||
| -rw-r--r-- | idatui/__init__.py | 6 | ||||
| -rw-r--r-- | idatui/__main__.py | 6 | ||||
| -rw-r--r-- | idatui/client.py | 636 | ||||
| -rw-r--r-- | idatui/domain.py | 6 | ||||
| -rw-r--r-- | idatui/pane.py | 136 | ||||
| -rw-r--r-- | idatui/tui.py | 45 | ||||
| -rw-r--r-- | idatui/worker.py | 2 | ||||
| -rw-r--r-- | pyproject.toml | 2 | ||||
| -rwxr-xr-x | spawn.sh | 24 | ||||
| -rw-r--r-- | systemd/idatui-mcp.service | 34 | ||||
| -rwxr-xr-x | systemd/install.sh | 23 | ||||
| -rw-r--r-- | tests/pane_smoke.py | 87 | ||||
| -rw-r--r-- | tests/rpc_smoke.py | 284 | ||||
| -rwxr-xr-x | tests/serverctl.sh | 65 | ||||
| -rw-r--r-- | tests/smoke_client.py | 162 | ||||
| -rw-r--r-- | tests/stress_client.py | 290 | ||||
| -rw-r--r-- | tests/stress_paging.py | 150 | ||||
| -rw-r--r-- | tests/test_domain.py | 191 | ||||
| -rw-r--r-- | tests/test_keepalive.py | 76 |
21 files changed, 67 insertions, 2236 deletions
@@ -3,11 +3,11 @@ A minimal, keyboard-first (mouse-capable) **TUI frontend for IDA Pro**, built with [Textual](https://textual.textualize.io/) and driving **idalib** (IDA headless). -Opening a binary now spawns our own **idalib worker** — a private subprocess -talking a unix socket (`idatui/worker.py` + `WorkerClient`), ~50–100× cheaper per -call than the old transport. The -[ida-pro-mcp](https://github.com/mrexodia/ida-pro-mcp) HTTP path is **deprecated** -(kept behind `--backend mcp` / `--db` attach) and slated for removal. +Opening a binary spawns our own **idalib worker** — a private subprocess talking a +unix socket (`idatui/worker.py` + `WorkerClient`), ~50–100× cheaper per call than +an HTTP transport. It reuses [ida-pro-mcp](https://github.com/mrexodia/ida-pro-mcp)'s +tool implementations in-process; the old ida-pro-mcp HTTP server/supervisor path +has been **removed**. ## ⚠️ Status: not ready for public consumption @@ -51,45 +51,32 @@ pulls in Textual + Pygments. ## Requirements - Python ≥ 3.11 -- A working **IDA Pro** with **idalib** and **ida-pro-mcp** installed. +- A working **IDA Pro** with **idalib** and **ida-pro-mcp** installed (the worker + reuses ida-pro-mcp's tool implementations in-process — no server runs). - Textual ≥ 8 and Pygments ≥ 2 for the TUI (`pip install -e '.[tui]'`). +Two python environments are expected: one with **textual + idapro** for the TUI +(`~/ida-venv`, override `$IDATUI_PYTHON`) and one with **idapro + ida_pro_mcp** +for the worker (auto-detected, override `$IDATUI_WORKER_PYTHON`). + ## Running -The caveman way — one command does all the plumbing (starts the supervisor if -it's down, recovers a binary wedged by a crashed worker, opens/adopts the -session, launches the TUI): +One command — it spawns a private idalib worker for the binary (which opens + +auto-analyzes it in its own process over a unix socket) and drops you into the +TUI behind a loading overlay: ```sh ./ida-tui /path/to/binary # open a binary and drive it — that's it -./ida-tui # attach to the sole open session -./ida-tui --db <session> # attach to a specific session ``` It uses `~/ida-venv/bin/python` for the TUI (override with `$IDATUI_PYTHON`) and resolves binary paths against your real cwd. The binary's directory must be writable (idalib writes a `.i64` there). -The manual way (if you want the pieces separate): - -```sh -# 1. Start the ida-pro-mcp supervisor (opens bin/ls by default). -./spawn.sh # supervisor on 127.0.0.1:8745 - -# 2. Launch the TUI (use a python that has textual + idapro). -python -m idatui.tui # auto-resolve the sole session -python -m idatui.tui --db <session> -python -m idatui.tui --open /abs/path/bin # dir must be WRITABLE (.i64) -``` - -> Recovering a wedged database by hand: if a worker was hard-killed it leaves -> unpacked `foo.id0/.id1/.id2/.nam/.til` next to `foo.i64`, and the `.i64` then -> refuses to reopen. Delete those stale files (never the `.i64`) and retry — -> `ida-tui` does this automatically. - -The `spawn.sh` host/port/target are overridable via `IDA_MCP_HOST`, -`IDA_MCP_PORT`, `IDA_MCP_TARGET`, `IDA_MCP_MAX_WORKERS`. A systemd unit is in -`systemd/`. +> Recovering a wedged database: if a worker was hard-killed it leaves unpacked +> `foo.id0/.id1/.id2/.nam/.til` next to `foo.i64`, and the `.i64` then refuses to +> reopen. Delete those stale files (never the `.i64`) and retry — `ida-tui` does +> this automatically. ## RPC / driving the TUI @@ -97,22 +84,30 @@ Give the TUI `--rpc <sock>` to expose a unix-socket control channel, then drive it from another pane: ```sh -python -m idatui.tui --db <s> --rpc /tmp/ida.sock +./ida-tui /abs/path/bin --rpc /tmp/ida.sock python -m idatui.drive where # ergonomic terse-text helper python -m idatui.drive pc main # pseudocode of main python -m idatui.drive rename sub_5BE0 foo # goto + rename ``` +Or let `idatui.pane` spawn + manage TUI panes in tmux (see the idatui-rpc skill): + +```sh +python -m idatui.pane spawn --open /abs/path/bin # -> {sock, pane, ready} +python -m idatui.pane list +python -m idatui.pane stop --sock <sock> +``` + See `docs/RPC.md` for the full protocol. ## Tests -Headless Textual `Pilot` suites live in `tests/` and need a live session id: +A headless Textual `Pilot` suite lives in `tests/`; it spawns a worker on the +given binary (default `targets/echo`): ```sh -python tests/test_scenarios.py --db <session> # UI suite -python tests/test_scenarios.py --db <s> --only hex,rename -python tests/test_domain.py --db <session> # domain/paging (stdlib) +python tests/test_scenarios.py targets/echo # full UI suite +python tests/test_scenarios.py --only hex,rename ``` ## Docs @@ -2,14 +2,11 @@ # Caveman launcher for the IDA TUI: # # ./ida-tui foo.elf # open a binary and drive it — that's it -# ./ida-tui # attach to the sole open session (mcp) -# ./ida-tui --db <id> # attach to a specific session (mcp) # -# Opening a binary now spins up our own idalib worker (a unix-socket subprocess; -# no HTTP, no supervisor). --db/attach still use the deprecated ida-pro-mcp -# path; pass --backend mcp to force it. Uses the venv python that has textual -# (override with $IDATUI_PYTHON); the worker auto-picks the python that has -# ida_pro_mcp (override with $IDATUI_WORKER_PYTHON). +# Opening a binary spins up our own idalib worker (a unix-socket subprocess; +# no HTTP, no supervisor). Uses the venv python that has textual (override with +# $IDATUI_PYTHON); the worker auto-picks the python that has ida_pro_mcp +# (override with $IDATUI_WORKER_PYTHON). set -eu SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) diff --git a/idatui/__init__.py b/idatui/__init__.py index 92d8b5d..2cbdde8 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -1,4 +1,5 @@ -"""idatui — a minimal keyboard-first TUI for IDA Pro over ida-pro-mcp (idalib).""" +"""idatui — a minimal keyboard-first TUI for IDA Pro, driving idalib via a +private unix-socket worker (idatui.worker / WorkerClient).""" from .errors import ( IDAError, @@ -10,7 +11,6 @@ from .errors import ( IDASessionError, Session, ) -from .client import IDAClient, KeepAlive # deprecated mcp transport from .domain import ( Program, FunctionIndex, @@ -35,7 +35,6 @@ __all__ = [ "Decompilation", "LIST_PAGE", "DISASM_BLOCK", - "IDAClient", "IDAError", "IDAConnectionError", "IDATimeoutError", @@ -44,5 +43,4 @@ __all__ = [ "IDAToolError", "IDASessionError", "Session", - "KeepAlive", ] diff --git a/idatui/__main__.py b/idatui/__main__.py index 29f5cf4..0090ace 100644 --- a/idatui/__main__.py +++ b/idatui/__main__.py @@ -1,6 +1,6 @@ -"""``python -m idatui`` -> client self-check (until the TUI app lands).""" +"""``python -m idatui`` -> the one-shot launcher (open a binary in the TUI).""" import sys -from .client import _main +from .launch import main -raise SystemExit(_main(sys.argv[1:])) +raise SystemExit(main(sys.argv[1:])) 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:])) diff --git a/idatui/domain.py b/idatui/domain.py index 4f72f5b..d7ebf3c 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -32,8 +32,8 @@ from typing import Callable, TYPE_CHECKING from .errors import IDAToolError -if TYPE_CHECKING: # type hint only; the runtime client is mcp IDAClient or worker - from .client import IDAClient # noqa: F401 +if TYPE_CHECKING: # type hint only + from .worker_client import WorkerClient # noqa: F401 # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 @@ -797,7 +797,7 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: IDAClient, prefetch_workers: int = 2): + def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" diff --git a/idatui/pane.py b/idatui/pane.py index af81982..e34b822 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -8,9 +8,6 @@ then close it — all without a human touching the keyboard. python -m idatui.pane spawn --open /abs/path/to/bin # -> {"sock": "/run/user/1000/idatui-3f2a.sock", "pane": "%7", "ready": true, ...} - # or attach to an existing session id - python -m idatui.pane spawn --db 80d83396 - # drive it (see docs/RPC.md / the idatui-rpc skill) python -m idatui.rpcclient --sock <sock> pseudocode target=main @@ -18,9 +15,9 @@ then close it — all without a human touching the keyboard. python -m idatui.pane list python -m idatui.pane stop --sock <sock> # graceful quit + kill pane -Requires: running inside tmux, and the ida-pro-mcp supervisor already up -(./spawn.sh). Uses ~/ida-venv/bin/python for the TUI (needs textual) unless ---python / IDATUI_PYTHON says otherwise. +Requires: running inside tmux. Each pane spawns its own private idalib worker +(no shared supervisor). Uses ~/ida-venv/bin/python for the TUI (needs textual) +unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -29,14 +26,11 @@ import json import os import secrets import signal -import socket import subprocess import sys import time from typing import Any -from urllib.parse import urlparse -from .client import DEFAULT_URL from .rpcclient import RpcClient, RpcError REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -81,24 +75,17 @@ def _tmux(*args: str) -> str: # --------------------------------------------------------------------------- # # idalib worker reaping # -# ``pane stop`` kills the TUI pane, but the ida-pro-mcp supervisor does not -# reliably reap the ``ida_pro_mcp.idalib_server`` worker it forked for that -# session. Leaked workers accumulate against IDA_MCP_MAX_WORKERS until the next -# ``spawn`` blocks forever waiting for a free slot (the TUI comes up but never -# becomes ready). A worker is only *safe* to reap when no idatui pane is live -# (then every worker is orphaned) — that mirrors the hand workaround -# `pkill -f ida_pro_mcp.idalib_server` and avoids killing an in-use analyser. +# ``pane stop`` kills the TUI pane, but a hard-killed pane can leave its private +# idalib worker (idatui/worker.py) running. A worker is only *safe* to reap when +# no idatui pane is live (then every worker is orphaned), which avoids killing an +# in-use analyser. # --------------------------------------------------------------------------- # -_WORKER_PATTERN = r"ida_pro_mcp\.idalib_server" +_WORKER_PATTERN = r"idatui/worker\.py" def _worker_pids() -> list[int]: - """PIDs of the supervisor's per-binary idalib worker processes. - - Matches the worker module invocation only (not the ``idalib-mcp`` - supervisor, whose command line does not contain the module path), and - never our own PID. - """ + """PIDs of our private per-pane idalib worker processes (idatui/worker.py), + never our own PID.""" try: out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], capture_output=True, text=True) @@ -138,73 +125,19 @@ def _reap_orphan_workers(force: bool = False) -> int: # --------------------------------------------------------------------------- # -# supervisor (ida-pro-mcp server) — auto-start if down -# --------------------------------------------------------------------------- # -def _server_addr(url: str) -> tuple[str, int]: - u = urlparse(url) - return (u.hostname or "127.0.0.1", u.port or 8745) - - -def _server_up(host: str, port: int, timeout: float = 0.75) -> bool: - """Is something listening on host:port? (Cheap TCP probe; the readiness poll - that follows catches a half-up server.)""" - try: - with socket.create_connection((host, port), timeout=timeout): - return True - except OSError: - return False - - -def _ensure_server(host: str, port: int, timeout: float, - detached: bool = True) -> dict[str, Any]: - """Make sure the supervisor is up; start ./spawn.sh in a tmux pane if not. - - Only auto-starts a *local* server (can't launch a remote one). spawn.sh binds - the port, so starting a duplicate is impossible — the probe guards that. - ``IDATUI_SERVER_CMD`` overrides the launch command (used by the tests). - """ - if _server_up(host, port): - return {"server_started": False, "server_up": True} - if host not in ("127.0.0.1", "localhost", "::1"): - return {"server_started": False, "server_up": False, - "error": f"server at {host}:{port} is down and not local; " - "cannot auto-start"} - cmd_str = os.environ.get("IDATUI_SERVER_CMD", "./spawn.sh") - cmd = f"cd {_q(REPO)} && exec {cmd_str}" - split = ["split-window", "-v", "-P", "-F", "#{pane_id}"] - if detached: - split += ["-d"] - anchor = os.environ.get("TMUX_PANE") - if anchor: - split += ["-t", anchor] - split.append(cmd) - pane = _tmux(*split) - deadline = time.time() + timeout - while time.time() < deadline: - if not _pane_alive(pane): - return {"server_started": True, "server_up": False, "server_pane": pane, - "error": "supervisor pane exited during startup (check it)"} - if _server_up(host, port): - return {"server_started": True, "server_up": True, "server_pane": pane} - time.sleep(0.5) - return {"server_started": True, "server_up": False, "server_pane": pane, - "error": "supervisor did not come up in time"} - - -# --------------------------------------------------------------------------- # # spawn # --------------------------------------------------------------------------- # def spawn(args) -> int: if not os.environ.get("TMUX"): print("error: not inside tmux (spawn creates a tmux pane)", file=sys.stderr) return 2 - if not args.open and not args.db: - print("error: pass --open <binary> or --db <session>", file=sys.stderr) + if not args.open: + print("error: pass --open <binary>", file=sys.stderr) return 2 sock = args.sock or os.path.join(_sockdir(), f"idatui-{secrets.token_hex(3)}.sock") - target = os.path.abspath(os.path.expanduser(args.open)) if args.open else args.db - if args.open and not os.path.exists(target): + target = os.path.abspath(os.path.expanduser(args.open)) + if not os.path.exists(target): print(f"error: no such binary: {target}", file=sys.stderr) return 2 @@ -216,26 +149,9 @@ def spawn(args) -> int: print(f"reaped {reaped} orphaned idalib worker(s) before spawn", file=sys.stderr) - # make sure the ida-pro-mcp supervisor is up (auto-start it if not) - srv: dict[str, Any] = {"server_started": False, "server_up": True} - if not args.no_ensure_server: - host, port = _server_addr(args.url or DEFAULT_URL) - srv = _ensure_server(host, port, args.server_timeout) - if srv.get("server_started"): - print(f"supervisor was down — started it ({srv.get('server_pane')})", - file=sys.stderr) - if not srv.get("server_up"): - print(json.dumps({"ready": False, **srv}), file=sys.stderr) - return 3 - - # the command the pane runs: become the TUI so kill-pane kills it cleanly - inner = [args.python, "-m", "idatui.tui", "--rpc", sock] - if args.open: - inner += ["--open", target] - else: - inner += ["--db", target] - if args.url: - inner += ["--url", args.url] + # the command the pane runs: the launcher spawns a private idalib worker for + # this binary and becomes the TUI, so kill-pane tears the whole thing down. + inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) split = ["split-window", "-v" if args.vertical else "-h", @@ -251,10 +167,7 @@ def spawn(args) -> int: pane = _tmux(*split) row = {"sock": sock, "pane": pane, "target": target, - "kind": "open" if args.open else "db", "started": time.time(), - "server_started": srv.get("server_started", False)} - if srv.get("server_pane"): - row["server_pane"] = srv["server_pane"] + "kind": "open", "started": time.time()} reg = [r for r in _load_registry() if r.get("sock") != sock] reg.append(row) _save_registry(reg) @@ -298,8 +211,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, why = ("RPC socket not created yet" if not os.path.exists(sock) else "TUI up but analysis not ready") print(f"still waiting ({int(time.time() - start)}s): {why}. If this " - f"hangs, the idalib worker may be stuck or IDA_MCP_MAX_WORKERS " - f"is full — try `python -m idatui.pane reap`.", file=sys.stderr) + f"hangs, the idalib worker may be stuck — try " + f"`python -m idatui.pane reap`.", file=sys.stderr) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -388,15 +301,10 @@ def main(argv: list[str]) -> int: sub = p.add_subparsers(dest="cmd", required=True) sp = sub.add_parser("spawn", help="open a TUI pane and wait until ready") - sp.add_argument("--open", metavar="PATH", help="binary to open (dir must be writable)") - sp.add_argument("--db", metavar="SESSION", help="attach to an existing session id") + sp.add_argument("--open", metavar="PATH", required=True, + help="binary to open (its dir must be writable)") sp.add_argument("--sock", help="RPC socket path (default: auto in $XDG_RUNTIME_DIR)") sp.add_argument("--python", default=DEFAULT_PY, help=f"python for the TUI ({DEFAULT_PY})") - sp.add_argument("--url", help="MCP server URL (default: idatui's default)") - sp.add_argument("--no-ensure-server", action="store_true", - help="don't auto-start ./spawn.sh if the supervisor is down") - sp.add_argument("--server-timeout", type=float, default=90.0, - help="seconds to wait for an auto-started supervisor") sp.add_argument("--vertical", action="store_true", help="split vertically (stacked)") sp.add_argument("--size", help="new pane size (tmux -l value, e.g. 60%% or 120)") sp.add_argument("--detached", action="store_true", help="don't focus the new pane") diff --git a/idatui/tui.py b/idatui/tui.py deleted file mode 100644 index b399ffe..0000000 --- a/idatui/tui.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Launcher for the idatui TUI. - - # attach to the single open session on a running server - python -m idatui.tui - - # attach to a specific session - python -m idatui.tui --db 80d83396 - - # open (or reopen) an arbitrary binary, then drive it - python -m idatui.tui --open /path/to/binary - -The server (supervisor) must already be running (see spawn.sh). --open creates a -session via idb_open; the binary's directory must be writable (idalib writes a -.i64 next to it). -""" -from __future__ import annotations - -import argparse -import os - -from .app import IdaTui -from .client import DEFAULT_URL - - -def main(argv: list[str] | None = None) -> int: - p = argparse.ArgumentParser(prog="idatui", description="Minimal TUI for IDA over MCP") - p.add_argument("--url", default=os.environ.get("IDA_MCP_URL", DEFAULT_URL), - help=f"MCP server URL (default {DEFAULT_URL})") - p.add_argument("--db", default=os.environ.get("IDA_MCP_DB"), - help="attach to an existing session id") - p.add_argument("--open", metavar="PATH", - help="open (or reopen) a binary and drive it (dir must be writable)") - p.add_argument("--no-keepalive", action="store_true", - help="Do not bump idle-TTL / run the keepalive heartbeat") - p.add_argument("--rpc", metavar="PATH", - help="listen for RPC on this unix socket path (puppeteer the TUI)") - args = p.parse_args(argv) - rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None - IdaTui(url=args.url, db=args.db, open_path=args.open, - keepalive=not args.no_keepalive, rpc_path=rpc_path).run() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/idatui/worker.py b/idatui/worker.py index 64cafb3..bfefa4a 100644 --- a/idatui/worker.py +++ b/idatui/worker.py @@ -62,7 +62,7 @@ def recv(sock: socket.socket): 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 - without spawn.sh having run server/patch_server.py first. Must run BEFORE + (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__))) diff --git a/pyproject.toml b/pyproject.toml index 9adde5d..30f8cc2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ tui = ["textual>=8", "pygments>=2"] # pygments ships with rich; explicit for th dev = ["pytest>=8"] [project.scripts] -idatui = "idatui.tui:main" +idatui = "idatui.launch:main" ida-tui = "idatui.launch:main" # caveman one-shot: `ida-tui foo.elf` [build-system] diff --git a/spawn.sh b/spawn.sh deleted file mode 100755 index a8763ae..0000000 --- a/spawn.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/sh -# Start the ida-pro-mcp supervisor. IDA_MCP_MAX_WORKERS raises the ceiling on -# concurrently-open binaries (default 4); idle workers self-exit and free slots. -# -# Location-independent: resolves paths relative to this script, not $PWD, so it -# behaves the same whether run by hand or from the systemd unit (see -# systemd/idatui-mcp.service). -set -eu - -SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) - -IDA_MCP_HOST="${IDA_MCP_HOST:-127.0.0.1}" -IDA_MCP_PORT="${IDA_MCP_PORT:-8745}" -IDA_MCP_TARGET="${IDA_MCP_TARGET:-$SCRIPT_DIR/bin/ls}" - -# Ensure idatui's extra server tools (del_type, needed by the struct editor) are -# present in the installed ida-pro-mcp. Idempotent; patches the api_types.py that -# the /usr/bin/python workers import. See server/patch_server.py. -/usr/bin/python "$SCRIPT_DIR/server/patch_server.py" || \ - echo "warn: server patch skipped (del_type may be unavailable)" - -cd "$SCRIPT_DIR" -exec env IDA_MCP_MAX_WORKERS="${IDA_MCP_MAX_WORKERS:-8}" \ - uv run idalib-mcp --host "$IDA_MCP_HOST" --port "$IDA_MCP_PORT" "$IDA_MCP_TARGET" diff --git a/systemd/idatui-mcp.service b/systemd/idatui-mcp.service deleted file mode 100644 index 9afcb73..0000000 --- a/systemd/idatui-mcp.service +++ /dev/null @@ -1,34 +0,0 @@ -[Unit] -Description=idatui ida-pro-mcp supervisor (idalib-mcp on 127.0.0.1:8745) -Documentation=https://github.com/mrexodia/ida-pro-mcp -# Wait for network in case a worker ever needs to resolve/pull anything. -After=network-online.target -Wants=network-online.target -# If it dies >5 times in 60s, stop flapping and surface a hard failure. -StartLimitIntervalSec=60 -StartLimitBurst=5 - -[Service] -Type=simple -WorkingDirectory=%h/dev/ida-tui-maybe -ExecStart=%h/dev/ida-tui-maybe/spawn.sh - -# Be explicit about PATH so the unit doesn't depend on an interactive login -# shell. %h/.local/bin is where `idalib-mcp` lives (uv tool install); uv itself -# and /usr/bin/python are on the system PATH. -Environment=PATH=%h/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -# Tunables (override with a drop-in: systemctl --user edit idatui-mcp). -Environment=IDA_MCP_MAX_WORKERS=8 -Environment=IDA_MCP_HOST=127.0.0.1 -Environment=IDA_MCP_PORT=8745 - -# Always bring it back; observe restarts/exits in `journalctl --user -u idatui-mcp`. -Restart=always -RestartSec=2 - -# Give the supervisor time to shut its workers down cleanly. -TimeoutStopSec=30 -KillMode=mixed - -[Install] -WantedBy=default.target diff --git a/systemd/install.sh b/systemd/install.sh deleted file mode 100755 index 0a118b7..0000000 --- a/systemd/install.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/sh -# Install & enable the idatui-mcp user service. Re-run after editing the unit. -set -eu - -UNIT=idatui-mcp.service -SRC=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)/$UNIT -DEST_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user" - -mkdir -p "$DEST_DIR" -ln -sf "$SRC" "$DEST_DIR/$UNIT" - -# Keep the service alive after logout / across reboots without a login session. -loginctl enable-linger "$(id -un)" || true - -systemctl --user daemon-reload -systemctl --user enable --now "$UNIT" - -echo -echo "Installed. Handy commands:" -echo " systemctl --user status idatui-mcp" -echo " systemctl --user restart idatui-mcp" -echo " journalctl --user -u idatui-mcp -f" -echo " systemctl --user edit idatui-mcp # drop-in overrides (ports, workers)" diff --git a/tests/pane_smoke.py b/tests/pane_smoke.py deleted file mode 100644 index bd7def6..0000000 --- a/tests/pane_smoke.py +++ /dev/null @@ -1,87 +0,0 @@ -#!/usr/bin/env python3 -"""Smoke test for idatui.pane's supervisor auto-start machinery (no IDA needed). - - python3 tests/pane_smoke.py - -Uses IDATUI_SERVER_CMD to launch a dummy port-binder in a tmux pane instead of -the real ./spawn.sh, so we can exercise _ensure_server end-to-end (start, detect, -idempotency, remote guard) without a real ida-pro-mcp server. Must run in tmux. -""" -import os -import socket -import subprocess -import sys -import tempfile - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui import pane # noqa: E402 - -PASS = FAIL = 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 _free_port() -> int: - s = socket.socket() - s.bind(("127.0.0.1", 0)) - p = s.getsockname()[1] - s.close() - return p - - -def main() -> int: - if not os.environ.get("TMUX"): - print("error: must run inside tmux", file=sys.stderr) - return 2 - - # a dummy "supervisor": bind the port and idle, so _server_up sees it. - fake = tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) - fake.write( - "import socket,sys,time\n" - "s=socket.socket(); s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)\n" - "s.bind(('127.0.0.1',int(sys.argv[1]))); s.listen(); time.sleep(300)\n") - fake.close() - port = _free_port() - os.environ["IDATUI_SERVER_CMD"] = f"{sys.executable} {fake.name} {port}" - - server_pane = None - try: - check("port starts down", not pane._server_up("127.0.0.1", port)) - - srv = pane._ensure_server("127.0.0.1", port, timeout=15.0) - server_pane = srv.get("server_pane") - check("ensure_server starts the supervisor and it comes up", - srv.get("server_started") and srv.get("server_up") and server_pane, - str(srv)) - check("port is now up", pane._server_up("127.0.0.1", port)) - - srv2 = pane._ensure_server("127.0.0.1", port, timeout=5.0) - check("ensure_server is idempotent when already up", - srv2.get("server_started") is False and srv2.get("server_up") is True, - str(srv2)) - - srv3 = pane._ensure_server("10.255.255.1", 9, timeout=2.0) - check("remote+down server is not auto-started", - srv3.get("server_started") is False and "local" in (srv3.get("error") or ""), - str(srv3)) - finally: - if server_pane: - subprocess.run(["tmux", "kill-pane", "-t", server_pane], - capture_output=True) - os.unlink(fake.name) - os.environ.pop("IDATUI_SERVER_CMD", None) - - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/rpc_smoke.py b/tests/rpc_smoke.py deleted file mode 100644 index a38162e..0000000 --- a/tests/rpc_smoke.py +++ /dev/null @@ -1,284 +0,0 @@ -#!/usr/bin/env python3 -"""End-to-end smoke test for the RPC server, in-process over a real unix socket. - - ~/ida-venv/bin/python tests/rpc_smoke.py --db <session_id> - -Boots the TUI headless with --rpc on a temp socket, connects a JSONL client over -that socket (same loop), and exercises the raw + introspection primitives. -""" -import asyncio -import json -import os -import sys -import tempfile - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.app import IdaTui # noqa: E402 -from idatui.client import DEFAULT_URL # noqa: E402 -from idatui._sync import wait_for # noqa: E402 -from textual.widgets import DataTable # noqa: E402 - -PASS = FAIL = 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}") - - -class Conn: - """Tiny JSONL client over a unix socket.""" - def __init__(self, r, w): - self.r, self.w = r, w - self._id = 0 - - async def call(self, method, **params): - self._id += 1 - self.w.write(json.dumps({"id": self._id, "method": method, - "params": params}).encode() + b"\n") - await self.w.drain() - line = await self.r.readline() - return json.loads(line.decode()) - - -async def run(db): - sock = os.path.join(tempfile.gettempdir(), f"idatui-rpc-{os.getpid()}.sock") - url = os.environ.get("IDA_MCP_URL", DEFAULT_URL) - app = IdaTui(url=url, db=db, keepalive=False, rpc_path=sock) - async with app.run_test(size=(140, 44)) as pilot: - # boot: functions loaded + socket up - await wait_for(lambda: app.query_one("#func-table", DataTable).row_count > 0, - pilot.pause, 60) - await wait_for(lambda: os.path.exists(sock), pilot.pause, 20) - if app._func_index is not None and not app._func_index.complete: - app._func_index.load_all() - - r, w = await asyncio.open_unix_connection(sock) - c = Conn(r, w) - - resp = await c.call("ping") - pong = resp.get("result", {}) - check("ping returns proto + readiness", - pong.get("proto") == 1 and pong.get("ready") is True and pong.get("functions", 0) > 0, - str(pong)) - - resp = await c.call("methods") - check("methods lists the verb table", - isinstance(resp.get("result"), dict) and "pseudocode" in resp["result"], - str(resp)[:80]) - - resp = await c.call("functions", limit=400) - funcs = resp.get("result", []) - check("functions lists entries", isinstance(funcs, list) and len(funcs) > 0, - str(resp)[:120]) - # a normally-named function (its name appears verbatim in its own decomp, - # unlike e.g. '.init_proc' which IDA renders as 'init_proc') - target = next((f for f in funcs if f["name"] == "main"), - next((f for f in funcs if f["name"].startswith("sub_")), funcs[0])) - - resp = await c.call("state") - st = resp.get("result", {}) - check("state has an active view", st.get("active") in ("decomp", "disasm", "hex"), - str(st)[:120]) - - # drive it like a user: goto a function by name via raw keys - keys = ["g"] + list(target["name"]) + ["enter"] - resp = await c.call("keys", keys=keys) - st = resp.get("result", {}) - check("keys(goto) navigates to the function", - st.get("function", {}).get("name") == target["name"], - f"got {st.get('function')}") - - resp = await c.call("view", lines=6) - v = resp.get("result", {}) - check("view returns visible lines with a cursor", - isinstance(v.get("lines"), list) and any(l.get("cur") for l in v["lines"]), - str(v)[:120]) - - resp = await c.call("screen") - scr = resp.get("result", {}) - check("screen returns a text grid", - isinstance(scr.get("text"), str) and target["name"] in scr["text"], - f"len={len(scr.get('text','')) if scr else 0}") - - # raw text primitive into the goto input, then escape out - from textual.widgets import Input - await c.call("keys", keys=["g"]) - await c.call("text", text="sub_", settle=False) - val = app.query_one("#goto", Input).value - check("text primitive types into the focused input", val == "sub_", f"val={val!r}") - await c.call("keys", keys=["escape"]) - - # --- semantic verbs ------------------------------------------------ # - resp = await c.call("goto", target=target["name"], delay_ms=0) - check("semantic goto lands on the function", - resp.get("result", {}).get("function", {}).get("name") == target["name"], - str(resp.get("result", {}).get("function"))) - - before = app._active - resp = await c.call("toggle_view") - check("toggle_view flips the active pane", - resp.get("result", {}).get("active") != before, f"still {before}") - await c.call("toggle_view") # flip back to a known state - - # fast movement: cursor should advance a few lines - line0 = (await c.call("state"))["result"]["cursor"].get("line") - resp = await c.call("move", dir="down", n=4) - line1 = resp["result"]["cursor"].get("line") - check("move(down,4) advances the cursor", - isinstance(line0, int) and isinstance(line1, int) and line1 > line0, - f"{line0} -> {line1}") - - # cursor_on: place the cursor on the function's own name by token - resp = await c.call("cursor_on", word=target["name"]) - cur = resp.get("result", {}) - check("cursor_on lands on the named token", - cur.get("found") is True and cur.get("cursor", {}).get("word") == target["name"], - str(cur.get("cursor"))) - - # rename using the ergonomic word= (cursor_on + prompt-fill in one call) - v = (await c.call("view", lines=1))["result"] - line0_text = v["lines"][0]["text"] if v.get("lines") else "" - col = line0_text.find(target["name"]) - if col >= 0: - newname = f"rpc_{os.getpid()}" - await c.call("rename", name=newname, word=target["name"], delay_ms=0) - got = app._func_index.by_addr(target["ea"]) - check("rename word= updates the function name", - got is not None and got.name == newname, - got.name if got else None) - app.program.client.call( - "rename", batch={"func": {"addr": hex(target["ea"]), "name": target["name"]}}) - else: - check("found the function name to rename", False, repr(line0_text[:60])) - - # --- structured introspection ------------------------------------- # - resp = await c.call("resolve", name=target["name"]) - check("resolve maps a name to its ea", - resp.get("result", {}).get("ea") == target["ea"], str(resp.get("result"))) - - resp = await c.call("pseudocode", target=target["name"]) - pc = resp.get("result", {}) - check("pseudocode returns the full body", - isinstance(pc.get("code"), str) and len(pc["code"]) > 0 and not pc["failed"], - f"failed={pc.get('failed')} len={len(pc.get('code') or '')}") - - resp = await c.call("disassembly", target=target["name"], max=50) - da = resp.get("result", {}) - check("disassembly returns lines with addresses", - isinstance(da.get("lines"), list) and len(da["lines"]) > 0 - and all("ea" in ln and "text" in ln for ln in da["lines"]), - f"total={da.get('total')} n={len(da.get('lines', []))}") - - # a function that is actually referenced, so xrefs_to is non-empty - callee = None - for f in funcs: - xr = (await c.call("xrefs_to", target=f["name"], limit=5)).get("result", []) - if xr: - callee = (f, xr) - break - check("xrefs_to returns structured references", - callee is not None and all("frm" in x for x in callee[1]), - "no referenced function found" if callee is None else "") - - # xrefs_from a function is whole-body (decomp refs), not just the entry: - # find a function that actually calls something. - caller = next((f["name"] for f in funcs if f["name"] == "main"), None) - xf = [] - for name in ([caller] if caller else []) + [f["name"] for f in funcs]: - xf = (await c.call("xrefs_from", target=name)).get("result", []) - if any(x.get("is_func") for x in xf): - caller = name - break - check("xrefs_from a function lists whole-body callees", - isinstance(xf, list) and len(xf) > 1 - and any(x.get("is_func") for x in xf) and all("to" in x for x in xf), - f"caller={caller} n={len(xf)}") - - # --- modal select: open xrefs on the callee, pick the first site --- # - if callee is not None: - f, xr = callee - await c.call("goto", target=f["name"], delay_ms=0) - # place the cursor on the function name so xrefs targets it - v = (await c.call("view", lines=1))["result"] - lt = v["lines"][0]["text"] if v.get("lines") else "" - col = lt.find(f["name"]) - if col >= 0: - await c.call("cursor", line=0, col=col + 1) - resp = await c.call("xrefs") - m = resp.get("result", {}).get("modal") or {} - check("xrefs opens the picker with items", - m.get("kind") == "XrefsScreen" and len(m.get("items", [])) > 0, - str(m)[:80]) - resp = await c.call("select", index=0) - check("select follows a picked xref (modal closes, we navigate)", - (resp.get("result", {}).get("modal") is None), str(resp.get("result", {}).get("modal"))) - - # --- in-view search ----------------------------------------------- # - await c.call("goto", target=target["name"], delay_ms=0) - resp = await c.call("search", term="return", delay_ms=0) - st = resp.get("result", {}) - check("search runs without error and returns state", - st.get("active") in ("decomp", "disasm"), str(st.get("active"))) - await c.call("keys", keys=["escape"]) - - # colored screen export (for an out-of-band web viewer) - resp = await c.call("screen", format="html") - html = resp.get("result", {}) - check("screen format=html returns an html document", - html.get("format") == "html" and "<" in (html.get("text") or ""), - str(html.get("format"))) - - # single-driver gate: a 2nd connection (while c is open) is refused - r2, w2 = await asyncio.open_unix_connection(sock) - c2 = Conn(r2, w2) - resp2 = await c2.call("ping") - check("second concurrent client is refused (single-driver)", - resp2.get("error") and "busy" in resp2["error"].get("message", ""), - str(resp2)) - w2.close() - - try: - w.close() - except Exception: # noqa: BLE001 - pass - - # --- ergonomic driver (idatui.drive) over the same socket ---------- # - # (c is closed now, so the single-driver gate lets drive connect.) - from idatui import drive - loop = asyncio.get_running_loop() - for cmd in (["where"], ["pc", target["name"]], ["callees", target["name"]], - ["names", "sub_", "3"]): - rc = await loop.run_in_executor( - None, lambda a=cmd: drive.main(["--sock", sock, *a])) - check(f"drive {cmd[0]} runs against the socket", rc == 0, f"rc={rc}") - - # --- graceful quit (last: it tears the app down) ------------------- # - r2, w2 = await asyncio.open_unix_connection(sock) - c2 = Conn(r2, w2) - resp = await c2.call("quit") - check("quit acknowledges before shutting down", - resp.get("result", {}).get("quitting") is True, str(resp)) - w2.close() - await wait_for(lambda: not app.is_running, pilot.pause, 10) - check("quit actually exits the app", not app.is_running, "still running") - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -def main(argv): - db = None - it = iter(argv) - for a in it: - if a == "--db": - db = next(it) - return asyncio.run(run(db)) - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/tests/serverctl.sh b/tests/serverctl.sh deleted file mode 100755 index a5bc098..0000000 --- a/tests/serverctl.sh +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env bash -# Control the ida-pro-mcp server for stress testing. Assumes ~/ida-venv + spawn.sh. -set -u -REPO="$(cd "$(dirname "$0")/.." && pwd)" -PORT=8745 -LOG=/tmp/ida-stress-spawn.log - -wait_port() { # wait_port <up|down> <secs> - local want="$1" secs="${2:-40}" i - for ((i=0; i<secs*2; i++)); do - if ss -ltn 2>/dev/null | grep -q ":$PORT "; then - [ "$want" = up ] && return 0 - else - [ "$want" = down ] && return 0 - fi - sleep 0.5 - done - return 1 -} - -wait_ready() { # wait until hexrays_ready via the client - local secs="${1:-60}" i - for ((i=0; i<secs; i++)); do - if python3 -c " -import sys; sys.path.insert(0,'$REPO') -from idatui.client import IDAClient -try: - c=IDAClient(); c.connect() - ss=[s for s in c.list_sessions() if s.session_id] - if ss: - c.set_db(ss[0].session_id) - if c.health().get('hexrays_ready'): print(ss[0].session_id); sys.exit(0) -except Exception: pass -sys.exit(1) -" 2>/dev/null; then return 0; fi - sleep 1 - done - return 1 -} - -case "${1:-}" in - start) - cd "$REPO" - source ~/ida-venv/bin/activate 2>/dev/null - nohup ./spawn.sh >"$LOG" 2>&1 & - wait_port up 40 || { echo "PORT_TIMEOUT"; tail -5 "$LOG"; exit 1; } - wait_ready 90 || { echo "READY_TIMEOUT"; tail -5 "$LOG"; exit 1; } - ;; - stop) - # Kill the whole tree: uv wrapper, supervisor, worker. - pkill -9 -f 'idalib_server' 2>/dev/null - pkill -9 -f 'idalib-mcp' 2>/dev/null - pkill -9 -f 'uv run idalib' 2>/dev/null - wait_port down 20 || { echo "STOP_TIMEOUT"; exit 1; } - ;; - kill9) - # Hard kill only the worker (simulate a crash of the analysis process). - pkill -9 -f 'idalib_server' 2>/dev/null - ;; - ready) - wait_ready "${2:-60}" - ;; - *) - echo "usage: $0 {start|stop|kill9|ready [secs]}"; exit 2;; -esac diff --git a/tests/smoke_client.py b/tests/smoke_client.py deleted file mode 100644 index 5466284..0000000 --- a/tests/smoke_client.py +++ /dev/null @@ -1,162 +0,0 @@ -#!/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 _query_data(payload): - res = payload.get("result", payload) if isinstance(payload, dict) else payload - if isinstance(res, list): - res = res[0] if res else {} - return res.get("data", []) if isinstance(res, dict) else [] - - -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__) - # Pick a real function (don't assume 'main' exists — libraries have none). - some = ida.call("list_funcs", queries=[{"filter": "sub_*", "count": 1}]) - target = _query_data(some)[0]["addr"] - dis = ida.call("disasm", addr=target, max_instructions=10) - lines = (dis.get("asm") or {}).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=target, 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:])) diff --git a/tests/stress_client.py b/tests/stress_client.py deleted file mode 100644 index 955700c..0000000 --- a/tests/stress_client.py +++ /dev/null @@ -1,290 +0,0 @@ -#!/usr/bin/env python3 -"""Adversarial stress tests for idatui.client. - -Kills/respawns the server, forces timeouts, drops connections, and hammers with -concurrency — asserting the client always fails *cleanly* (no hangs, no -deadlocks) and recovers where recovery is possible. - - python3 tests/stress_client.py # run all scenarios - python3 tests/stress_client.py timeout # run one by name - -Requires ~/ida-venv + spawn.sh (via tests/serverctl.sh). Leaves the server UP. -""" -import concurrent.futures -import os -import subprocess -import sys -import threading -import time - -HERE = os.path.dirname(os.path.abspath(__file__)) -REPO = os.path.dirname(HERE) -sys.path.insert(0, REPO) - -from idatui.client import ( # noqa: E402 - IDAClient, - IDAConnectionError, - IDASessionError, - IDATimeoutError, - IDAToolError, - IDAError, -) - -URL = "http://127.0.0.1:8745/mcp" -CTL = os.path.join(HERE, "serverctl.sh") -PASS = FAIL = 0 - - -def log(m): - print(m, flush=True) - - -def check(name, cond, detail=""): - global PASS, FAIL - if cond: - PASS += 1 - log(f" ok {name}") - else: - FAIL += 1 - log(f" FAIL {name} {detail}") - - -def ctl(*args, timeout=180): - r = subprocess.run([CTL, *args], capture_output=True, text=True, timeout=timeout) - return r.returncode, (r.stdout + r.stderr).strip() - - -def ensure_up(): - rc, out = ctl("ready", "5") - if rc != 0: - log(" (server down; starting...)") - rc, out = ctl("start") - assert rc == 0, f"could not start server: {out}" - rc, sid = ctl("ready", "60") - return sid.strip().splitlines()[-1] if sid.strip() else None - - -def deadline(fn, secs): - """Run fn() in a thread; return (finished_bool, result_or_exc).""" - box = {} - def run(): - try: - box["r"] = fn() - except BaseException as e: # noqa: BLE001 - box["e"] = e - t = threading.Thread(target=run, daemon=True) - t.start() - t.join(secs) - if t.is_alive(): - return False, None - return True, box.get("e", box.get("r")) - - -# --------------------------------------------------------------------------- # -def s_connect_refused(): - """Connecting to a dead port fails fast and cleanly, never hangs.""" - c = IDAClient("http://127.0.0.1:59999/mcp", db="x", timeout=3.0) - fin, res = deadline(c.connect, 8) - check("dead-port connect returns (no hang)", fin) - check("dead-port raises IDAConnectionError", isinstance(res, IDAConnectionError), - f"{type(res).__name__}: {res}") - - -def s_timeout(): - """A too-tight timeout on the slow decompile raises IDATimeoutError, and the - client stays usable for subsequent calls.""" - sid = ensure_up() - c = IDAClient(URL, db=sid) - c.connect() - # force_recompile clears the hexrays cache so decompile is genuinely slow. - try: - c.call("force_recompile", addr="main") - except IDAError: - pass - fin, res = deadline(lambda: c.call("decompile", addr="main", timeout=0.001), 10) - check("tight-timeout returns (no hang)", fin) - check("tight-timeout raises IDATimeoutError", isinstance(res, IDATimeoutError), - f"{type(res).__name__}: {res}") - # Client must still work afterwards (pool not poisoned). - fin2, res2 = deadline(lambda: c.call("list_funcs", queries=[{"count": 1}]), 10) - check("client usable after timeout", fin2 and not isinstance(res2, Exception), - f"{type(res2).__name__}: {res2}") - c.close() - - -def s_worker_crash(): - """Hard-kill the analysis worker mid-use: calls must raise cleanly (no hang), - and after a full restart the client recovers by re-resolving the session.""" - sid = ensure_up() - c = IDAClient(URL, db=sid) - c.connect() - check("baseline call ok", isinstance(c.call("list_funcs", queries=[{"count": 1}]), (list, dict))) - - ctl("kill9") # kill only the worker; supervisor may stay up - time.sleep(1.0) - - def probe(): - try: - return c.call("list_funcs", queries=[{"count": 1}]) - except IDAError as e: - return e - fin, res = deadline(probe, 15) - check("post-crash call returns (no hang)", fin, "call hung after worker kill") - check("post-crash raises IDAError (clean)", isinstance(res, IDAError) or res is not None, - f"{type(res).__name__}: {res}") - - # Full restart and recovery. - ctl("stop") - rc, _ = ctl("start") - check("server restart ok", rc == 0) - newsid = ensure_up() - # A fresh client should just work. - c2 = IDAClient(URL, db=None) - c2.connect() - fin3, res3 = deadline(lambda: (c2.set_db(c2.resolve_db()), c2.health())[1], 20) - check("fresh client recovers after restart", fin3 and isinstance(res3, dict), - f"{type(res3).__name__}: {res3}") - check("session id stable across restart" if newsid == sid else "session id changed (expected-ok)", - True, f"{sid} -> {newsid}") - c.close(); c2.close() - - -def s_full_restart_same_client(): - """Keep ONE auto-injected client across a full server restart. The FIRST - naive call after restart must transparently self-heal the stale db pin (no - internal poking, no manual re-resolve).""" - sid = ensure_up() - c = IDAClient(URL, db=None) # auto-resolve -> auto_recover eligible - c.connect() - old = c.resolve_db() - check("pre-restart call ok", isinstance(c.health(), dict)) - - ctl("stop") - # While down: calls must fail cleanly, fast. - fin, res = deadline(lambda: c.health(), 8) - check("while-down call returns (no hang)", fin) - check("while-down raises IDAError", isinstance(res, IDAError), f"{type(res).__name__}: {res}") - - ctl("start") - ensure_up() - # NAIVE call: no internal poking. Auto-recovery must kick in on the stale - # "Session not found" and retry against the fresh session. - fin2, res2 = deadline(lambda: c.health(), 25) - check("naive call auto-recovers after restart", fin2 and isinstance(res2, dict), - f"{type(res2).__name__}: {res2}") - check("db pin was re-resolved to new session", c.db is not None and c.db != old, - f"old={old} new={c.db}") - c.close() - - -def s_no_autoswitch_when_explicit(): - """With auto_recover disabled (or an explicitly pinned db), a stale session - must NOT be silently switched — it raises so the caller stays in control.""" - sid = ensure_up() - c = IDAClient(URL, db=sid, auto_recover_session=False) - c.connect() - check("pre-restart call ok", isinstance(c.health(), dict)) - ctl("stop"); ctl("start"); newsid = ensure_up() - fin, res = deadline(lambda: c.health(), 15) - check("explicit-pin call returns (no hang)", fin) - check("explicit-pin raises IDAToolError (no silent switch)", - isinstance(res, IDAToolError), f"{type(res).__name__}: {res}") - check("db pin unchanged when recovery disabled", c.db == sid, f"{c.db} vs {sid}") - c.close() - - -def s_concurrency_high(): - """Heavy concurrency: 300 calls over 16 threads, all succeed, no deadlock.""" - sid = ensure_up() - c = IDAClient(URL, db=sid, pool_size=8) - c.connect() - N = 300 - errors = [] - def work(i): - try: - tool = ("list_funcs", "disasm")[i % 2] - if tool == "list_funcs": - return bool(c.call("list_funcs", queries=[{"count": 3}])) - return bool(c.call("disasm", addr="main", max_instructions=5)) - except IDAError as e: - errors.append(e); return False - def run_all(): - with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex: - return sum(ex.map(work, range(N))) - fin, ok = deadline(run_all, 60) - check("300/16 concurrent finished (no deadlock)", fin, "pool deadlocked") - check("300/16 all succeeded", ok == N, f"{ok}/{N}, errs={errors[:3]}") - c.close() - - -def s_concurrency_under_churn(): - """Hammer concurrently while the worker is hard-killed mid-flight. No hang, - no deadlock; errors are clean; after respawn calls succeed again.""" - sid = ensure_up() - c = IDAClient(URL, db=sid, pool_size=8) - c.connect() - stop = threading.Event() - stats = {"ok": 0, "err": 0, "weird": 0} - lock = threading.Lock() - def spinner(): - while not stop.is_set(): - try: - c.call("list_funcs", queries=[{"count": 1}]) - with lock: stats["ok"] += 1 - except IDAError: - with lock: stats["err"] += 1 - except Exception: # noqa: BLE001 -- any non-IDAError is a bug - with lock: stats["weird"] += 1 - time.sleep(0.01) - threads = [threading.Thread(target=spinner, daemon=True) for _ in range(12)] - for t in threads: t.start() - time.sleep(1.0) - ctl("kill9") # crash the worker under load - time.sleep(2.0) - stop.set() - for t in threads: t.join(10) - alive = [t for t in threads if t.is_alive()] - check("no spinner thread hung on worker kill", not alive, f"{len(alive)} stuck") - check("no non-IDAError leaked during churn", stats["weird"] == 0, str(stats)) - log(f" churn stats: {stats}") - # Recover. - ctl("stop"); ctl("start"); ensure_up() - c.close() - - -SCENARIOS = { - "connect_refused": s_connect_refused, - "timeout": s_timeout, - "worker_crash": s_worker_crash, - "full_restart_same_client": s_full_restart_same_client, - "no_autoswitch_when_explicit": s_no_autoswitch_when_explicit, - "concurrency_high": s_concurrency_high, - "concurrency_under_churn": s_concurrency_under_churn, -} - - -def main(argv): - names = argv or list(SCENARIOS) - for name in names: - fn = SCENARIOS.get(name) - if not fn: - log(f"unknown scenario: {name} (have: {', '.join(SCENARIOS)})") - return 2 - log(f"\n[{name}]") - t0 = time.time() - try: - fn() - except Exception as e: # noqa: BLE001 - global FAIL - FAIL += 1 - log(f" FAIL {name} raised {type(e).__name__}: {e}") - log(f" ({time.time() - t0:.1f}s)") - log("\n== ensuring server is UP for subsequent work ==") - sid = ensure_up() - log(f" server ready, db={sid}") - log(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/tests/stress_paging.py b/tests/stress_paging.py deleted file mode 100644 index b821354..0000000 --- a/tests/stress_paging.py +++ /dev/null @@ -1,150 +0,0 @@ -#!/usr/bin/env python3 -"""Stress the 'many lines of text' problem: paging huge listings/functions. - -Validates the windowing strategy the TUI will use so we never hand a widget more -than a viewport-sized slice. Prints concise stats only. - - python3 tests/stress_paging.py --db <session_id> -""" -import os -import statistics -import sys -import time - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.client import IDAClient # noqa: E402 - -URL = "http://127.0.0.1:8745/mcp" -# Verified server cap: list_* honors count up to ~700, then silently collapses -# to a 10-item default. Clamp with margin. -SAFE_PAGE = 500 - - -def ms(fn): - t = time.time() - r = fn() - return (time.time() - t) * 1e3, r - - -def q0(payload): - """Unwrap a *_query/list_* payload: {'result':[{'data':[...],'next_offset':N}]}.""" - res = payload.get("result", payload) - if isinstance(res, list): - res = res[0] if res else {} - return res.get("data", []), res.get("next_offset") - - -def page_all_funcs(c, page=SAFE_PAGE): - """Enumerate the entire function list correctly: advance by len(data), NOT by - next_offset (which is just offset+count and skips over the per-call cap).""" - funcs, offset, pages, t0 = [], 0, 0, time.time() - while True: - payload = c.call("list_funcs", queries=[{"offset": offset, "count": page}]) - data, _ = q0(payload) - funcs.extend(data) - pages += 1 - if len(data) < page: # short page => reached the end - break - offset += len(data) - return funcs, pages, (time.time() - t0) * 1e3 - - -def disasm_meta(c, addr): - # total_instructions is a TOP-LEVEL field, not under 'asm'. - payload = c.call("disasm", addr=addr, max_instructions=1, include_total=True) - return payload.get("total_instructions", payload.get("instruction_count")) - - -def main(argv): - db = None - it = iter(argv) - for a in it: - if a == "--db": - db = next(it) - c = IDAClient(URL, db=db) - c.connect() - if db is None: - c.set_db(c.resolve_db()) - print(f"db={c.db} ({c.health().get('module')})") - - # ---- 1. Full function list via cursor pagination ------------------- # - print("\n[1] page entire function list (cursor)") - funcs, pages, total_ms = page_all_funcs(c, page=SAFE_PAGE) - n = len(funcs) - print(f" {n} functions in {pages} pages, {total_ms:.0f}ms " - f"({total_ms / max(n,1):.3f}ms/func)") - - def sz(f): - s = f["size"] - return int(s, 16) if isinstance(s, str) else s - biggest = sorted(funcs, key=sz, reverse=True)[:5] - print(" biggest by bytes:") - for f in biggest: - print(f" {f['addr']:>12} {sz(f):#8x} {f['name']}") - - # ---- 2. Instruction counts of the biggest -------------------------- # - print("\n[2] instruction totals of biggest funcs") - fattest = None - fattest_n = 0 - for f in biggest: - dt, total = ms(lambda f=f: disasm_meta(c, f["addr"])) - print(f" {f['name']:<28} {str(total):>8} insns (meta {dt:.1f}ms)") - if isinstance(total, int) and total > fattest_n: - fattest, fattest_n = f, total - - # ---- 3. Windowed paging INTO the fattest function ------------------ # - print(f"\n[3] window-scroll fattest func {fattest['name']} ({fattest_n} insns)") - WIN = 60 # a viewport - offsets = list(range(0, max(fattest_n - WIN, 1), max((fattest_n // 12), 1))) - times = [] - for off in offsets: - dt, payload = ms(lambda off=off: c.call( - "disasm", addr=fattest["addr"], offset=off, max_instructions=WIN)) - lines = payload.get("asm", {}).get("lines", []) - times.append(dt) - assert len(lines) <= WIN, f"got {len(lines)} > window {WIN}" - print(f" {len(offsets)} windowed reads @win={WIN}: " - f"min={min(times):.1f} med={statistics.median(times):.1f} " - f"max={max(times):.1f} ms (payload capped to <= {WIN} lines)") - - # ---- 3b. O(offset) latency curve (the key hazard) ------------------ # - print("\n[3b] disasm offset-latency curve (offset paging is O(offset))") - for off in [0, 1000, 5000, 20000, min(50000, fattest_n - WIN)]: - if off < 0: - continue - dt, _ = ms(lambda off=off: c.call( - "disasm", addr=fattest["addr"], offset=off, max_instructions=WIN)) - print(f" offset={off:>6}: {dt:6.1f}ms") - - # ---- 4. Simulated 'hold page-down' throughput ---------------------- # - print("\n[4] rapid sequential scroll (hold page-down)") - reads, t0, off = 0, time.time(), 0 - while time.time() - t0 < 2.0: - c.call("disasm", addr=fattest["addr"], offset=off, max_instructions=WIN) - off = (off + WIN) % max(fattest_n - WIN, 1) - reads += 1 - dur = time.time() - t0 - print(f" {reads} windowed reads in {dur:.1f}s = {reads / dur:.0f} reads/s " - f"(~{reads / dur * WIN:.0f} lines/s)") - - # ---- 5. Decompile: truncation AND hard-failure on monsters --------- # - print("\n[5] decompile on fattest func (may fail on huge funcs)") - dt, payload = ms(lambda: c.call("decompile", addr=fattest["addr"])) - code = payload.get("code") if isinstance(payload, dict) else None - if not code: - err = payload.get("error") if isinstance(payload, dict) else payload - print(f" decompile {dt:.0f}ms -> FAILED (soft error, not raised): {err}") - else: - marker = "chars total]" in code[-40:] - print(f" decompile {dt:.0f}ms, returned {len(code)} chars, " - f"server-truncated={marker}") - if marker: - print(f" tail: ...{code[-60:]!r}") - - c.close() - print("\nOK") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/tests/test_domain.py b/tests/test_domain.py deleted file mode 100644 index 24cb781..0000000 --- a/tests/test_domain.py +++ /dev/null @@ -1,191 +0,0 @@ -#!/usr/bin/env python3 -"""Validate the domain/paging layer against a real, large binary. - -Run with a session open on a big binary (e.g. libcrypto.so.3, 10k funcs): - - python3 tests/test_domain.py --db <session_id> - -Checks: correct full pagination (advance by len, not next_offset), viewport -slicing across block boundaries, window caching (revisit is instant), prefetch -warming, cached instruction totals, decompile success + hard-failure handling, -and address resolution. -""" -import os -import sys -import time - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.client import IDAClient # noqa: E402 -from idatui.domain import DISASM_BLOCK, LIST_PAGE, Program # noqa: E402 - -URL = "http://127.0.0.1:8745/mcp" -PASS = FAIL = 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 ms(fn): - t = time.time() - r = fn() - return (time.time() - t) * 1e3, r - - -def main(argv): - db = None - it = iter(argv) - for a in it: - if a == "--db": - db = next(it) - c = IDAClient(URL, db=db) - c.connect() - if db is None: - c.set_db(c.resolve_db()) - prog = Program(c) - module = c.health().get("module") - total_funcs = c.call("survey_binary").get("statistics", {}).get("total_functions") - print(f"db={c.db} module={module} survey_total_functions={total_funcs}") - - # ---- function index: full enumeration correctness ------------------ # - print("\n[function index]") - idx = prog.functions() - dt, _ = ms(lambda: idx.load_all()) - n = len(idx) - check("enumerated all functions (matches survey)", n == total_funcs, - f"got {n} vs survey {total_funcs}") - print(f" loaded {n} funcs in {dt:.0f}ms ({dt / max(n,1):.3f}ms/func), " - f"page={LIST_PAGE}") - # no duplicates, monotonic-ish uniqueness by addr - addrs = [idx.get(i).addr for i in range(min(n, 3000))] - check("no duplicate addrs in first 3000", len(addrs) == len(set(addrs))) - # viewport slice (adapt to binary size) - wstart = min(1000, max(n - 50, 0)) - wlen = min(50, n - wstart) - w = idx.window(wstart, wlen) - check(f"window({wstart},{wlen}) returns {wlen}", len(w) == wlen, str(len(w))) - - # ---- filtered index uses server-side glob -------------------------- # - print("\n[filtered index]") - sub = prog.functions(filter="sub_*") - sub.ensure(10) - check("filter sub_* yields sub_ names", all(f.name.startswith("sub_") for f in sub.window(0, 10)), - str([f.name for f in sub.window(0, 5)])) - - # ---- pick the fattest function for disasm stress ------------------- # - print("\n[disasm windowing]") - fattest = max((idx.get(i) for i in range(n)), key=lambda f: f.size) - dm = prog.disasm(fattest.addr, fattest.name) - dt, total = ms(dm.total) - check("total() returns positive count", total > 0, str(total)) - dt2, total2 = ms(dm.total) - check("total() cached (2nd call ~instant)", dt2 < dt / 2 + 1, f"{dt:.0f}ms -> {dt2:.1f}ms") - print(f" fattest {fattest.name}: {total} insns, total() {dt:.0f}ms then {dt2:.1f}ms") - - # viewport across a block boundary - start = DISASM_BLOCK - 5 - win = dm.lines(start, 60, prefetch=False) - check("viewport spans block boundary, right length", - len(win) == min(60, max(total - start, 0)), f"got {len(win)}") - # addresses strictly increasing and contiguous slice - eas = [ln.ea for ln in win] - check("viewport addrs strictly increasing", all(b > a for a, b in zip(eas, eas[1:])), - str(eas[:4])) - - # deep window: first slow (O(offset)), revisit instant (cached) - if total > DISASM_BLOCK * 4: - deep = (total // DISASM_BLOCK - 1) * DISASM_BLOCK - dt_cold, a = ms(lambda: dm.lines(deep, 60, prefetch=False)) - dt_warm, b = ms(lambda: dm.lines(deep, 60, prefetch=False)) - check("deep window revisit is cached/instant", dt_warm < dt_cold / 2 + 1, - f"cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms") - check("cached window identical", [l.ea for l in a] == [l.ea for l in b]) - print(f" deep@{deep}: cold={dt_cold:.0f}ms warm={dt_warm:.1f}ms") - - # prefetch warms the next block - print("\n[prefetch]") - dm2 = prog.disasm(fattest.addr + 0) # same model (cached by ea) - fresh = prog.disasm(idx.get(0).addr, idx.get(0).name) - fresh.lines(0, 60, prefetch=True) # should prefetch block 1 - time.sleep(0.3) - check("prefetch warmed a neighbor block", fresh.cached_blocks() >= 2, - f"cached_blocks={fresh.cached_blocks()}") - - # ---- decompile: success + hard-failure ----------------------------- # - print("\n[decompile]") - # a small function likely decompiles - small = min((idx.get(i) for i in range(n)), key=lambda f: f.size if f.size > 4 else 1 << 30) - d_small = prog.decompile(small.addr) - check("small func decompiles or fails cleanly", isinstance(d_small.failed, bool)) - dt_c, _ = ms(lambda: prog.decompile(small.addr)) - check("decompile cached (2nd ~instant)", dt_c < 5, f"{dt_c:.1f}ms") - # the monster should hard-fail as a soft error (not raise) - d_big = prog.decompile(fattest.addr) - check("monster decompile handled (failed flag, no raise)", - d_big.failed or d_big.code is not None, - f"failed={d_big.failed} err={d_big.error}") - print(f" small {small.name}: failed={d_small.failed} " - f"trunc={d_small.truncated} chars={d_small.total_chars}") - print(f" monster {fattest.name}: failed={d_big.failed} err={d_big.error}") - - # ---- resolve ------------------------------------------------------- # - print("\n[resolve]") - check("resolve hex", prog.resolve(hex(fattest.addr)) == fattest.addr) - check("resolve int passthrough", prog.resolve(fattest.addr) == fattest.addr) - named = next((idx.get(i) for i in range(n) if not idx.get(i).name.startswith("sub_")), None) - if named: - try: - r = prog.resolve(named.name) - check("resolve symbol name", r == named.addr, f"{hex(r)} vs {hex(named.addr)} ({named.name})") - except KeyError as e: - check("resolve symbol name", False, str(e)) - - # ---- flat listing (code + data + undefined heads) ------------------ # - print("\n[listing]") - seg = prog.segment_bounds(fattest.addr) - check("segment_bounds finds the .text segment", seg is not None and seg[0] <= fattest.addr < seg[1], - str(seg)) - lm = prog.listing(fattest.addr) - check("listing() returns a model for a mapped address", lm is not None) - if lm is not None: - lm.ensure(20) - w = lm.window(0, 20) - check("listing window returns heads", len(w) == 20, str(len(w))) - eas = [h.ea for h in w] - check("listing head addrs strictly increasing", all(b > a for a, b in zip(eas, eas[1:])), - str(eas[:4])) - check("listing heads carry a kind", all(h.kind in ("code", "data", "unknown") for h in w), - str({h.kind for h in w})) - check("listing head sizes positive", all(h.size >= 1 for h in w), - str([h.size for h in w[:6]])) - # random access to a mid-segment address lands on the containing head - mid = w[10].ea - li = lm.ensure_ea(mid) - check("ensure_ea lands on the exact head for a head address", - li >= 0 and lm.get(li).ea == mid, f"idx={li}") - # a mid-item byte resolves to its containing head - if w[4].size > 1: - inside = w[4].ea + 1 - li2 = lm.ensure_ea(inside) - check("ensure_ea resolves a mid-item byte to its head", - li2 >= 0 and lm.get(li2).ea == w[4].ea, f"idx={li2} ea={w[4].ea:#x}") - dt_cold, _ = ms(lambda: lm.window(0, 20)) - check("listing window revisit is cached/instant", dt_cold < 5, f"{dt_cold:.1f}ms") - dt_all, _ = ms(lambda: lm.load_all()) - check("listing load_all completes the segment", lm.complete and len(lm) > 20, - f"n={len(lm)} complete={lm.complete}") - print(f" {seg[2]}: {len(lm)} heads walked in {dt_all:.0f}ms") - - prog.close() - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/tests/test_keepalive.py b/tests/test_keepalive.py deleted file mode 100644 index 2910e1a..0000000 --- a/tests/test_keepalive.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -"""Prove that an interactive session can 'just chill' without the worker idling -out. Uses a deliberately short worker TTL to keep the test fast. - - python3 tests/test_keepalive.py - -Needs ~/ida-venv + a running server (tests/serverctl.sh) and a writable target. -""" -import os -import sys -import time - -HERE = os.path.dirname(os.path.abspath(__file__)) -REPO = os.path.dirname(HERE) -sys.path.insert(0, REPO) -from idatui.client import IDAClient, IDAToolError # noqa: E402 - -PASS = FAIL = 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 alive(c, sid): - try: - c.set_db(sid) - c.health() - return True - except IDAToolError: - return False - - -def main(): - target = os.path.join(REPO, "targets", "ls_ttl") - if not os.path.exists(target): - src = os.path.join(REPO, "bin", "ls") - os.makedirs(os.path.dirname(target), exist_ok=True) - import shutil - shutil.copy(src, target) - - c = IDAClient(timeout=300) - c.connect() - SHORT = 12 # worker self-exits after ~12s idle unless kept alive - - print("[heartbeat keeps a short-TTL worker alive]") - sid = c.call("idb_open", input_path=target, idle_ttl_sec=SHORT)["session"]["session_id"] - c.set_db(sid) - ka = c.keepalive(interval=4.0).start() - time.sleep(SHORT * 2 + 2) # idle well past the TTL, but heartbeat is beating - ok = alive(c, sid) - check("alive past 2x TTL with heartbeat", ok, f"beats={ka.beats}") - check("heartbeat actually beat", ka.beats >= 4, f"beats={ka.beats}") - check("heartbeat had no failures", ka.failures == 0, f"failures={ka.failures}") - ka.stop() - - print("\n[bump_idle_ttl makes it effectively immortal]") - sid = c.call("idb_open", input_path=target, idle_ttl_sec=SHORT)["session"]["session_id"] - c.set_db(sid) - c.bump_idle_ttl() # ~1e9 seconds - time.sleep(SHORT * 2 + 6) # zero requests during this window - check("alive past 2x TTL after bump, zero requests", alive(c, sid)) - - c.close() - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main()) |
