From f3715d8de0d255c8b14710acfa120ccb9ea953fd Mon Sep 17 00:00:00 2001 From: Duncan Ogilvie Date: Thu, 20 Aug 2026 23:42:42 +0200 Subject: Adopt idb_events and remote module features from ida-codemode --- idatui/codemode_client.py | 1371 +++++++-------------------------------------- 1 file changed, 214 insertions(+), 1157 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 8eb66b5..e95238f 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -5,24 +5,20 @@ 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. -The Code Mode transport intentionally exposes one broad operation, -``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric -operations needed by the paging layer into self-contained snippets. The -snippets prefer the public ``ida-domain`` ``db`` object. A handful of features -that ida-domain does not currently expose (IDA-coloured listing rows, creating -instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the -IDAPython modules that Code Mode deliberately makes importable. +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 hashlib -import json import os import shlex import threading import time -from pathlib import Path -from textwrap import dedent +from collections.abc import Callable from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -68,7 +64,8 @@ def _require_codemode() -> None: "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 + "database." + ) from _CODEMODE_ERROR def database_owner(idb_path: str, staged_path: str | None = None): @@ -155,1106 +152,119 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: return processor, loading_address, file_type -#: Key of the pre-serialised payload envelope. See _script(). -_PACKED = "__idatui_json__" - -#: Serialise the answer INSIDE the database process and hand back one string. -#: -#: Written when Code Mode ran to_jsonable() over every snippet result, walking -#: the whole structure in Python to make it JSON-safe: a 200-row listing page is -#: ~10k small objects, which cost 66ms to walk -- 72% of the page's total cost, -#: and 114x what json.dumps of the same data cost (0.58ms). -#: -#: ida-codemode 0.3.2 removed that reason: serialization.dumps_json now hands -#: the structure straight to the C encoder and only falls back to the walker for -#: values json.dumps rejects. Re-measured against 0.3.2, packing buys 0.97x on -#: that same page (experiments/bench_pack_trace.py) -- i.e. nothing, because the -#: dodged walk is replaced by a double encode. -#: -#: It is kept anyway, on correctness rather than speed: packing pins OUR encoder -#: settings (compact separators, default=str) inside the database process, so an -#: un-encodable IDA object degrades to repr() at a point we control instead of -#: depending on the runtime's fallback. Delete it if that stops being worth a -#: protocol step -- it is no longer load-bearing for performance. -_PACK_EPILOGUE = ( - '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' -) - +class IDBEventListener: + """Debounced, closeable delivery of another client's IDB changes. -def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code. - - This used to also run the body with Code Mode's trace hook detached - (sys.settrace(None) + restore), because the runtime wrapped every - execute_python in a trace function that returned ITSELF -- enabling line - tracing in every frame it saw, so every line of every function we called - paid a Python-level callback (ida_bytes.get_flags: 0.106us -> 5.49us, 52x). - - ida-codemode 0.3.2 deleted that hook; cancellation is now a C-level thread - interrupt (runtime._interrupt_thread) that costs nothing while idle. The - workaround measured 0.99x on a 200-row listing page against 0.3.2 -- pure - noise -- so it is gone, and with it the caveat that a pure-Python loop in a - snippet escaped its deadline. See experiments/bench_pack_trace.py. + 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. """ - encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - head = f"import json\na = json.loads({encoded!r})\n" - return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" - -_OPERATIONS: dict[str, str] = { - "list_funcs": r''' -import fnmatch -queries = a.get("queries") or [{}] -q = queries[0] -offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500))) -pattern = str(q.get("filter") or "").lower() -if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*" -rows = [] -for fn in db.functions.get_all(): - name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" - if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue - rows.append({"addr": hex(int(fn.start_ea)), "name": name, - "size": int(fn.end_ea) - int(fn.start_ea)}) -page = rows[offset:offset + count] -result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]} -result -''', - "disasm": r''' -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"instructions": [], "total_instructions": 0, "instruction_count": 0} -else: - instructions = list(db.functions.get_instructions(fn)) - limit = max(1, int(a.get("max_instructions", len(instructions) or 1))) - rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)} - for insn in instructions[:limit]] - result = {"instructions": rows, "total_instructions": len(instructions), - "instruction_count": len(instructions)} -result -''', - "file_regions": r''' -import idaapi -rows = [] -for seg in db.segments.get_all(): - try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea)) - except Exception: file_off = -1 - if file_off < 0 or file_off >= (1 << 48): file_off = -1 - rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), - "file_off": file_off, "name": db.segments.get_name(seg) or ""}) -result = {"regions": rows} -result -''', - "read_raw": r''' -import ida_bytes -ea, size = int(str(a["addr"]), 16), max(0, int(a["size"])) -raw = ida_bytes.get_bytes(ea, size) or b"" -raw = raw[:size] + b"\xff" * max(0, size - len(raw)) -data = bytearray(raw) -for index, value in enumerate(data): - if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0 -result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)} -result -''', - "get_bytes": r''' -rows = [] -for region in a.get("regions", []): - ea, size = int(str(region["addr"]), 16), int(region["size"]) - raw = db.bytes.get_bytes_at(ea, size) or b"" - rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)}) -result = {"result": rows} -result -''', - "search_structs": r''' -needle = str(a.get("filter") or "").lower() -rows = [] -for tif in db.types.get_all(): - name = tif.get_type_name() or "" - if not name or needle not in name.lower() or not tif.is_udt(): continue - members = list(db.types.get_udt_members(tif)) - rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), - "cardinality": len(members), "ordinal": int(tif.get_ordinal())}) -result = {"result": rows} -result -''', - "type_inspect": r''' -rows = [] -for query in a.get("queries", []): - name = str(query.get("name") or "") - tif = db.types.get_by_name(name) - if tif is None: - rows.append({"name": name, "error": "type not found"}); continue - members = [{"name": m.name, "type": m.type.dstr() or str(m.type), - "offset": int(m.offset), "size": int(m.size)} - for m in db.types.get_udt_members(tif)] if tif.is_udt() else [] - rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), - "members": members}) -result = {"result": rows} -result -''', - "declare_type": r''' -import ida_typeinf -decls = a.get("decls", "") -if isinstance(decls, str): decls = [decls] -rows = [] -for declaration in decls: - try: - errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration)) - rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})}) - except Exception as exc: - rows.append({"ok": False, "error": str(exc)}) -result = {"result": rows} -result -''', - "del_type": r''' -import ida_typeinf -name = str(a["name"]) -ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE)) -result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})} -result -''', - "func_types": r''' -import ida_typeinf -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"addr": a["addr"], "error": "no function at address"} -else: - pseudo = db.pseudocode.decompile(fn) - name = db.functions.get_name(fn) or "" - tif = pseudo.get_func_type() - try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else "" - except Exception: prototype = tif.dstr() if tif else "" - lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "", - "is_arg": bool(var.is_arg)} for var in pseudo.local_variables] - result = {"addr": hex(int(fn.start_ea)), "name": name, - "prototype": (prototype or "").strip(), "lvars": lvars} -result -''', - "set_lvar_type": r''' -import ida_typeinf -ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"]) -fn = db.functions.get_at(ea) -if fn is None: - result = {"error": "no function at address"} -else: - pseudo = db.pseudocode.decompile(fn) - var = pseudo.find_local_variable(variable) - if var is None: - result = {"error": f"local variable {variable!r} not found"} - else: + 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: - tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) - accepted = bool(var.set_type(tif)) - saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False - result = {"addr": hex(int(fn.start_ea)), "variable": variable, - "type": declaration, "ok": accepted and saved} - except Exception as exc: - result = {"error": f"bad type {declaration!r}: {exc}"} -result -''', - "set_type": r''' -from ida_domain.types import TypeApplyFlags -rows = [] -for edit in a.get("edits", []): - ea = int(str(edit["addr"]), 16) - declaration = str(edit.get("signature") or edit.get("type") or "") - try: - ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE)) - rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})}) - except Exception as exc: - rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) -result = {"result": rows} -result -''', - "data_type": r''' -ea = int(str(a["addr"]), 16) -try: - tif = db.types.get_at(ea) - fn = db.functions.get_at(ea) - result = {"addr": hex(ea), "name": db.names.get_at(ea) or "", - "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0, - "is_func": bool(fn)} -except Exception as exc: - result = {"addr": hex(ea), "error": str(exc)} -result -''', - "force_recompile": r''' -import ida_hexrays -rows = [] -for item in a.get("items", []): - ea = int(str(item["addr"]), 16) - ida_hexrays.mark_cfunc_dirty(ea, False) - rows.append({"addr": hex(ea), "ok": True}) -result = {"result": rows} -result -''', - "undefine": r''' -import ida_bytes -rows = [] -for item in a.get("items", []): - ea = int(str(item["addr"]), 16) - size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1)) - ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)) - rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})}) -result = {"result": rows} -result -''', - "define_code": r''' -import ida_ua -rows = [] -for item in a.get("items", []): - ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea)) - rows.append({"addr": hex(ea), "ok": size > 0, "size": size, - **({} if size > 0 else {"error": "instruction did not decode"})}) -result = {"result": rows} -result -''', - "define_func": r''' -rows = [] -for item in a.get("items", []): - ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea)) - rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})}) -result = {"result": rows} -result -''', - "make_data": r''' -import ida_bytes, ida_idaapi, ida_typeinf -from ida_domain.types import TypeApplyFlags -rows = [] -for item in a.get("items", []): - ea, declaration = int(str(item["addr"]), 16), str(item["type"]) - try: - tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) - size = max(1, int(tif.get_size())) - saved_names = [(addr, name) for addr, name in db.names.get_all() - if ea <= int(addr) < ea + size] - ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES, - max(size, int(ida_bytes.get_item_size(ea) or 1))) - created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR)) - ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE)) - for address, name in saved_names: - db.names.set_name(int(address), name) - if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"]))) - rows.append({"addr": hex(ea), "ok": ok, "size": size, - **({} if ok else {"error": "IDA rejected the data type"})}) - except Exception as exc: - rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) -result = {"result": rows} -result -''', - "make_string": r''' -from ida_domain.strings import StringType -ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0))) -kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32, - "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C) -import ida_bytes -try: - ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1) -except Exception: - pass -try: - ok = bool(db.bytes.create_string_at(ea, length or None, kind)) - text = db.bytes.get_string_at(ea) or "" if ok else "" - result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text} -except Exception as exc: - result = {"addr": hex(ea), "ok": False, "error": str(exc)} -result -''', - "list_strings": r''' -from ida_domain.strings import StringListConfig -offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4))) -if offset == 0 or a.get("refresh"): - from ida_domain.strings import StringType - db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len, - only_ascii_7bit=False)) -items = list(db.strings.get_all()) -page = items[offset:offset + count] -rows = [] -for item in page: - try: text = str(item) - except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else "" - rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name}) -result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)} -result -''', - # Everything a person ADDED to the database: comments, non-dummy names, and - # the prototypes they set. - # - # Names come from IDA's name list, which is already an index -- no scan at - # all. Comments have no index, so they need a walk, and the walk is over - # HEADS: `next_that`'s predicate is a *Python* callback (SWIG calls it with - # one argument, so `f_has_cmt` does not even fit), which would be one call - # per BYTE -- 400 million of them on a big image. `max_scan` bounds it and - # reports `truncated` rather than sitting there. - "list_annotations": r''' -import ida_bytes, ida_funcs, ida_lines, ida_nalt, ida_name -import ida_segment, ida_typeinf, idautils -limit = max(1, int(a.get("limit", 4000))) -max_scan = max(1000, int(a.get("max_scan", 2000000))) -comments, names = [], [] -scanned = 0 - -def _line(ea): - try: - txt = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) - except Exception: - txt = "" - return " ".join((txt or "").split()) - -for ea, nm in idautils.Names(): - if len(names) >= limit: - break - if not nm or not ida_bytes.has_user_name(ida_bytes.get_flags(ea)): - continue - fn = ida_funcs.get_func(ea) - is_fn = fn is not None and int(fn.start_ea) == int(ea) - proto = None - if is_fn: + 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: - ti = ida_typeinf.tinfo_t() - if ida_nalt.get_tinfo(ti, ea): - proto = str(ti) - except Exception: - proto = None - seg = ida_segment.getseg(ea) - names.append({"addr": hex(int(ea)), "name": nm, "func": is_fn, - "size": (int(fn.end_ea - fn.start_ea) if is_fn else 0), - "proto": proto, - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) - -for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(comments) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(comments) >= limit or scanned >= max_scan: - break - if not ida_bytes.has_cmt(ida_bytes.get_flags(ea)): - continue - for rep in (False, True): - text = ida_bytes.get_cmt(ea, rep) - if text: - fn = ida_funcs.get_func(ea) - comments.append({ - "addr": hex(int(ea)), "text": text, "repeatable": rep, - "line": _line(ea), "seg": ida_segment.get_segm_name(seg), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None)}) - -# Whole-function comments are not on the byte flags, so the scan cannot see them. -for fn_ea in idautils.Functions(): - fn = ida_funcs.get_func(fn_ea) - if fn is None or len(comments) >= limit: - continue - for rep in (False, True): - text = ida_funcs.get_func_cmt(fn, rep) - if text: - seg = ida_segment.getseg(fn_ea) - comments.append({"addr": hex(int(fn_ea)), "text": text, - "repeatable": rep, "line": "", "whole_func": True, - "seg": (ida_segment.get_segm_name(seg) if seg else ""), - "func": ida_funcs.get_func_name(fn_ea), - "func_addr": hex(int(fn_ea))}) -result = {"comments": comments, "names": names, "scanned": scanned, - "truncated": (len(comments) >= limit or len(names) >= limit - or scanned >= max_scan)} -result -''', - # The findings journal (idatui/journal.py). A netnode blob rides along in - # the .i64, so "what did I work out here" survives closing the database. - "journal_get": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal")) -blob = n.getblob(0, "I") if ida_netnode.exist(n) else None -result = {"data": blob.decode("utf-8", "replace") if blob else ""} -result -''', - "journal_put": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal"), 0, True) -payload = (a.get("data") or "").encode("utf-8") -n.setblob(payload, 0, "I") -result = {"ok": True, "bytes": len(payload)} -result -''', - # Database-wide search (Ctrl+F), two kinds. - # - # BYTES uses IDA's own `find_bytes`, which already understands the pattern - # language people expect -- "B8 ? ? ? ? 90", nibble wildcards ("48 8? ??") - # and quoted literals -- so we neither parse nor match anything ourselves. - # Iterating is match+1, per its documented contract. - "search_bytes": r''' -import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment -pat = str(a.get("pattern", "")).strip() -limit = max(1, int(a.get("limit", 500))) -lo = int(a.get("start", 0)) -hi = int(a.get("end", 0)) or ida_idaapi.BADADDR -flags = ida_bytes.BIN_SEARCH_FORWARD | ida_bytes.BIN_SEARCH_NOSHOW -if a.get("case"): - flags |= ida_bytes.BIN_SEARCH_CASE -rows, err, ea = [], None, lo -while len(rows) < limit: - try: - hit = ida_bytes.find_bytes(pat, range_start=ea, range_end=hi, flags=flags) - except Exception as exc: - err = str(exc) or exc.__class__.__name__ - break - if hit is None or hit == ida_idaapi.BADADDR: - break - head = ida_bytes.get_item_head(hit) - fn = ida_funcs.get_func(hit) - seg = ida_segment.getseg(hit) - try: - line = ida_lines.generate_disasm_line(head, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - line = "" - rows.append({"addr": hex(int(hit)), "head": hex(int(head)), - "line": " ".join(line.split()), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) - ea = int(hit) + 1 -result = {"hits": rows, "error": err, "truncated": len(rows) >= limit} -result -''', - # TEXT walks the listing the way a person reads it: every head's rendered - # disassembly line, which is why it finds "call cs:__isoc99_scanf" and - # "0deadbeefh" alike. Bounded by max_scan, so a 400MB image reports partial - # results instead of stalling. - "search_text": r''' -import ida_lines, ida_funcs, ida_segment, idautils -import re as _re -q = str(a.get("query", "")) -limit = max(1, int(a.get("limit", 500))) -max_scan = max(1000, int(a.get("max_scan", 3000000))) -ci = (not a.get("case")) and q.islower() # smartcase, like the in-view search -rx, err = None, None -if a.get("regex"): - try: - rx = _re.compile(q, _re.I if ci else 0) - except Exception as exc: - err = "bad regex: " + str(exc) -needle = q.lower() if ci else q -rows, scanned = [], 0 -if err is None and q: - for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(rows) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(rows) >= limit or scanned >= max_scan: - break + 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: - line = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - continue - # Match what the user SEES, not IDA's column padding: nobody types - # "call" + four spaces + "cs:getenv_ptr". - line = " ".join(line.split()) - hay = line.lower() if ci else line - if (rx.search(line) if rx is not None else (needle in hay)): - fn = ida_funcs.get_func(ea) - rows.append({"addr": hex(int(ea)), "head": hex(int(ea)), - "line": line, - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": ida_segment.get_segm_name(seg)}) -result = {"hits": rows, "error": err, "scanned": scanned, - "truncated": len(rows) >= limit or scanned >= max_scan} -result -''', - "list_linkage": r''' -imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name} - for item in db.imports.get_all_imports() if item.name] -exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)} - for item in db.entries.get_all() if item.name] -result = {"imports": imports, "exports": exports, - "n_imports": len(imports), "n_exports": len(exports)} -result -''', - "lookup_funcs": r''' -rows = [] -for query in a.get("queries", []): - raw = str(query) - try: ea = int(raw, 16) - except ValueError: - fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None - else: fn = db.functions.get_at(ea) - if fn is None: - rows.append({"query": raw, "fn": None}) - else: - rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)), - "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}", - "size": int(fn.end_ea) - int(fn.start_ea)}}) -result = {"result": rows} -result -''', - "resolve_names": r''' -import ida_idaapi, ida_name -rows = [] -for query in a.get("queries", []): - name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name) - rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None}) -result = {"result": rows} -result -''', - # Ours: the coarse code/data type plus a fine `kind` (call/jump/flow, - # read/write/offset/text/info) that the xref dialog draws its badges from. - # Deliberately NOT sorted -- the dialog lists xrefs in IDA's own order. - "xref_types": r''' -import idaapi, idautils, ida_bytes, ida_funcs, ida_xref -code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", ida_xref.fl_JF: "jump", - ida_xref.fl_JN: "jump", ida_xref.fl_F: "flow"} -data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", ida_xref.dr_R: "read", - ida_xref.dr_T: "text", ida_xref.dr_I: "info"} -def _kind(xr): - return (code_kind if xr.iscode else data_kind).get(xr.type, "code" if xr.iscode else "data") -def _fn(ea): - f = ida_funcs.get_func(ea) - return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None -queries = a.get("queries") or [] -all_results = [] -for query in queries: - query = query if isinstance(query, dict) else {"addr": query} - raw = str(query.get("addr", "")).strip() - direction = str(query.get("direction", "to") or "to").lower() - include_fn = bool(query.get("include_fn", True)) - dedup = bool(query.get("dedup", True)) - try: count = int(query.get("count", 2000) or 2000) - except (TypeError, ValueError): count = 2000 - try: target = int(raw, 16) - except ValueError: target = idaapi.get_name_ea(idaapi.BADADDR, raw) - rows = [] - if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): - if direction in ("to", "both"): - for xr in idautils.XrefsTo(target, 0): - row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), - "to": hex(int(target)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} - if include_fn: row["fn"] = _fn(xr.frm) - rows.append(row) - if direction in ("from", "both"): - for xr in idautils.XrefsFrom(target, 0): - row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), - "to": hex(int(xr.to)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} - if include_fn: row["fn"] = _fn(xr.to) - rows.append(row) - if dedup: - seen, deduped = set(), [] - for r in rows: - k = (r["direction"], r["from"], r["to"], r["kind"]) - if k in seen: continue - seen.add(k); deduped.append(r) - rows = deduped - rows = rows[:count] - all_results.append({"query": raw, "data": rows, "next_offset": None}) -result = {"result": all_results} -result -''', - # Mirrors the tool ida-tui was written against, ORDER INCLUDED. The rows are - # sorted by the far-end address and deduped by default, and the pseudocode - # follow's address fallback silently depends on it: at a call site the raw - # IDA order yields the ordinary-flow xref (the next instruction) first, so an - # unsorted result makes "follow the call" land on the following line instead. - "xref_query": r''' -import idaapi, idautils, ida_bytes, ida_funcs -def _fn(ea): - f = ida_funcs.get_func(ea) - return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None -queries = a.get("queries") or [] -all_results = [] -for query in queries: - raw = str(query.get("addr", "")).strip() - direction = str(query.get("direction", "both") or "both").lower() - if direction not in ("to", "from", "both"): direction = "both" - xref_type = str(query.get("xref_type", "any") or "any").lower() - if xref_type not in ("any", "code", "data"): xref_type = "any" - include_fn = bool(query.get("include_fn", True)) - dedup = bool(query.get("dedup", True)) - sort_by = str(query.get("sort_by", "addr") or "addr") - descending = bool(query.get("descending", False)) - try: offset = max(0, int(query.get("offset", 0) or 0)) - except (TypeError, ValueError): offset = 0 - try: count = max(0, min(int(query.get("count", 200) or 200), 5000)) - except (TypeError, ValueError): count = 200 - try: - try: target = int(raw, 16) - except ValueError: - target = idaapi.get_name_ea(idaapi.BADADDR, raw) - if target == idaapi.BADADDR: raise ValueError(f"Failed to resolve address/name: {raw}") - if not ida_bytes.is_mapped(target): raise ValueError(f"Address not mapped: {raw}") - rows = [] - if direction in ("to", "both"): - for xr in idautils.XrefsTo(target, 0): - kind = "code" if xr.iscode else "data" - if xref_type != "any" and kind != xref_type: continue - row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), - "to": hex(int(target)), "type": kind} - if include_fn: row["fn"] = _fn(xr.frm) - rows.append(row) - if direction in ("from", "both"): - for xr in idautils.XrefsFrom(target, 0): - kind = "code" if xr.iscode else "data" - if xref_type != "any" and kind != xref_type: continue - row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), - "to": hex(int(xr.to)), "type": kind} - if include_fn: row["fn"] = _fn(xr.to) - rows.append(row) - if dedup: - seen, deduped = set(), [] - for row in rows: - key = (row["direction"], row["from"], row["to"], row["type"]) - if key in seen: continue - seen.add(key); deduped.append(row) - rows = deduped - if sort_by == "type": - rows.sort(key=lambda r: (str(r.get("type", "")), int(str(r["addr"]), 16)), reverse=descending) - else: - rows.sort(key=lambda r: int(str(r["addr"]), 16), reverse=descending) - page = rows[offset:offset + count] if count else rows[offset:] - nxt = offset + len(page) - all_results.append({"target": raw, "resolved_addr": hex(int(target)), "direction": direction, - "xref_type": xref_type, "data": page, - "next_offset": nxt if nxt < len(rows) else None, - "total": len(rows), "error": None}) - except Exception as exc: - all_results.append({"target": raw, "resolved_addr": None, "direction": direction, - "xref_type": xref_type, "data": [], "next_offset": None, - "total": 0, "error": str(exc)}) -result = {"result": all_results} -result -''', - # A comment must land in BOTH views, and the pseudocode half is not a - # simple set: db.comments.set_at() alone leaves the pseudocode unchanged. - # Hex-Rays comments are anchored to a ctree location (treeloc_t), and an - # anchor the ctree does not actually own is dropped as an "orphan" -- so the - # itp slot has to be searched until one sticks, exactly as IDA's own UI does. - # Without it a comment silently never appears in the decompilation. - "set_comments": r''' -import idaapi, idc, ida_hexrays -rows = [] -for item in a.get("items", []): - addr_s = str(item.get("addr", "")) - text = str(item.get("comment") or "") - try: - ea = int(addr_s, 16) - if not idaapi.set_cmt(ea, text, False): - rows.append({"addr": addr_s, - "error": f"Failed to set disassembly comment at {hex(ea)}"}) - continue - if not ida_hexrays.init_hexrays_plugin(): - rows.append({"addr": addr_s}); continue - try: - cfunc = ida_hexrays.decompile(ea) - except Exception: - cfunc = None - if cfunc is None: - rows.append({"addr": addr_s}); continue - if ea == cfunc.entry_ea: - # The signature line carries no ctree item: it is a function comment. - idc.set_func_cmt(ea, text, True) - cfunc.refresh_func_ctext() - rows.append({"addr": addr_s}); continue - eamap = cfunc.get_eamap() - if ea not in eamap: - rows.append({"addr": addr_s, - "error": f"Failed to set decompiler comment at {hex(ea)}"}) - continue - nearest_ea = eamap[ea][0].ea - if cfunc.has_orphan_cmts(): - cfunc.del_orphan_cmts(); cfunc.save_user_cmts() - tl = idaapi.treeloc_t(); tl.ea = nearest_ea - placed = False - for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON): - tl.itp = itp - cfunc.set_user_cmt(tl, text) - cfunc.save_user_cmts() - cfunc.refresh_func_ctext() - if not cfunc.has_orphan_cmts(): - placed = True; break - cfunc.del_orphan_cmts(); cfunc.save_user_cmts() - rows.append({"addr": addr_s} if placed else - {"addr": addr_s, - "error": f"Failed to set decompiler comment at {hex(ea)}"}) - except Exception as exc: - rows.append({"addr": addr_s, "error": str(exc)}) -result = {"result": rows} -result -''', - # Every category takes EITHER one edit or a LIST of them, and the answer is - # one row per edit. The port accepted only a single dict, so any batch path - # (rpc rename_many applying a whole symbol file, which is the entire point of - # that verb) died with "list indices must be integers or slices, not str" and - # reported the failure against addr=null. Mirrors the real tool: conflict - # detection before the write, dry_run/allow_overwrite/stop_on_error, per-row - # addr/old/name, and a summary counting EDITS rather than categories. - "rename": r''' -import idaapi, ida_hexrays, ida_name -batch = a.get("batch") or {} -dry_run = bool(batch.get("dry_run", False)) -allow_overwrite = bool(batch.get("allow_overwrite", False)) -stop_on_error = bool(batch.get("stop_on_error", False)) - -def _items(value): - if value is None: return [] - if isinstance(value, dict): return [value] - if isinstance(value, list): return [i for i in value if isinstance(i, dict)] - return [] - -def _set_name_checked(ea, new): - conflict = idaapi.get_name_ea(idaapi.BADADDR, new) - if conflict != idaapi.BADADDR and conflict != ea and not allow_overwrite: - return False, f"can't rename at {hex(ea)} as {new!r}: name already used at {hex(conflict)}" - if dry_run: - return True, None - flags = idaapi.SN_CHECK - if allow_overwrite: flags |= int(getattr(idaapi, "SN_FORCE", 0)) - if not idaapi.set_name(ea, new, flags): - return False, (f"Rename failed at {hex(ea)}: IDA rejected name {new!r} " - "(invalid identifier or internal conflict)") - return True, None - -def _refresh_ctext(fn_addr): - # A renamed function must invalidate Hex-Rays' cache, which is per function - # and persisted in the .i64: without this the pseudocode keeps calling the - # old name forever while every other readback reports the new one. - if not ida_hexrays.init_hexrays_plugin(): return - failure = ida_hexrays.hexrays_failure_t() - cfunc = ida_hexrays.decompile_func(fn_addr, failure, ida_hexrays.DECOMP_WARNINGS) - if cfunc: cfunc.refresh_func_ctext() - -out = {}; ok_count = failed = 0; halted = False -for category in ("func", "data", "local", "stack"): - if category not in batch: continue - rows = [] - for edit in _items(batch.get(category)): - try: - if category == "func": - addr_text = edit.get("addr") or edit.get("func_addr") or edit.get("func") - new = edit.get("name") or edit.get("new") or edit.get("new_name") - if not addr_text or not new: - row = {"addr": addr_text, "name": new, - "error": "Function rename requires addr + name"} - else: - ea = int(str(addr_text), 16) - fn = idaapi.get_func(ea) - if fn is None: - row = {"addr": addr_text, "name": new, "error": "Function not found"} - else: - old = idaapi.get_name(fn.start_ea) or None - ok, err = _set_name_checked(fn.start_ea, str(new)) - row = {"addr": addr_text, "old": old, "name": str(new)} - if err: row["error"] = err - if dry_run: row["dry_run"] = True - if ok and not dry_run: _refresh_ctext(fn.start_ea) - elif category == "data": - addr_text = edit.get("addr") - old = edit.get("old") or edit.get("old_name") - new = edit.get("new") or edit.get("new_name") or edit.get("name") - if not new and new != "": - row = {"old": old, "new": None, - "error": "Global rename requires target and new name"} - else: - if addr_text is not None: - ea = int(str(addr_text), 16) - old = old or (idaapi.get_name(ea) or None) - else: - ea = idaapi.get_name_ea(idaapi.BADADDR, str(old or "")) - if ea == idaapi.BADADDR: - row = {"old": old, "new": str(new), "error": f"Global {old!r} not found"} - else: - # An empty new name CLEARS the label; that is a real - # request (tests revert with it), not a missing argument. - if str(new) == "": - ok = bool(ida_name.set_name(ea, "", idaapi.SN_CHECK)) - err = None if ok else f"Failed to clear the name at {hex(ea)}" - else: - ok, err = _set_name_checked(ea, str(new)) - row = {"addr": hex(ea), "old": old, "new": str(new)} - if err: row["error"] = err - if dry_run: row["dry_run"] = True - else: - fa, old, new = edit.get("func_addr"), edit.get("old"), edit.get("new") - if not fa or not old or not new: - row = {"old": old, "new": new, - "error": f"{category} rename requires func_addr + old + new"} - else: - ea = int(str(fa), 16) - pseudo = db.pseudocode.decompile(ea) - var = pseudo.find_local_variable(str(old)) - if var is None: - row = {"func_addr": fa, "old": old, "new": new, - "error": f"no local {old!r} in that function"} - elif dry_run: - row = {"func_addr": fa, "old": old, "new": new, "dry_run": True} - else: - var.set_user_name(str(new)) - ok = bool(pseudo.save_local_variable_info(var, save_name=True)) - row = {"func_addr": fa, "old": old, "new": new} - if not ok: row["error"] = "IDA rejected the local variable name" - except Exception as exc: - row = {"addr": edit.get("addr"), "error": str(exc)} - rows.append(row) - if row.get("error"): failed += 1 - else: ok_count += 1 - if row.get("error") and stop_on_error: - halted = True; break - out[category] = rows - if halted: break -out["summary"] = {"ok": ok_count, "failed": failed} -if dry_run: out["summary"]["dry_run"] = True -if halted: out["summary"]["halted"] = True -result = out -result -''', -} - - -_OPERATIONS["define_code_run"] = r''' -import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi -ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) -seg = ida_segment.getseg(ea) -if seg is None: - result = {"addr": a["addr"], "error": "no segment", "count": 0} -else: - start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) - while count < limit: - if ea >= hi: stopped = "segment"; break - flags = ida_bytes.get_flags(ea) - if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break - size = int(ida_ua.create_insn(ea)) - if size <= 0: stopped = "undecodable"; break - count += 1 - insn = ida_ua.insn_t() - if ida_ua.decode_insn(insn, ea) > 0: - try: is_ret = bool(ida_idp.is_ret_insn(insn)) - except Exception: is_ret = False - if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): - ea += size; stopped = "flow"; break - ea += size - result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} -result -''' - - -_OPERATIONS["define_func_run"] = r''' -import ida_bytes, ida_funcs, ida_segment -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is not None and int(fn.start_ea) == ea: - result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} -else: - automatic = bool(db.functions.create(ea)) - if not automatic: - seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea - while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): - nxt = int(ida_bytes.get_item_end(end)) - if nxt <= end: break - end = nxt - ok = bool(end > ea and ida_funcs.add_func(ea, end)) - else: ok = True - fn = db.functions.get_at(ea) - result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), - "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} - if ok and fn is not None else - {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) -result -''' - - -_OPERATIONS["set_thumb"] = r''' -import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs -ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") -seg = ida_segment.getseg(ea) -if treg is None or treg < 0: - result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} -elif seg is None: - result = {"addr": hex(ea), "error": "no segment"} -else: - current = ida_segregs.get_sreg(ea, treg) - current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) - want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) - changed = False - if want and seg.bitness != 1: - ida_segment.set_segm_addressing(seg, 1); changed = True - size = max(int(ida_bytes.get_item_size(ea)), 2) - ida_bytes.del_items(ea, 0, size) - ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) - now = ida_segregs.get_sreg(ea, treg) - result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, - "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, - "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} -result -''' - - -_OPERATIONS["thumb_scan"] = r''' -import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua -lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) -apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) -treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo -while cursor + 4 <= hi and len(found) < limit: - at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 - if not value & 1: continue - target = value & ~1; seg = ida_segment.getseg(target) - if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue - flags = ida_bytes.get_flags(target) - if ida_bytes.is_data(flags): continue - item = {"at": hex(at), "value": hex(value), "target": hex(target), - "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) - if not apply: continue - if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) - if not ida_bytes.is_code(ida_bytes.get_flags(target)): - ida_bytes.del_items(target, 0, 2) - if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue - item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 -result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} -result -''' - - -_OPERATIONS["decomp_error"] = r''' -import ida_hexrays, ida_ida -ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) -result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} -if fn is None: - result["reason"] = "no function here" -else: - try: - failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) - if cfunc is not None: result["reason"] = "" - else: - result.update({"reason": failure.desc() or f"error {failure.code}", - "code": int(failure.code), "errea": hex(int(failure.errea))}) - except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" -result -''' - -# `heads` and the operand-format tools are the port's IDAPython island: the -# continuous listing's presentation model (undefined runs, colour spans, operand -# extents, banners, struct members, the digest protocol) and IDA/Hex-Rays number -# formats have no ida-domain surface. Rather than paraphrase ~1100 lines of -# performance-tuned, behaviour-sensitive code into string literals, they stay -# real, diffable source in idatui/remote_tools.py and are shipped to the database -# process as text. Read once at import; the file ships beside this module. -_REMOTE_LIB = (Path(__file__).with_name("remote_tools.py")).read_text(encoding="utf-8") - -#: Versioned by content, so editing remote_tools.py re-installs it instead of -#: silently running the copy a long-lived worker already has. -_REMOTE_MODULE = "_idatui_remote_" + hashlib.sha1( - _REMOTE_LIB.encode("utf-8")).hexdigest()[:12] - -#: Sent back when the database process has not got the library yet; the client -#: installs it and retries once. Amortised, a worker receives it exactly once. -_NEED_LIB = "__idatui_needs_remote_lib__" - -#: Installs the library as a real module in the database process. Persisting it -#: in sys.modules is what makes the module-level caches (the tag maps, and the -#: line-render lru_cache the listing's throughput depends on) survive between -#: calls -- execute_python builds a fresh namespace every time, so a library -#: exec'd inline is rebuilt, and its caches thrown away, on every single call. -_INSTALL_LIB = f''' -import sys, types -_m = types.ModuleType({_REMOTE_MODULE!r}) -exec(compile(a["source"], {_REMOTE_MODULE!r}, "exec"), _m.__dict__) -sys.modules[{_REMOTE_MODULE!r}] = _m -result = True -result -''' - - -def _remote_op(call: str) -> str: - """A snippet that calls one of the carried-over tools by its real signature. - - Costs one short request: the library is imported from the database process's - own sys.modules, not shipped again. - """ - return (f"import sys\n" - f"_m = sys.modules.get({_REMOTE_MODULE!r})\n" - f"result = {{{_NEED_LIB!r}: True}} if _m is None else _m.{call}\n" - f"result\n") - - -_OPERATIONS["op_format"] = _remote_op( - 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' - ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') -_OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') -_OPERATIONS["decompile"] = _remote_op( - 'decompile(addr=a["addr"],' - ' include_addresses=bool(a.get("include_addresses", True)))') -_OPERATIONS["decomp_map"] = _remote_op('decomp_map(addr=a["addr"])') -_OPERATIONS["pc_num_format"] = _remote_op( - 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' - ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' - ' ea=a.get("ea", ""), opnum=int(a.get("opnum", -1)))') - -# The listing walker itself. Replaces the port's re-implementation, which -# rendered no per-operand extents (so no keypress could say which literal it -# would reformat) and had no digest/expect support (so every page was re-sent -# after any edit), and whose span walk was the per-character loop our own -# version had already been rewritten to avoid. -#: Row count + seek anchors for a whole segment, in ONE call. See -#: remote_tools.segment_index: the alternative is fetching every row. -_OPERATIONS["segment_index"] = _remote_op( - 'segment_index(addr=a["addr"], end=a.get("end", ""),' - ' page_rows=int(a.get("page_rows", 500)), detail=bool(a.get("detail", False)))') - -_HEADS = _remote_op( - 'heads(addr=a["addr"], count=int(a.get("count", 200)),' - ' offset=int(a.get("offset", 0)), end=a.get("end", ""),' - ' back=bool(a.get("back", False)), annotate=bool(a.get("annotate", False)),' - ' expect=a.get("expect", ""), text=bool(a.get("text", True)))') - - -# The graph view's only backend call. Blocks are address RANGES, never text: -# the client re-renders them with `heads`, so boxes reuse the exact listing rows -# (colours, operand marks, trail painting) instead of growing a second renderer. -# -# ida-domain exposes no basic-block/edge-kind surface, so this stays on ida_gdl. -_OPERATIONS["flowchart"] = r''' -import ida_funcs, ida_gdl -ea = int(str(a["addr"]), 16) -fn = ida_funcs.get_func(ea) -if fn is None: - result = {"addr": hex(ea), "error": "no function at that address", "blocks": []} -else: - fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS) - index, order = {}, [] - for bb in fc: - index[bb.start_ea] = len(order) - order.append(bb) - blocks = [] - for bb in order: - sl = [s for s in bb.succs() if s.start_ea in index] - succs = [] - for s in sl: - # Edge kind is what the graph view colours by: an n-way dispatch is - # "switch", a successor that is literally the next address falls - # through, anything else is a taken branch. - if len(sl) > 2: kind = "switch" - elif s.start_ea == bb.end_ea: kind = "fall" - else: kind = "jump" - succs.append([index[s.start_ea], kind]) - blocks.append({"id": index[bb.start_ea], "start": hex(int(bb.start_ea)), - "end": hex(int(bb.end_ea)), "succs": succs}) - result = {"addr": hex(ea), - "func": {"addr": hex(int(fn.start_ea)), "end": hex(int(fn.end_ea)), - "name": ida_funcs.get_func_name(fn.start_ea) or ""}, - "entry": index.get(fn.start_ea, 0), "blocks": blocks} -result -''' - -# Only ever reached as domain.py's fallback when file_regions yields nothing. -_OPERATIONS["survey_binary"] = r''' -segments = [] -for seg in db.segments.get_all(): - segments.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), - "name": db.segments.get_name(seg) or ""}) -result = {"segments": segments} -result -''' + 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: @@ -1277,7 +287,9 @@ class CodeModeClient: 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._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 @@ -1289,10 +301,17 @@ class CodeModeClient: def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": _require_codemode() with self._connect_lock: - if self._handle is not None and self._handle.connected: - return self + 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)}…") + 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 @@ -1321,7 +340,9 @@ class CodeModeClient: if not self._new_database or time.monotonic() >= deadline: raise if progress: - progress("waiting for the previous Code Mode lease to close…") + progress( + "waiting for the previous Code Mode lease to close…" + ) owner = find_database_owner( self._path, output_database=self._output_database, @@ -1336,7 +357,9 @@ class CodeModeClient: time.sleep(0.2) if progress: backend = handle.instance.backend - progress(f"attached to {backend} database; waiting for auto-analysis…") + 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 @@ -1360,61 +383,60 @@ class CodeModeClient: def backend(self) -> str | None: return self._handle.instance.backend if self._handle is not None else None - def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + 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: - response = handle.execute_python(code, timeout=timeout) + return remote(handle, **args) except RemoteError as exc: - details = exc.details or {} message = str(exc) - if details.get("traceback"): - message += f"\n{details['traceback']}" + if exc.details.get("traceback"): + message += f"\n{exc.details['traceback']}" if exc.code == "operation_timeout": raise IDATimeoutError(message) from exc - raise IDAToolError("execute_python", message) from exc + raise IDAToolError(name, message) from exc except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: raise self._connection_error(exc) from exc - if not isinstance(response, dict) or "result" not in response: - raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") - return response["result"] - - @staticmethod - def _unpack(answer: Any) -> Any: - """Undo _PACK_EPILOGUE. Anything else passes through untouched.""" - if isinstance(answer, dict) and _PACKED in answer: - return json.loads(answer[_PACKED]) - return answer - - def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: - """Execute one TUI domain operation through Code Mode.""" - if operation in ("idb_save", "save"): - return self.save_database() - if operation in ("server_health", "ping", "health", "state"): - return self.health() - body = _HEADS if operation == "heads" else _OPERATIONS.get(operation) - if body is None: - raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") - try: - answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) - if isinstance(answer, dict) and answer.get(_NEED_LIB): - # First call against this database process (or a restarted one). - self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), - timeout=timeout) - answer = self._unpack( - self.execute_python(_script(args, body), timeout=timeout)) - return answer - except IDAToolError as exc: - if exc.tool == "execute_python": - raise IDAToolError(operation, exc.message) from exc - raise - - # Temporary source compatibility for external drivers/tests that used the - # old WorkerClient. Application code uses the accurately named invoke(). - call = invoke def save_database(self) -> dict[str, Any]: if not self.connected: @@ -1429,6 +451,35 @@ class CodeModeClient: 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() @@ -1463,8 +514,14 @@ class CodeModeClient: 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)] + 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 -- cgit v1.3.1-sl0p