From f4c1d9b5497fd0d38137b6b345e8171b307c1117 Mon Sep 17 00:00:00 2001 From: Duncan Ogilvie Date: Thu, 20 Aug 2026 23:55:52 +0200 Subject: Switch to ida-nexus --- idatui/codemode_client.py | 551 ---------------------------------------------- 1 file changed, 551 deletions(-) delete mode 100644 idatui/codemode_client.py (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py deleted file mode 100644 index e95238f..0000000 --- a/idatui/codemode_client.py +++ /dev/null @@ -1,551 +0,0 @@ -"""Client adapter from ida-tui's domain operations to IDA Code Mode. - -``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered -GUI database, reuses a shared managed idalib worker, or starts one when needed. -The TUI never owns or terminates an IDA process. Closing this client releases -only its lease. - -Remote operations are ordinary typed Python functions declared in -``idatui.remote_ops``. Code Mode installs their content-addressed modules once -per handle; subsequent calls send only encoded arguments. The optimized -IDAPython listing/decompiler implementation remains real source in -``idatui.remote_tools`` and is installed through the same module interface. -""" - -from __future__ import annotations - -import os -import shlex -import threading -import time -from collections.abc import Callable -from typing import Any - -from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session - -# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost. -# -# The paging/graph/trace layers and their offline test suites must keep importing -# `idatui` on a machine with no IDA and no Code Mode installed -- that is the -# house rule the stdlib-only worker client used to satisfy for free, and -# `tests/run.py --fast` (380 checks, any python3) depends on it. A hard top-level -# import here makes the whole package unimportable, so the failure is deferred to -# the first operation that genuinely needs the library. -_CODEMODE_ERROR: Exception | None = None -try: - from ida_codemode import ( - CodeModeConnectionError, - DatabaseBusyError, - DatabaseDisconnectedError, - DatabaseHandle, - DatabaseInstance, - DatabaseOpenOptions, - RemoteError, - find_database_owner, - wait_database_released, - ) -except ImportError as _exc: # library absent: usable only for offline layers - _CODEMODE_ERROR = _exc - # Bound to None rather than left undefined so the names stay patchable: the - # offline contract tests inject a fake DatabaseHandle here. - CodeModeConnectionError = DatabaseDisconnectedError = RemoteError = None # type: ignore[assignment,misc] - DatabaseBusyError = DatabaseHandle = DatabaseInstance = None # type: ignore[assignment,misc] - DatabaseOpenOptions = find_database_owner = wait_database_released = None # type: ignore[assignment] - - -def _require_codemode() -> None: - """Raise an actionable error when the Code Mode library is missing. - - Gated on the binding, not on the original import result, so a test that - injects a fake ``DatabaseHandle`` exercises the real adapter logic. - """ - if DatabaseHandle is None: - raise IDAConnectionError( - "ida-codemode is not installed in this environment " - f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " - "`pip install ida-codemode`) so ida-tui can lease a " - "database." - ) from _CODEMODE_ERROR - - -def database_owner(idb_path: str, staged_path: str | None = None): - """The Code Mode instance that owns ``idb_path``/``staged_path``, else None. - - Returns None when the Code Mode library is absent: with no library there is - no client in this environment that could be holding the database, and the - IDA-free layers (project staging) must keep working. Discovery errors with - the library installed still propagate because unknown ownership is unsafe. - """ - if DatabaseHandle is None: - return None - if staged_path: - owner = find_database_owner( - staged_path, - output_database=idb_path, - timeout=0.5, - ) - return owner or find_database_owner(staged_path, timeout=0.5) - return find_database_owner(idb_path, timeout=0.5) - - -def registered_database(path: str, output_database: str | None = None) -> bool: - """Whether a live/lock-held Code Mode instance owns this target.""" - _require_codemode() - return ( - find_database_owner( - path, - output_database=output_database, - timeout=0.5, - ) - is not None - ) - - -class _NoopKeepAlive: - """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat.""" - - def __init__(self) -> None: - self.beats = self.failures = 0 - - def start(self) -> "_NoopKeepAlive": - return self - - def stop(self) -> None: - pass - - -def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: - """Translate ida-tui's legacy first-open switches to Code Mode options. - - Code Mode has typed options for processor, natural loading address and file - type. It deliberately has no arbitrary command-line escape hatch; reject - switches we cannot represent instead of silently loading a blob wrongly. - """ - processor: str | None = None - loading_address: int | None = None - file_type: str | None = None - unsupported: list[str] = [] - try: - words = shlex.split(value or "", posix=os.name != "nt") - except ValueError as exc: - raise ValueError(f"invalid IDA load options: {exc}") from exc - for word in words: - if word.startswith("-p") and len(word) > 2: - processor = word[2:] - elif word.startswith("-b") and len(word) > 2: - try: - # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the - # natural address, which is the safer public API. - loading_address = int(word[2:], 16) << 4 - except ValueError as exc: - raise ValueError(f"invalid IDA loading address: {word!r}") from exc - elif word.startswith("-T") and len(word) > 2: - file_type = word[2:] - else: - unsupported.append(word) - if unsupported: - joined = " ".join(unsupported) - raise ValueError( - "ida-codemode cannot represent arbitrary IDA load options: " - f"{joined!r}; use processor/base/file type options instead" - ) - return processor, loading_address, file_type - - -class IDBEventListener: - """Debounced, closeable delivery of another client's IDB changes. - - Code Mode's subscription is a blocking iterator, so one daemon thread reads - it and a second waits for a quiet period before handing a batch to the UI. - Keeping the debounce here avoids a permanent Textual worker (which would - make the app's worker-idle contract impossible) and bounds refresh work to - one pass per edit burst. - """ - - def __init__( - self, - client: "CodeModeClient", - callback: Callable[[tuple[dict[str, Any], ...]], None], - *, - on_error: Callable[[BaseException], None] | None = None, - debounce: float = 0.2, - ) -> None: - self._client = client - self._callback = callback - self._on_error = on_error - self._debounce = max(float(debounce), 0.0) - self._condition = threading.Condition() - self._closed = False - self._subscription = None - self._pending: list[dict[str, Any]] = [] - self._deadline = 0.0 - self._reader = threading.Thread( - target=self._read, name="idatui-idb-events", daemon=True - ) - self._deliverer = threading.Thread( - target=self._deliver, name="idatui-idb-refresh", daemon=True - ) - self._deliverer.start() - self._reader.start() - - def _report(self, error: BaseException) -> None: - disconnected = DatabaseDisconnectedError - if isinstance(disconnected, type) and isinstance(error, disconnected): - error = self._client._connection_error(error) - with self._condition: - closed = self._closed - if not closed and self._on_error is not None: - self._on_error(error) - - def _read(self) -> None: - try: - subscription = self._client.subscribe_idb_events() - except Exception as exc: # noqa: BLE001 -- surfaced through on_error - self._report(exc) - with self._condition: - self._closed = True - self._pending.clear() - self._condition.notify_all() - return - with self._condition: - if self._closed: - subscription.close() - return - self._subscription = subscription - try: - for event in subscription: - with self._condition: - if self._closed: - break - if self._client.owns_event(event): - continue - with self._condition: - if self._closed: - break - self._pending.append(event) - self._deadline = time.monotonic() + self._debounce - self._condition.notify_all() - except Exception as exc: # noqa: BLE001 -- stream failures are recoverable - self._report(exc) - finally: - subscription.close() - with self._condition: - if self._subscription is subscription: - self._subscription = None - self._closed = True - self._pending.clear() - self._condition.notify_all() - - def _deliver(self) -> None: - while True: - with self._condition: - while not self._closed and not self._pending: - self._condition.wait() - if self._closed: - return - remaining = self._deadline - time.monotonic() - if remaining > 0: - self._condition.wait(remaining) - continue - batch = tuple(self._pending) - self._pending.clear() - try: - self._callback(batch) - except Exception as exc: # noqa: BLE001 -- keep the stream alive - self._report(exc) - - def close(self) -> None: - """Stop delivery and unblock the subscription reader.""" - with self._condition: - if self._closed: - return - self._closed = True - self._pending.clear() - subscription = self._subscription - self._condition.notify_all() - if subscription is not None: - subscription.close() - - -class CodeModeClient: - """A leased GUI/idalib database accessed through ``ida_codemode``.""" - - def __init__( - self, - binary_path: str, - *, - ttl: int = 0, - load_args: str = "", - processor: str | None = None, - loading_address: int | None = None, - file_type: str | None = None, - output_database: str | None = None, - spawn: bool = True, - new_database: bool = False, - ) -> None: - del ttl # managed-worker lifetime is lease-based, not idle-TTL based - self._path = os.path.abspath(os.path.expanduser(binary_path)) - parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args) - self._processor = processor or parsed_processor - self._loading_address = ( - loading_address if loading_address is not None else parsed_address - ) - self._file_type = file_type or parsed_file_type - self._output_database = output_database - self._spawn = spawn - self._new_database = new_database - self._handle: DatabaseHandle | None = None - self._last_instance: DatabaseInstance | None = None - self._connect_lock = threading.Lock() - - def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": - _require_codemode() - with self._connect_lock: - handle = self._handle - if handle is not None: - if handle.connected: - return self - raise IDAConnectionError( - "Code Mode database disconnected; explicit rediscovery required" - ) - if progress: - progress( - f"discovering Code Mode database for {os.path.basename(self._path)}…" - ) - try: - # A Ctrl+L reload releases its current managed-worker lease, but - # that worker remains registered during Code Mode's final-lease - # grace period. Retry only that known handoff window. A GUI or - # another long-lived client remains busy and yields a clear - # failure rather than being modified underneath its owner. - deadline = time.monotonic() + min(timeout, 60.0) - while True: - try: - handle = DatabaseHandle.open( - self._path, - options=DatabaseOpenOptions( - spawn=self._spawn, - startup_timeout=max(0.1, timeout), - output_database=self._output_database, - processor=self._processor, - # The natural byte address is converted to IDA's - # paragraph-based -b value by Code Mode. - image_base=self._loading_address, - file_type=self._file_type, - new_database=self._new_database, - ), - ) - break - except DatabaseBusyError: - if not self._new_database or time.monotonic() >= deadline: - raise - if progress: - progress( - "waiting for the previous Code Mode lease to close…" - ) - owner = find_database_owner( - self._path, - output_database=self._output_database, - timeout=0.5, - ) - if owner is not None: - wait_database_released( - owner, - max(0.0, deadline - time.monotonic()), - ) - else: - time.sleep(0.2) - if progress: - backend = handle.instance.backend - progress( - f"attached to {backend} database; waiting for auto-analysis…" - ) - handle.wait_autoanalysis(timeout=timeout) - except Exception as exc: # normalize the dependency's transport errors - raise self._connection_error(exc) from exc - self._handle = handle - self._last_instance = handle.instance - return self - - @staticmethod - def _connection_error(exc: BaseException) -> IDAConnectionError: - return IDAConnectionError(str(exc) or type(exc).__name__) - - @property - def connected(self) -> bool: - return self._handle is not None and self._handle.connected - - @property - def pid(self) -> int | None: - return self._handle.instance.pid if self._handle is not None else None - - @property - def backend(self) -> str | None: - return self._handle.instance.backend if self._handle is not None else None - - def owns_event(self, event: dict[str, Any]) -> bool: - """Whether ``event`` was produced through this client's handle.""" - handle = self._handle - return handle is not None and handle.owns_event(event) - - def subscribe_idb_events(self): - """Open Code Mode's closeable IDB-change iterator.""" - if not self.connected: - self.connect() - handle = self._handle - if handle is None: - raise IDAConnectionError("Code Mode database is not connected") - try: - return handle.subscribe_idb_events() - except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: - raise self._connection_error(exc) from exc - - def watch_idb_events( - self, - callback: Callable[[tuple[dict[str, Any], ...]], None], - *, - on_error: Callable[[BaseException], None] | None = None, - debounce: float = 0.2, - ) -> IDBEventListener: - """Deliver external IDB changes in debounced batches.""" - return IDBEventListener(self, callback, on_error=on_error, debounce=debounce) - - def call(self, operation: Callable[..., Any], /, **args) -> Any: - """Execute one source-backed remote declaration through this client.""" - name = getattr(operation, "__name__", "remote operation") - try: - from .remote_ops import bind - - remote = bind(operation) - except KeyError as exc: - raise IDAToolError( - name, f"remote operation {name!r} is not registered" - ) from exc - if not self.connected: - self.connect() - handle = self._handle - if handle is None: - raise IDAConnectionError("Code Mode database is not connected") - try: - return remote(handle, **args) - except RemoteError as exc: - message = str(exc) - if exc.details.get("traceback"): - message += f"\n{exc.details['traceback']}" - if exc.code == "operation_timeout": - raise IDATimeoutError(message) from exc - raise IDAToolError(name, message) from exc - except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: - raise self._connection_error(exc) from exc - - def save_database(self) -> dict[str, Any]: - if not self.connected: - self.connect() - handle = self._handle - if handle is None: - raise IDAConnectionError("Code Mode database is not connected") - try: - return handle.save_database() - except RemoteError as exc: - raise IDAToolError("save_database", str(exc)) from exc - except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: - raise self._connection_error(exc) from exc - - def discard_database(self, timeout: float = 5.0) -> bool: - """Discard a final managed-worker lease; otherwise transfer finalization. - - ``False`` is an expected ownership result: a GUI owns its session, or - another lease still shares the managed worker. A busy final worker is - retried briefly so background reads finishing during quit do not turn a - real discard into an implicit save. - """ - handle = self._handle - if handle is None or not handle.connected: - return False - entry = handle.instance - if entry.backend != "idalib" or not getattr(entry, "managed", False): - return False - deadline = time.monotonic() + max(float(timeout), 0.0) - while True: - try: - handle.shutdown_database(save=False) - return True - except RemoteError as exc: - if exc.code in ("instance_shared", "shutdown_not_supported"): - return False - if exc.code == "instance_busy" and time.monotonic() < deadline: - time.sleep(0.05) - continue - raise IDAToolError("shutdown_database", str(exc)) from exc - except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: - raise self._connection_error(exc) from exc - - def health(self) -> dict[str, Any]: - if not self.connected: - self.connect() - assert self._handle is not None - entry = self._handle.instance - module = os.path.basename(entry.exe_path or entry.idb_path or self._path) - return { - "ok": self._handle.connected, - "module": module, - "backend": entry.backend, - "record_id": entry.record_id, - "input_path": entry.exe_path, - "idb_path": entry.idb_path, - } - - def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: - del interval - return _NoopKeepAlive() - - def resolve_db(self) -> str: - if not self.connected: - self.connect() - assert self._handle is not None - return self._handle.instance.record_id - - def set_db(self, db: str | None) -> None: - del db # one handle is permanently bound to one registered database - - def list_sessions(self) -> list[Session]: - if not self.connected: - self.connect() - assert self._handle is not None - entry = self._handle.instance - path = entry.exe_path or entry.idb_path or self._path - return [ - Session( - session_id=entry.record_id, - filename=os.path.basename(path), - input_path=path, - is_active=True, - ) - ] - - def close(self, grace: float = 0.0) -> None: - del grace - with self._connect_lock: - handle, self._handle = self._handle, None - if handle is not None: - self._last_instance = handle.instance - handle.close() # release our lease; never close a GUI/other client's DB - - def wait_released(self, timeout: float = 45.0) -> bool: - """Wait until a managed instance releases its lifetime lock. - - Normal application shutdown must not wait: another client may retain the - worker. This is an explicit test/maintenance helper for deleting a - temporary IDB safely after this client closes. GUI instances return - ``False`` immediately because clients never own their lifetime. - """ - instance = self._last_instance - if instance is None or instance.backend != "idalib": - return False - return wait_database_released(instance, timeout) - - def __enter__(self) -> "CodeModeClient": - return self.connect() - - def __exit__(self, *exc) -> None: - self.close() -- cgit v1.3.1-sl0p