diff options
Diffstat (limited to 'idatui')
| -rw-r--r-- | idatui/__init__.py | 5 | ||||
| -rw-r--r-- | idatui/app.py | 4220 | ||||
| -rw-r--r-- | idatui/codemode_client.py | 1550 | ||||
| -rw-r--r-- | idatui/diag.py | 121 | ||||
| -rw-r--r-- | idatui/domain.py | 964 | ||||
| -rw-r--r-- | idatui/drive.py | 103 | ||||
| -rw-r--r-- | idatui/edit_ctl.py | 738 | ||||
| -rw-r--r-- | idatui/errors.py | 10 | ||||
| -rw-r--r-- | idatui/findings.py | 384 | ||||
| -rw-r--r-- | idatui/formats.py | 15 | ||||
| -rw-r--r-- | idatui/graph.py | 750 | ||||
| -rw-r--r-- | idatui/highlight.py | 156 | ||||
| -rw-r--r-- | idatui/journal.py | 107 | ||||
| -rw-r--r-- | idatui/kittygfx.py | 271 | ||||
| -rw-r--r-- | idatui/launch.py | 121 | ||||
| -rw-r--r-- | idatui/pane.py | 456 | ||||
| -rw-r--r-- | idatui/pool.py | 100 | ||||
| -rw-r--r-- | idatui/project.py | 28 | ||||
| -rw-r--r-- | idatui/prompt.py | 117 | ||||
| -rw-r--r-- | idatui/remote_tools.py | 1548 | ||||
| -rw-r--r-- | idatui/rpc.py | 604 | ||||
| -rw-r--r-- | idatui/rpcclient.py | 10 | ||||
| -rw-r--r-- | idatui/search.py | 112 | ||||
| -rw-r--r-- | idatui/trace.py | 495 | ||||
| -rw-r--r-- | idatui/trace_ctl.py | 421 | ||||
| -rw-r--r-- | idatui/worker.py | 233 | ||||
| -rw-r--r-- | idatui/worker_client.py | 234 |
27 files changed, 11760 insertions, 2113 deletions
diff --git a/idatui/__init__.py b/idatui/__init__.py index 2cbdde8..7da28b3 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -1,5 +1,4 @@ -"""idatui — a minimal keyboard-first TUI for IDA Pro, driving idalib via a -private unix-socket worker (idatui.worker / WorkerClient).""" +"""idatui — a keyboard-first TUI using shared IDA Code Mode databases.""" from .errors import ( IDAError, @@ -11,6 +10,7 @@ from .errors import ( IDASessionError, Session, ) +from .codemode_client import CodeModeClient from .domain import ( Program, FunctionIndex, @@ -25,6 +25,7 @@ from .domain import ( ) __all__ = [ + "CodeModeClient", "Program", "FunctionIndex", "DisasmModel", diff --git a/idatui/app.py b/idatui/app.py index dd7c222..8baff74 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -10,8 +10,8 @@ Design notes: without ever materializing 52k lines in a widget. * All network/domain work runs in Textual worker threads; the UI never blocks. * An address-history stack backs Enter (follow) / Esc (back), IDA-style. -* On startup we bump the worker idle-TTL and run a keepalive heartbeat so the - session never gets reaped while we chill. +* Database lifecycle is lease-based through ida_codemode: matching GUI sessions + are reused, otherwise a shared managed idalib worker is opened on demand. """ from __future__ import annotations @@ -20,7 +20,9 @@ import asyncio import os import re import subprocess +import time from dataclasses import dataclass, field +from enum import StrEnum from rich.align import Align from rich.segment import Segment @@ -30,7 +32,7 @@ from textual import work from textual.app import App, ComposeResult from textual.binding import Binding from textual.command import DiscoveryHit, Hit, Provider -from textual.containers import Grid, Horizontal, Vertical, VerticalScroll +from textual.containers import Horizontal, Vertical, VerticalScroll from textual.geometry import Region, Size from textual.message import Message from textual.reactive import reactive @@ -43,16 +45,40 @@ from textual.widgets import ( ) from textual.widgets.option_list import Option -from .highlight import highlight_c +from . import graph +from . import kittygfx +from .edit_ctl import EditController +from .prompt import PromptBar +from .trace_ctl import TraceController +from . import findings, search +from .highlight import CTextArea, highlight_c +from .journal import Journal from .errors import IDAToolError, IDAConnectionError -from .worker_client import WorkerClient -from .domain import DisasmModel, Func, Head, ListingModel, Program, Struct +from .codemode_client import CodeModeClient, registered_database +from .domain import Func, Head, ListingModel, Program, Struct # Styles for the disassembly listing. _S_ADDR = Style(color="#6b7684") _S_LABEL = Style(color="#7aa2f7", bold=True) _S_INSN = Style(color="#c3cad3") +#: IDA's token kinds -> the measured palette. The rule that keeps a dense +#: disassembly readable: NEUTRALS for the machine (mnemonic brightest because you +#: scan down that column, registers at body weight because they're most of the +#: text), HUES only where they mean something (numbers, strings, symbols), +#: structure recedes so brackets and commas stop competing with operands. +_S_SPAN = { + "insn": Style(color="#e8ecf2"), # 15.3:1 mnemonic / directive + "reg": Style(color="#c3cad3"), # 11.0:1 registers = body weight + "num": Style(color="#d8a657"), # 8.2:1 immediates, offsets + "str": Style(color="#9ece6a"), # 9.9:1 string literals + "name": Style(color="#7aa2f7"), # 7.2:1 symbols / xref targets + "seg": Style(color="#93aee0"), # 8.1:1 segment names + "cmt": Style(color="#7c8b9e", italic=True), # 5.2:1 + "punct": Style(color="#626c7a"), # 3.4:1 brackets, commas, +/- + "err": Style(color="#c9762f"), # IDA's own error marker + "text": Style(color="#c3cad3"), # 11.0:1 anything unclassified +} _S_MNEM = Style(color="#e8ecf2") _S_OPBYTES = Style(color="#5e6875") # raw opcode bytes column _S_DATA = Style(color="#d8a657") @@ -66,6 +92,42 @@ _OP_LIMIT = 8 # opcode bytes shown in the 'limited' column mode _JUMP_CONTEXT = 4 # lines of context kept above a jump target (cursor stays on it) _SPLIT_MIN_WIDTH = 100 # need room for two usable code panes side by side + +class ViewMode(StrEnum): + """Which pane is showing (in split: which one has focus). ``IdaTui._active``. + + A ``StrEnum`` rather than an enum, deliberately: ``_active`` is handed + straight to drivers as ``cursor.kind`` over the RPC socket, and the pilot + suite compares it to plain strings. Members ARE strings, so every existing + comparison and every JSON payload keeps working -- what this buys is one + place that says which modes exist, and an AttributeError instead of silence + when one is misspelled. + + There used to be a fifth value, ``"disasm"``, assigned on exactly one path + (a decompile that failed with nowhere to return to) and meaning the same + widget as ``LISTING``. Four sites handled it and five compared against + ``"listing"`` alone, so it silently took the wrong branch in half the app -- + Tab out of a failed decompile flipped to the listing instead of retrying the + decompiler, and rpc.py carried a workaround for a mode change that never + arrived. It is gone; the failure path lands on LISTING like every other + route into the listing. + + Adding a mode means auditing every ``_active`` comparison. Prefer the + predicates on IdaTui (``is_listing``/``in_code``/...) over bare ``==`` so + the next one has fewer places to reach. + """ + + LISTING = "listing" # the unified continuous listing (code + data) + DECOMP = "decomp" # Hex-Rays pseudocode + HEX = "hex" # the hex viewer + GRAPH = "graph" # the CFG graph view + + #: The two that show a code view over a NavEntry, i.e. where a follow, an + #: xref or a rename makes sense. + @classmethod + def code_modes(cls) -> frozenset["ViewMode"]: + return frozenset({cls.LISTING, cls.DECOMP, cls.GRAPH}) + # Tokens that look like identifiers but aren't renamable symbols (so 'n' on them # in the listing names the address instead of trying to rename the token). _ASM_KEYWORDS = frozenset({ @@ -74,11 +136,27 @@ _ASM_KEYWORDS = frozenset({ "gs", "ss", "align", "public", "assume", "end", }) _S_CURSOR = Style(bgcolor="#2a313c") +#: Execution trails. Deliberately faint: they sit UNDER the code palette and +#: must not compete with it — the trail says "you came through here", the text +#: still has to be readable as code. Now is the loudest because there is exactly +#: one of it. +_S_TRAIL_NOW = Style(bgcolor="#3f3410") +_S_TRAIL_PAST = Style(bgcolor="#2b1c17") # warm: behind you +_S_TRAIL_FUTURE = Style(bgcolor="#152230") # cool: ahead of you +#: Hex with a trace loaded: bytes the trace SAW at this timestamp vs bytes we're +#: still showing from the file. The distinction matters more than the values — +#: one is evidence, the other is an assumption. +_S_HEX_LIVE = Style(color="#9ece6a") +_S_HEX_STALE = Style(color="#5e6875") _S_DIM = Style(color="#7c8b9e", italic=True) _S_MATCH = Style(bgcolor="#7a5c00") # all search matches _S_MATCH_CUR = Style(bgcolor="#d0a215", color="#12161c") # the current match _S_NAME_MATCH = Style(bgcolor="#d0a215", color="#12161c") # filter match in a name _S_WORD = Style(bgcolor="#2a3f5f") # identifier under the cursor +#: The operand/literal under the cursor — what `o` would reformat. Distinct from +#: _S_WORD (which marks every occurrence of an identifier): this marks ONE span, +#: the thing a keypress acts on, so it reads as a selection rather than a match. +_S_OPERAND = Style(bgcolor="#3a3560", underline=True) _S_CELL = Style(reverse=True) # the block cursor cell _S_LINENO = Style(color="#626c7a") # pseudocode line-number gutter _S_LINENO_CUR = Style(color="#c3cad3", bold=True) # gutter on the cursor line @@ -97,15 +175,15 @@ _ADDR_MARK_STRIP_RE = re.compile(r"\s*/\*\s*0x[0-9A-Fa-f]+\s*\*/") @dataclass class BinaryState: """Everything that makes one project binary's session resumable across a - switch. Addresses outlive the worker, so nav history survives eviction; the - Program/index only survive while that worker is still resident.""" + switch. Addresses outlive a database lease, so nav history survives eviction; + the Program/index only survive while that lease remains resident.""" label: str program: object | None = None func_index: object | None = None nav: list = field(default_factory=list) cur: object | None = None - active: str = "listing" + active: ViewMode = ViewMode.LISTING split: bool = False filter_term: str = "" dirty: bool = False @@ -132,6 +210,8 @@ class ViewAnchor: top_ea: int | None = None # first visible address cursor_x: int = 0 flash: str | None = None + #: The edit changed which functions exist, so the index must be rebuilt. + refresh_functions: bool = False @dataclass @@ -218,6 +298,17 @@ class MakeDataRequested(Message): self.view = view +class OpFormatRequested(Message): + """A code view asks to change how the literal under the cursor is DISPLAYED + (IDA's 'o'): hex, decimal, binary, character, offset. ``mode`` is 'cycle', + 'back', or a format by name.""" + + def __init__(self, view, mode: str = "cycle") -> None: # type: ignore[no-untyped-def] + super().__init__() + self.view = view + self.mode = mode + + class NavMixin: """Follow / xrefs actions shared by the code views (bindings live on each view since Textual only merges BINDINGS from DOMNode subclasses).""" @@ -503,13 +594,22 @@ class ColumnCursor: self._page_scroll(direction * (self._visible_height() // 2)) def _apply_scroll(self, y: int, x: int = 0) -> None: - """Set the scroll offset reliably after a (re)load. Applied now and again - after the next refresh — when the view was just shown its size isn't - computed yet, so an immediate scroll_to clamps to 0; the deferred pass - re-applies it and forces a repaint so the pane never shows a stale frame. + """Set the scroll offset reliably after a (re)load. Applied now and, if + that didn't take, again after the next refresh — when the view was just + shown its size isn't computed yet, so an immediate scroll_to clamps to 0; + the deferred pass re-applies it and forces a repaint so the pane never + shows a stale frame. + + The deferred pass is only scheduled when the scroll actually clamped. + ``refresh(layout=True)`` re-arranges the whole screen, and paying that on + every scroll that already landed cost ~25% of the time it takes to move + through a view. """ y = max(0, y) self.scroll_to(x=x, y=y, animate=False) + off = self.scroll_offset + if round(off.y) == y and round(off.x) == x: + return # max_scroll was current, the offset is already where we want def _fix(yy: int = y, xx: int = x) -> None: self.scroll_to(x=xx, y=yy, animate=False) @@ -518,6 +618,73 @@ class ColumnCursor: self.call_after_refresh(_fix) +class _MatchRanges: + """Which lines matched, plus where in each line — the *where* computed lazily. + + Searching a big segment for one character matches most of it: `c` over bash + hits 177 000 lines at 310 000 places. Building a range list for every one of + them costs more than finding them did, and all but the forty on screen are + thrown away unread. + + So the line set is eager (search needs it to count and to jump) and the + offsets within a line are worked out when that line is painted, or when the + cursor lands on it, and cached from then on. + + Quacks like the ``{line: [(start, end), ...]}`` dict it replaced: ``in``, + ``get``, ``[]``, ``items``, ``len``. Assigning a plain ``{}`` to reset stays + valid, because every reader only uses that same subset. + """ + + __slots__ = ("_lines", "_needle", "_n", "_ci", "_text", "_cache") + + def __init__(self, lines, needle: str, n: int, ci: bool, text) -> None: + self._lines = lines # set[int] + self._needle = needle # already case-folded when ci + self._n = n # len(term); the needle may be folded + self._ci = ci + self._text = text # callable: line index -> str | None + self._cache: dict[int, list[tuple[int, int]]] = {} + + def _find(self, i: int) -> list[tuple[int, int]]: + s = self._text(i) + if not s: + return [] + hay = s.lower() if self._ci else s + needle, n = self._needle, self._n + out = [] + j = hay.find(needle) + while j >= 0: + out.append((j, j + n)) + j = hay.find(needle, j + n) + return out + + def __contains__(self, i) -> bool: + return i in self._lines + + def __len__(self) -> int: + return len(self._lines) + + def __iter__(self): + return iter(self._lines) + + def get(self, i, default=None): + if i not in self._lines: + return default + got = self._cache.get(i) + if got is None: + got = self._cache[i] = self._find(i) + return got + + def __getitem__(self, i): + got = self.get(i) + if got is None: + raise KeyError(i) + return got + + def items(self): + return ((i, self.get(i)) for i in sorted(self._lines)) + + class SearchMixin: """Vim-style in-view search shared by the disasm and pseudocode views. @@ -540,6 +707,80 @@ class SearchMixin: _matches: list[int] _ranges: dict[int, list[tuple[int, int]]] + #: What ``_matches`` was last computed from: (term, case-fold, row count, + #: id(line source)). Lets an as-you-type search narrow the previous hits + #: instead of rescanning the segment; see :meth:`_compute_matches`. Any of + #: those changing under us discards it. + _matched_key: tuple | None = None + + #: The searchable body as ONE string (see :meth:`_search_haystack`), its + #: case-folded twin, the character offset each line starts at, and the key + #: they were built for. Dropped when the search ends. + _hay_key: tuple | None = None + _hay: tuple | None = None + + def _search_source_id(self) -> int: + """Identity of whatever supplies the line text. Changes when the view is + pointed at a different model/body, which invalidates a narrowing.""" + return id(getattr(self, "model", None) or getattr(self, "_texts", None)) + + def _reset_search_cache(self, body: bool = False) -> None: + """Forget the narrowing key, and with ``body=True`` the joined body too. + + Every place that resets ``_matches``/``_ranges`` must call this: a stale + prefix would make the next search narrow from an empty list. + + The body is a different question. It is keyed by (row count, line source) + so it invalidates itself when the view is pointed somewhere else or more + rows stream in — which means ending a search does NOT have to throw it + away, and the next `/` over the same segment is then instant instead of + re-joining a quarter of a million lines. It DOES have to go when the + plain text of a row changes without either of those moving, which is + exactly what toggling the opcode-bytes column does. + """ + self._matched_key = None + if body: + self._hay_key = None + self._hay = None + + def _search_haystack(self, count: int, src: int): + """``(starts, blob, blob_folded)`` for the whole body, or None. + + Every line joined with newlines, so finding a term is one C-level + ``str.find`` walk over the segment instead of a python loop that rebuilds + and case-folds 200 000 lines per keystroke. ``starts[i]`` is where line + ``i`` begins; a term can never straddle a line because it cannot contain + a newline. + + Returns None (and the caller falls back to the per-line loop) if folding + changes the length — a couple of unicode codepoints grow when lowered, + and then every offset after them would be wrong. + """ + key = (count, src) + if self._hay_key == key: + return self._hay + starts: list[int] = [] + parts: list[str] = [] + pos = 0 + chunk = 4096 + for base in range(0, count, chunk): + for s in self._search_line_texts(base, min(chunk, count - base)): + if not s: + s = "" + starts.append(pos) + parts.append(s) + pos += len(s) + 1 + while len(starts) < count: # a short window: keep the indices lined up + starts.append(pos) + parts.append("") + pos += 1 + blob = "\n".join(parts) + folded = blob.lower() + hay = None if len(folded) != len(blob) else (starts, blob, folded) + self._hay_key = key + self._hay = hay + return hay + # --- hooks a subclass implements --- def _search_line_count(self) -> int: raise NotImplementedError @@ -547,6 +788,12 @@ class SearchMixin: def _search_line_text(self, i: int) -> str | None: raise NotImplementedError + def _search_line_texts(self, start: int, count: int) -> list: + """``count`` line texts from ``start``. Overridable so a view whose rows + come from a locked model can fetch a window in one go.""" + text_of = self._search_line_text + return [text_of(i) for i in range(start, start + count)] + def _search_ensure(self, done) -> None: """Ensure all line texts are available, then call ``done()`` on the UI thread. Default: assume ready.""" @@ -571,6 +818,7 @@ class SearchMixin: self._term = "" self._matches = [] self._ranges = {} + self._reset_search_cache() self.cursor = getattr(self, "_search_origin", self.cursor) self.cursor_x = getattr(self, "_search_origin_x", self.cursor_x) self.refresh() @@ -609,6 +857,7 @@ class SearchMixin: self._term = "" self._matches = [] self._ranges = {} + self._reset_search_cache() self.cursor = getattr(self, "_search_origin", self.cursor) self.cursor_x = getattr(self, "_search_origin_x", self.cursor_x) self.scroll_to(y=max(self.cursor - self._visible_height() // 2, 0), animate=False) @@ -626,26 +875,70 @@ class SearchMixin: def _compute_matches(self) -> None: term = self._term - needle = term.lower() if getattr(self, "_ci", True) else term + ci = getattr(self, "_ci", True) + needle = term.lower() if ci else term + count = self._search_line_count() + src = self._search_source_id() + # Typing forward can only ever REMOVE lines: a line holding "mov" holds + # "mo". So when the term just grew (and nothing else moved -- same + # case-folding, same body, same number of rows) rescan only the previous + # hits. Search is driven a keystroke at a time, and a segment of bash is + # 224k rows; this is the difference between rescanning all of them per + # keypress and looking at a few thousand. + # + # The row count is part of the key because the listing streams in behind + # the search: rows that arrived after the last pass have never been + # looked at, and narrowing would silently never find them. + n = len(term) matches: list[int] = [] - ranges: dict[int, list[tuple[int, int]]] = {} - for i in range(self._search_line_count()): - s = self._search_line_text(i) - if not s: - continue - hay = s.lower() if getattr(self, "_ci", True) else s - pos, rs = 0, [] - while True: - j = hay.find(needle, pos) - if j < 0: - break - rs.append((j, j + len(term))) - pos = j + len(term) - if rs: + ranges: object = {} + hay = self._search_haystack(count, src) + if hay is not None: + starts, blob, folded = hay + body = folded if ci else blob + nlines = len(starts) + blen = len(body) + line = 0 + j = body.find(needle) + while j >= 0: + # find() walks forward, so the line only ever advances. + while line + 1 < nlines and starts[line + 1] <= j: + line += 1 + matches.append(line) + # Only the LINE matters here; the offsets inside it are worked + # out on demand. So skip the rest of this line rather than + # finding every further occurrence in it. + nxt = starts[line + 1] if line + 1 < nlines else blen + j = body.find(needle, nxt) + ranges = _MatchRanges(set(matches), needle, n, ci, + self._search_line_text) + else: + # Typing forward can only ever REMOVE lines: a line holding "mov" + # holds "mo". So when the term just grew (and nothing else moved -- + # same case-folding, same body, same number of rows) rescan only the + # previous hits. + # + # The row count is part of the key because the listing streams in + # behind the search: rows that arrived after the last pass have never + # been looked at, and narrowing would silently never find them. + rows: object = range(count) + prev = self._matched_key + if (prev is not None and prev[2] == count and prev[3] == src + and prev[1] == ci and term.startswith(prev[0]) and prev[0]): + rows = self._matches + text_of = self._search_line_text + for i in rows: + s = text_of(i) + if not s: + continue + h = s.lower() if ci else s + if h.find(needle) < 0: + continue matches.append(i) - ranges[i] = rs + ranges = _MatchRanges(set(matches), needle, n, ci, text_of) self._matches = matches self._ranges = ranges + self._matched_key = (term, ci, count, src) def search_repeat(self, direction: int, include_current: bool = False) -> None: if not getattr(self, "_term", ""): @@ -679,6 +972,7 @@ class SearchMixin: self._term = "" self._matches = [] self._ranges = {} + self._reset_search_cache() self.refresh() def _match_style(self, idx: int) -> Style: @@ -693,316 +987,12 @@ class SearchMixin: # --------------------------------------------------------------------------- # # Virtualized disassembly view # --------------------------------------------------------------------------- # -class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True): - """A line-virtualized disassembly listing for a single function.""" - - BINDINGS = [ - Binding("j,down", "cursor_down", "Down", show=False), - Binding("k,up", "cursor_up", "Up", show=False), - Binding("ctrl+d", "half_page(1)", "½↓", show=False), - Binding("ctrl+u", "half_page(-1)", "½↑", show=False), - Binding("pagedown", "page(1)", "PgDn", show=False), - Binding("pageup", "page(-1)", "PgUp", show=False), - Binding("home", "goto_top", "Top", show=False), - Binding("G,end", "goto_bottom", "Bottom", show=False), - Binding("o", "toggle_opcodes", "Opcodes", show=False), - Binding("c", "define_code", "Code", show=False), - Binding("p", "define_func", "Func", show=False), - Binding("u", "undefine", "Undef", show=False), - Binding("tab,shift+tab", "app.toggle_view", "Pseudocode", priority=True), - Binding("f5", "app.toggle_view", "Decompile", priority=True), - Binding("L", "app.continuous_here", "Listing"), - *SearchMixin.SEARCH_BINDINGS, - *NavMixin.NAV_BINDINGS, - *ColumnCursor.COL_BINDINGS, - ] - - cursor = reactive(0, repaint=False) - cursor_x = reactive(0, repaint=False) - - class CursorMoved(Message): - """Posted when the disasm cursor moves; carries the instruction ea.""" - - def __init__(self, index: int, ea: int | None) -> None: - super().__init__() - self.index = index - self.ea = ea - - def __init__(self) -> None: - super().__init__() - self.model: DisasmModel | None = None - self.total = 0 - self._name = "" - self._pending_scroll_y: int | None = None - self._term = "" - self._matches: list[int] = [] - self._ranges: dict[int, list[tuple[int, int]]] = {} - self._search_texts: list[str] | None = None - self._show_ops = True - self._op_w = 0 # char width of the hex-bytes field (excl. trailing gap) - - def _op_field(self, line) -> str: # type: ignore[no-untyped-def] - """The padded opcode-bytes column (empty when hidden). Kept identical - between the rendered strip and the plain text so cursor/search offsets - line up.""" - if not self._show_ops or self._op_w <= 0: - return "" - raw = line.raw or b"" - return " ".join(f"{b:02X}" for b in raw).ljust(self._op_w) + " " - - def _line_plain(self, idx: int) -> str | None: - if self.model is None: - return None - line = self.model.cached_line(idx) - if line is None: - return None - s = f"{line.ea:08X} " + self._op_field(line) - if line.label: - s += f"{line.label}: " - return s + line.text - - # -- public API -------------------------------------------------------- # - def load(self, model: DisasmModel, name: str, cursor: int = 0, - cursor_x: int = 0, scroll_y: int | None = None) -> None: - self.model = model - self._name = name - self.total = 0 - self.cursor = cursor - self.cursor_x = cursor_x - self._pending_scroll_y = scroll_y - # NB: don't zero virtual_size here — that snaps the scroll to 0 and - # causes a visible jump before _on_primed restores the target scroll. - self._matches = [] - self._ranges = {} - self._search_texts = None - self._prime() - - # -- search hooks ------------------------------------------------------ # - def _fmt(self, line) -> str: # type: ignore[no-untyped-def] - s = f"{line.ea:08X} " + self._op_field(line) - if line.label: - s += f"{line.label}: " - return s + line.text - - def _search_line_count(self) -> int: - return self.total - - def _search_line_text(self, i: int) -> str | None: - t = self._search_texts - return t[i] if t is not None and 0 <= i < len(t) else None - - def _search_ensure(self, done) -> None: - if self._search_texts is not None: - done() - return - self._app_status(f"/{self._term}/ indexing {self.total} lines…") - self._index_for_search(done) - - @work(thread=True, exclusive=True, group="search-index") - def _index_for_search(self, done) -> None: - model = self.model - if model is None: - self.app.call_from_thread(done) - return - texts: list[str] = [] - off, total = 0, self.total - # A freshly-loaded blob is one row per undefined byte, so "all lines" can - # be millions of `db 4Ah` that nobody searches for. Index a bounded - # prefix rather than hang; say so instead of silently finding nothing. - LIMIT = 400_000 - capped = total > LIMIT - total = min(total, LIMIT) - while off < total: - lines = model.lines(off, min(DisasmModel.BLOCK, total - off), - prefetch=False) - if not lines: - break - texts.extend(self._fmt(ln) for ln in lines) - off += len(lines) - self._search_texts = texts - if capped: - self.app.call_from_thread( - self._app_status, - f"search covers the first {LIMIT:,} lines of {self.total:,}") - self.app.call_from_thread(done) - - @work(thread=True, exclusive=True, group="disasm-prime") - def _prime(self) -> None: - model = self.model - if model is None: - return - total = model.total() - height = max(self.size.height, 1) - model.lines(0, min(total, height + DisasmModel.BLOCK), prefetch=True) - if self.cursor: - model.lines(max(self.cursor - 2, 0), height, prefetch=True) - self.app.call_from_thread(self._on_primed, total) - - def _on_primed(self, total: int) -> None: - self.total = total - self.virtual_size = Size(0, total) - self._update_op_w() # provisional width from the primed window - self._scan_op_width() # settle it against the whole function - self._clamp_x() # cursor line is now cached; keep the column in range - if self._pending_scroll_y is not None and self._pending_scroll_y >= 0: - self._apply_scroll(min(self._pending_scroll_y, max(total - 1, 0))) - else: - self._scroll_cursor_into_view() - self._pending_scroll_y = None - self.refresh() - - # -- rendering --------------------------------------------------------- # - def render_line(self, y: int) -> Strip: - model = self.model - width = self.size.width - if model is None or self.total == 0: - return Strip([Segment("".ljust(width), _S_DIM)]) - top = round(self.scroll_offset.y) - if y == 0: # once per refresh: warm the visible window + a little ahead - self._ensure_window(top) - idx = top + y - if idx >= self.total: - return Strip([Segment("".ljust(width), _S_INSN)]) - line = model.cached_line(idx) - is_cursor = idx == self.cursor - if line is None: - strip = Strip([Segment(f" {idx:>8} …", _S_DIM)]) - else: - segs: list[Segment] = [Segment(f"{line.ea:08X} ", _S_ADDR)] - op = self._op_field(line) - if op: - segs.append(Segment(op, _S_OPBYTES)) - if line.label: - segs.append(Segment(f"{line.label}: ", _S_LABEL)) - mnem, _, rest = line.text.partition(" ") - segs.append(Segment(mnem, _S_MNEM)) - if rest: - segs.append(Segment(" " + rest, _S_INSN)) - strip = Strip(segs) - if idx in self._ranges: - strip = _overlay_ranges(strip, self._ranges[idx], self._match_style(idx)) - if is_cursor: - strip = _cursor_decorate(strip, self._line_plain(idx) or "", self.cursor_x) - return strip.adjust_cell_length(width, _S_INSN) - - def _update_op_w(self) -> bool: - """Recompute the opcode column width from the model's widest instruction. - Returns True if it changed.""" - w = 0 - if self.model is not None and self._show_ops: - mx = self.model.max_raw_len() - w = max(mx * 3 - 1, 0) if mx > 0 else 0 - if w != self._op_w: - self._op_w = w - return True - return False - - @work(thread=True, exclusive=True, group="disasm-opwidth") - def _scan_op_width(self) -> None: - model = self.model - if model is None: - return - model.scan_bytes() # fetch all blocks -> stable widest instruction - self.app.call_from_thread(self._settle_op_w) - - def _settle_op_w(self) -> None: - if self._update_op_w(): - self._search_texts = None # layout changed -> stale offsets - self.refresh() - - def action_toggle_opcodes(self) -> None: - self._show_ops = not self._show_ops - self._update_op_w() - self._search_texts = None # column layout changed -> reindex on next search - self._ranges = {} - self._clamp_x() - self.refresh() - self._app_status("opcodes " + ("on" if self._show_ops else "off")) - - def _ensure_window(self, top: int) -> None: - if self.model is None: - return - height = max(self.size.height, 1) - start = max(top - DisasmModel.BLOCK, 0) - count = height + 2 * DisasmModel.BLOCK - if not self.model.is_cached(top, height): - self._fetch_window(start, count) - else: - self.model.ensure_async(start, count) # warm neighbors - - @work(thread=True, exclusive=False, group="disasm-fetch") - def _fetch_window(self, start: int, count: int) -> None: - model = self.model - if model is None: - return - model.lines(start, count, prefetch=True) - self.app.call_from_thread(self.refresh) - - # -- navigation -------------------------------------------------------- # - def _visible_height(self) -> int: - return max(self.size.height, 1) - - def _scroll_cursor_into_view(self) -> None: - height = self._visible_height() - top = round(self.scroll_offset.y) - if self.cursor < top: - self.scroll_to(y=self.cursor, animate=False) - elif self.cursor >= top + height: - self.scroll_to(y=max(self.cursor - height + 1, 0), animate=False) - - def _move(self, delta: int) -> None: - if self.total == 0: - return - old = self.cursor - before = round(self.scroll_offset.y) - self.cursor = max(0, min(self.total - 1, self.cursor + delta)) - self._clamp_x() - self._scroll_cursor_into_view() - if round(self.scroll_offset.y) != before: - self.refresh() # scrolled: the whole viewport shifted - else: - _refresh_lines(self, old, self.cursor) # only the two changed rows - self._refresh_hl() - self.post_message(DisasmView.CursorMoved(self.cursor, self._cursor_ea())) - - def _after_cursor_move(self) -> None: - self.post_message(DisasmView.CursorMoved(self.cursor, self._cursor_ea())) - - def _cursor_ea(self) -> int | None: - if self.model is None: - return None - line = self.model.cached_line(self.cursor) - return line.ea if line else None - - # -- item structure edits (IDA c/p/u) --------------------------------- # - def action_define_code(self) -> None: - self.post_message(EditItemRequested(self, "code")) - - def action_define_func(self) -> None: - self.post_message(EditItemRequested(self, "func")) - - def action_undefine(self) -> None: - self.post_message(EditItemRequested(self, "undef")) - - def action_cursor_down(self) -> None: - self._move(1) - - def action_cursor_up(self) -> None: - self._move(-1) - - def action_goto_top(self) -> None: - self._move(-self.total) - - def action_goto_bottom(self) -> None: - self._move(self.total) - - # --------------------------------------------------------------------------- # # Virtualized flat listing view (code + data + undefined, per segment) # --------------------------------------------------------------------------- # class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True): """A line-virtualized *flat* listing over one segment: code, data and - undefined heads interleaved (IDA's disassembly view), unlike ``DisasmView`` + undefined heads interleaved (IDA's disassembly view), unlike ``DisasmModel`` which is bounded to one function. Backed by ``ListingModel`` (the ``heads`` server tool). Used for non-function regions and raw segment browsing. """ @@ -1019,12 +1009,18 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru Binding("end", "col_end", "eol", show=False), Binding("G,ctrl+end", "goto_bottom", "Bottom", show=False), Binding("ctrl+home", "goto_top", "Top", show=False), - Binding("o", "toggle_opcodes", "Opcodes"), + # `o` is IDA's operand-format key, and that muscle memory is worth more + # than the opcode column's old claim on it (moved to B, for bytes). + Binding("o", "op_format('cycle')", "Format"), + Binding("O", "op_format('back')", "Format \u2190", show=False), + Binding("B", "toggle_opcodes", "Bytes", show=False), Binding("c", "define_code", "Code", show=False), Binding("d", "make_data", "Data", show=False), Binding("a", "make_string", "Str", show=False), Binding("p", "define_func", "Func", show=False), Binding("u", "undefine", "Undef", show=False), + Binding("t", "toggle_thumb", "ARM/Thumb", show=False), + Binding("T", "thumb_scan", "Scan vectors", show=False), *SearchMixin.SEARCH_BINDINGS, *NavMixin.NAV_BINDINGS, *ColumnCursor.COL_BINDINGS, @@ -1052,6 +1048,9 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._name = "" self._pending_scroll_y: int | None = None self._pending_focus: str | None = None + #: Operand index to put the cursor column back on once rows land — a + #: reformat can change an operand's width, moving the ones after it. + self._pending_op: int | None = None self._term = "" self._matches: list[int] = [] self._ranges: dict[int, list[tuple[int, int]]] = {} @@ -1060,6 +1059,8 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._search_loading = False self._search_pending: list = [] # done-callbacks awaiting the load self._link_rows: set[int] = set() # split-view: linked instruction rows + #: {address: 'now'|'past'|'future'} painted under the code (trace mode). + self.trail: dict[int, str] = {} # -- text helpers ------------------------------------------------------ # def _head(self, idx: int) -> Head | None: @@ -1071,11 +1072,34 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru def _op_bytes_text(self, h: Head) -> str: """Hex bytes for ``h``, truncated with an ellipsis in 'limited' mode so a - long x86-64 instruction doesn't blow out the column.""" + long x86-64 instruction doesn't blow out the column. + + ``bytes.hex(" ")`` rather than a per-byte f-string generator: this is + called for every row of every plain line, and search builds the plain + line for the whole segment. + """ raw = h.raw or b"" if self._op_mode == 1 and len(raw) > _OP_LIMIT: - return " ".join(f"{b:02X}" for b in raw[:_OP_LIMIT]) + "\u2026" - return " ".join(f"{b:02X}" for b in raw) + return raw[:_OP_LIMIT].hex(" ").upper() + "\u2026" + return raw.hex(" ").upper() + + @staticmethod + def _span_segments(h: Head, fallback: Style): + """Segments for a row's disassembly text. + + Uses IDA's own token classification when Code Mode supplies it; falls + back to the mnemonic/rest split when spans are absent or disagree with + the plain text. + """ + if h.spans: + return [Segment(t, _S_SPAN.get(k, fallback)) for k, t in h.spans] + if h.kind == "code": + mnem, _, rest = h.text.partition(" ") + segs = [Segment(mnem, _S_MNEM)] + if rest: + segs.append(Segment(" " + rest, fallback)) + return segs + return [Segment(h.text, fallback)] def _op_field(self, h: Head) -> str: """The padded opcode-bytes column (empty when hidden). Shared format so @@ -1084,8 +1108,7 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru return "" return self._op_bytes_text(h).ljust(self._op_w) + " " - def _line_plain(self, idx: int) -> str | None: - h = self._head(idx) + def _plain_of(self, h: Head | None) -> str | None: if h is None: return None # Function headers and code labels sit at depth 0 (with the address); @@ -1098,6 +1121,22 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru extra = _LST_INDENT if h.kind == "member" else "" return base + self._op_field(h) + extra + self._name_prefix(h) + h.text + def _line_plain(self, idx: int) -> str | None: + return self._plain_of(self._head(idx)) + + def _search_line_texts(self, start: int, count: int) -> list: + """A window of plain lines in one model call. + + Building the search body row by row took the model's lock and bisected + its row table a quarter of a million times; ``window`` does both once + for the whole window. + """ + model = self.model + if model is None: + return [] + plain = self._plain_of + return [plain(h) for h in model.window(start, count)] + def _insn_col(self, idx: int) -> int: """Column where the instruction/content text begins, past the address + opcode-bytes gutter — the shift+home target. Mirrors ``_line_plain``'s @@ -1125,15 +1164,19 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru def load(self, model: ListingModel, name: str, cursor: int = 0, cursor_x: int = 0, scroll_y: int | None = None, focus: str | None = None) -> None: - self.model = model + previous, self.model = self.model, model self._name = name self.total = 0 self.cursor = cursor self.cursor_x = cursor_x self._pending_scroll_y = scroll_y self._pending_focus = focus # token to land the cursor column on + self._pending_op = None self._matches = [] self._ranges = {} + # Navigating inside the same segment reuses the same model, and the + # searchable body with it; only a different model invalidates it. + self._reset_search_cache(body=model is not previous) self._prime() @work(thread=True, exclusive=True, group="listing-prime") @@ -1166,6 +1209,16 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru if occ: self.cursor_x = occ[0][0] self._pending_focus = None + if self._pending_op is not None: + # Stay on the operand that was just reformatted: its text can change + # width, which moves every operand after it out from under the + # cursor (and the next press would then hit a different one). + h = self._head(self.cursor) + for lo, _hi, n in (h.ops if h is not None else None) or (): + if n == self._pending_op: + self.cursor_x = self._insn_col(self.cursor) + lo + break + self._pending_op = None self._clamp_x() if self._pending_scroll_y is not None and self._pending_scroll_y >= 0: self._apply_scroll(min(self._pending_scroll_y, max(total - 1, 0))) @@ -1198,6 +1251,8 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._op_mode = (self._op_mode + 1) % 3 self._update_op_w() self._ranges = {} # column layout changed -> stale match offsets + # ...and which rows match at all: the opcode hex is searchable text. + self._reset_search_cache(body=True) self._clamp_x() self.refresh() self._app_status("opcodes: " + {0: "off", 1: f"limited ({_OP_LIMIT} bytes)", @@ -1301,21 +1356,21 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru segs.append(Segment(_LST_INDENT, _S_MEMBER)) if h.name: segs.append(Segment(f"{h.name} ", _S_LABEL)) - if h.kind == "code": - mnem, _, rest = h.text.partition(" ") - segs.append(Segment(mnem, _S_MNEM)) - if rest: - segs.append(Segment(" " + rest, _S_INSN)) - elif h.kind == "data": - segs.append(Segment(h.text, _S_DATA)) - elif h.kind == "member": + if h.kind == "member": segs.append(Segment(h.text, _S_MEMBER)) else: - segs.append(Segment(h.text, _S_UNK)) + base = {"code": _S_INSN, "data": _S_DATA}.get(h.kind, _S_UNK) + segs.extend(self._span_segments(h, base)) strip = Strip(segs) linked = idx in self._link_rows if linked: strip = strip.apply_style(_S_LINK) # split-view companion band + if self.trail and h is not None: + kind = self.trail.get(h.ea) + if kind is not None: + strip = strip.apply_style( + _S_TRAIL_NOW if kind == "now" else + _S_TRAIL_PAST if kind == "past" else _S_TRAIL_FUTURE) plain = self._line_plain(idx) if (self._hl_word or idx == self.cursor) else None if idx in self._ranges: strip = _overlay_ranges(strip, self._ranges[idx], self._match_style(idx)) @@ -1325,8 +1380,30 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru strip = _overlay_ranges(strip, occ, _S_WORD) if idx == self.cursor: strip = _cursor_decorate(strip, plain or "", self.cursor_x) + # The operand 'o' would reformat, marked before you press it. A line + # can hold several literals and the cursor picks one; showing which + # is the difference between an edit you chose and one you got. + # Drawn LAST on purpose: _cursor_decorate paints the word under the + # cursor, and that word is usually the literal itself — so painting + # this first just loses to it. + span = self._cursor_operand(idx) + if span is not None: + strip = _overlay_over(strip, [span], _S_OPERAND) return strip.adjust_cell_length(width, _S_LINK if linked else _S_INSN) + def _cursor_operand(self, idx: int) -> tuple[int, int] | None: + """Screen columns of the operand under the cursor on row ``idx``, or None. + + The extents come from the worker (IDA's own operand markers) and are + offsets into the head's text, so they shift by the same gutter the + cursor column is measured against.""" + h = self._head(idx) + if h is None or not h.ops: + return None + base = self._insn_col(idx) + got = h.op_at(self.cursor_x - base) + return (base + got[0], base + got[1]) if got else None + # -- split-view link highlight ---------------------------------------- # def set_link(self, rows) -> None: rows = set(rows) if rows else set() @@ -1401,12 +1478,30 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru def action_undefine(self) -> None: self.post_message(EditItemRequested(self, "undef")) + def action_toggle_thumb(self) -> None: + self.post_message(EditItemRequested(self, "thumb")) + + def action_thumb_scan(self) -> None: + self.post_message(EditItemRequested(self, "thumbscan")) + def action_make_data(self) -> None: self.post_message(MakeDataRequested(self)) def action_make_string(self) -> None: self.post_message(EditItemRequested(self, "string")) + # -- literal display format (IDA 'o') --------------------------------- # + def action_op_format(self, mode: str = "cycle") -> None: + self.post_message(OpFormatRequested(self, mode)) + + def op_col(self) -> int: + """Cursor column inside the head's OWN text — the same string the worker + renders, so it can say which operand the cursor is standing on. -1 when + the cursor is left of it (in the address/opcode gutter), which means "you + didn't pick one, take the first literal".""" + base = self._insn_col(self.cursor) + return self.cursor_x - base if self.cursor_x >= base else -1 + def cur_head(self) -> Head | None: return self._head(self.cursor) @@ -1490,6 +1585,10 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True Binding("end", "col_end", "eol", show=False), Binding("ctrl+home", "goto_top", "Top", show=False), Binding("G,ctrl+end", "goto_bottom", "Bottom", show=False), + # Hex-Rays keeps number formats of its own, so `o` works here too — on + # the C literal under the cursor, not on the instruction's operand. + Binding("o", "op_format('cycle')", "Format"), + Binding("O", "op_format('back')", "Format \u2190", show=False), *SearchMixin.SEARCH_BINDINGS, *NavMixin.NAV_BINDINGS, *ColumnCursor.COL_BINDINGS, @@ -1518,13 +1617,44 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self._gutter = 0 # line-number gutter width (cells) self._line_eas: list[int | None] = [] # per-line address (marker stripped) self._link_line: int | None = None # split-view: linked pseudocode line + #: {line index: 'now'|'past'|'future'} — the execution trail, mapped from + #: instructions onto pseudocode via decomp_map. + self.trail: dict[int, str] = {} self._term = "" self._matches: list[int] = [] self._ranges: dict[int, list[tuple[int, int]]] = {} + #: {line: [(x0, x1, value, ea, opnum)]} — where the number literals are, + #: so the one under the cursor can be marked (Hex-Rays keeps formats per + #: literal, and a C line often has several). + self._nums: dict[int, list[tuple[int, int, str, int, int]]] = {} + #: (ea, opnum) to put the cursor back on once the nums land — a reformat + #: reflows the line, so the old column points at the wrong literal. + self._keep_lit: tuple[int, int] | None = None def _line_plain(self, idx: int) -> str | None: return self._texts[idx] if 0 <= idx < len(self._texts) else None + def set_nums(self, nums: dict) -> None: + self._nums = nums or {} + keep, self._keep_lit = self._keep_lit, None + if keep is not None: + # Land the cursor back on the literal that was just reformatted: + # `48` becoming `0x30` moves everything after it, so holding the + # column would put the next press on a different literal. + for x0, _x1, _v, ea, opnum in self._nums.get(self.cursor, ()): + if (ea, opnum) == keep: + self.cursor_x = x0 + self._hscroll() + break + self.refresh() + + def _cursor_literal(self, idx: int) -> tuple[int, int] | None: + """Columns of the number literal under the cursor on ``idx``, or None.""" + for x0, x1, _v, _ea, _op in self._nums.get(idx, ()): + if x0 <= self.cursor_x < x1: + return (x0, x1) + return None + def _col_offset(self) -> int: return self._gutter @@ -1558,6 +1688,10 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self.cursor_x = cursor_x self._matches = [] self._ranges = {} + # A whole new body: drop the joined haystack outright rather than trust + # id(self._texts) to differ, since the list it replaces is freed here and + # its address can be handed straight back. + self._reset_search_cache(body=True) # Gutter wide enough for the largest line number + a trailing space. self._gutter = (len(str(total)) + 1) if total else 0 maxw = max((s.cell_length for s in self._strips), default=0) @@ -1651,6 +1785,11 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True base = self._strips[idx] if linked: base = base.apply_style(_S_LINK) # split-view companion band + kind = self.trail.get(idx) if self.trail else None + if kind is not None: + base = base.apply_style( + _S_TRAIL_NOW if kind == "now" else + _S_TRAIL_PAST if kind == "past" else _S_TRAIL_FUTURE) if idx in self._ranges: base = _overlay_ranges(base, self._ranges[idx], self._match_style(idx)) if self._hl_word: @@ -1659,6 +1798,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True base = _overlay_ranges(base, occ, _S_WORD) if idx == self.cursor: base = _cursor_decorate(base, self._texts[idx], self.cursor_x) + span = self._cursor_literal(idx) # the literal `o` would reformat + if span is not None: # (last: see ListingView) + base = _overlay_over(base, [span], _S_OPERAND) code_w = max(width - gw, 0) code = base.crop(x, x + code_w).adjust_cell_length( code_w, _S_LINK if linked else None) @@ -1695,6 +1837,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True if round(old_value) != round(new_value): self.post_message(DecompView.Scrolled()) + def action_op_format(self, mode: str = "cycle") -> None: + self.post_message(OpFormatRequested(self, mode)) + def action_col_code_home(self) -> None: """shift+home: first non-blank column — past the C indentation, the pseudocode analogue of the listing's skip-the-address-gutter.""" @@ -1713,7 +1858,7 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True best, best_ea = i, e return best - # -- navigation (mirrors DisasmView) ---------------------------------- # + # -- navigation -------------------------------------------------------- # def _visible_height(self) -> int: return max(self.size.height, 1) @@ -1803,6 +1948,10 @@ class HexView(ScrollView, can_focus=True): super().__init__(id="hex") self.model = None self.total = 0 + #: Trace to read memory from, and the timestamp to read it at. When set, + #: the dump shows what memory HELD then rather than what the file holds. + self.trace = None + self.trace_idx = 0 self._internal_top: int | None = None # scroll target we set ourselves # -- public API -------------------------------------------------------- # @@ -1855,7 +2004,11 @@ class HexView(ScrollView, can_focus=True): self.refresh(layout=True) self.scroll_to(y=y, animate=False) - self.call_after_refresh(_fix) + # Only re-apply when the scroll clamped (the view's size wasn't computed + # yet). See ColumnCursor._apply_scroll: the deferred pass drags a full + # layout with it, which is far too expensive to do on every scroll. + if round(self.scroll_offset.y) != y: + self.call_after_refresh(_fix) def _scroll_to_cursor(self, center: bool = False) -> None: height = self._visible_height() @@ -1989,6 +2142,15 @@ class HexView(ScrollView, can_focus=True): return Strip([Segment("".ljust(width), _S_HEX)]) va, data = model.row(r) cur_row, cur_col = self.cursor // 16, self.cursor % 16 + # With a trace loaded the row shows what memory HELD at the current + # timestamp, not what the file contains. Only the bytes the trace + # actually saw are overlaid: the rest stay the database's, dimmed, so + # you can always tell evidence from the file's idea of the world. + tmem = tknown = None + if self.trace is not None and data is not None: + tmem, tknown = self.trace.memory(va, len(data), self.trace_idx) + if not any(tknown): + tmem = tknown = None fo = model.file_offset(va) fo_str = f"{fo:08X}" if fo is not None else "--------" segs: list[Segment] = [ @@ -1998,28 +2160,908 @@ class HexView(ScrollView, can_focus=True): if data is None: segs.append(Segment("… fetching", _S_DIM)) else: + # Emit RUNS, not one segment per byte. A row is 32 cells whose style + # almost never changes (one cursor cell, or a trace boundary), and a + # segment per cell made every hex frame 1540 segments for the + # compositor to cut and merge again. n = len(data) + run: list[str] = [] + run_st = None + + def flush(st=None, _segs=segs) -> None: + nonlocal run, run_st + if run: + _segs.append(Segment("".join(run), run_st)) + run = [] + run_st = st + + def put(text: str, st) -> None: + nonlocal run_st + if st is not run_st: + flush(st) + run.append(text) + + run_st = _S_HEX for i in range(16): if i == 8: - segs.append(Segment(" ", _S_HEX)) + put(" ", _S_HEX) if i < n: - st = _S_CELL if (r == cur_row and i == cur_col) else _S_HEX - segs.append(Segment(f"{data[i]:02X} ", st)) + live = tknown is not None and tknown[i] + val = tmem[i] if live else data[i] + if r == cur_row and i == cur_col: + st = _S_CELL + elif tknown is None: + st = _S_HEX + else: + st = _S_HEX_LIVE if live else _S_HEX_STALE + put(f"{val:02X} ", st) else: - segs.append(Segment(" ", _S_HEX)) - segs.append(Segment(" |", _S_DIM)) + put(" ", _S_HEX) + put(" |", _S_DIM) for i in range(16): if i < n: - ch = chr(data[i]) if 32 <= data[i] < 127 else "." - st = _S_CELL if (r == cur_row and i == cur_col) else _S_ASCII + live = tknown is not None and tknown[i] + val = tmem[i] if live else data[i] + ch = chr(val) if 32 <= val < 127 else "." + if r == cur_row and i == cur_col: + st = _S_CELL + elif tknown is None: + st = _S_ASCII + else: + st = _S_HEX_LIVE if live else _S_HEX_STALE else: ch, st = " ", _S_ASCII - segs.append(Segment(ch, st)) - segs.append(Segment("|", _S_DIM)) + put(ch, st) + put("|", _S_DIM) + flush() return Strip(segs).adjust_cell_length(width, _S_HEX) # --------------------------------------------------------------------------- # +# Graph view +# --------------------------------------------------------------------------- # +_S_GBORDER = Style(color="#4b5565") # box border, idle +_S_GBORDER_CUR = Style(color="#7aa2f7", bold=True) # box border, cursor block +_S_GLABEL = Style(color="#7aa2f7", bold=True) # loc_XXXX in the border +_S_GLABEL_CUR = Style(color="#c0caf5", bold=True) +_S_GDIM = Style(color="#5e6875") +_S_GENTRY = Style(color="#9ece6a", bold=True) # the entry block's label +#: Edge colours follow IDA's convention: green = branch taken, red = falls +#: through, blue = the block's only successor, purple = loops back. +_S_EDGE = { + graph.E_TRUE: Style(color="#5fbf5f"), + graph.E_FALSE: Style(color="#cf5f5f"), + graph.E_UNCOND: Style(color="#5f87d7"), + graph.E_SWITCH: Style(color="#c9a227"), + graph.E_BACK: Style(color="#a06fd0"), +} +#: Same hues, brightened: the edges touching the block you're on. +_S_EDGE_HOT = { + graph.E_TRUE: Style(color="#8ff08f", bold=True), + graph.E_FALSE: Style(color="#ff8f8f", bold=True), + graph.E_UNCOND: Style(color="#9fc4ff", bold=True), + graph.E_SWITCH: Style(color="#ffd75f", bold=True), + graph.E_BACK: Style(color="#d0a0ff", bold=True), +} +_S_MINI_BG = Style(bgcolor="#161b22", color="#3b4453") +_S_MINI_NODE = Style(bgcolor="#161b22", color="#5f7799") +_S_MINI_CUR = Style(bgcolor="#161b22", color="#9ece6a", bold=True) +_S_MINI_VIEW = Style(bgcolor="#233044", color="#c0caf5") +_S_MINI_EDGE = Style(bgcolor="#161b22", color="#2f3945") + +_GPAD = 1 # columns of padding inside a box +_MINI_W, _MINI_H = 30, 14 + + +#: Memo for ``Style + Style``. Rich rebuilds a whole Style on every ``+``, and +#: the graph merges an overlay (cursor band, trail, highlight) over a run of +#: cells that share only a handful of base styles -- so the same pair is +#: recombined hundreds of times per frame. +_STYLE_SUM: dict[tuple, Style] = {} + + +class _CellRow: + """A row of (char, style) cells that coalesces into a Strip. + + The graph is drawn per screen row from three independent sources -- edge + cells, boxes, then the minimap -- which overwrite each other. Composing into + a flat cell array and coalescing once at the end is both simpler and cheaper + than splicing Strips three times. + """ + + __slots__ = ("ch", "st", "width") + + def __init__(self, width: int, base: Style) -> None: + self.width = max(width, 0) + self.ch = [" "] * self.width + self.st = [base] * self.width + + def put(self, i: int, ch: str, style: Style) -> None: + if 0 <= i < self.width: + self.ch[i] = ch + self.st[i] = style + + def text(self, i: int, s: str, style: Style) -> None: + """Write ``s`` at cell ``i``, clipped to the row. + + Slice assignment rather than a call per character: a graph row is drawn + from box borders, instruction text and the minimap, and doing it a cell + at a time made painting one frame thousands of bound-method calls. + Assigning a str to a list slice expands it to characters in C. + """ + if not s: + return + a = i if i > 0 else 0 + b = i + len(s) + if b > self.width: + b = self.width + if b <= a: + return + self.ch[a:b] = s[a - i:b - i] + self.st[a:b] = [style] * (b - a) + + def restyle(self, a: int, b: int, style: Style) -> None: + """Merge ``style`` over the cells in [a, b) (keeps the characters).""" + if a < 0: + a = 0 + if b > self.width: + b = self.width + st = self.st + combine = _STYLE_SUM + for i in range(a, b): + base = st[i] + key = (base, style) + got = combine.get(key) + if got is None: + got = combine[key] = base + style + st[i] = got + + def strip(self) -> Strip: + segs: list[Segment] = [] + if not self.width: + return Strip([]) + run_start = 0 + cur = self.st[0] + for i in range(1, self.width): + if self.st[i] is not cur and self.st[i] != cur: + segs.append(Segment("".join(self.ch[run_start:i]), cur)) + run_start, cur = i, self.st[i] + segs.append(Segment("".join(self.ch[run_start:]), cur)) + return Strip(segs) + + +class GraphView(NavMixin, ScrollView, can_focus=True): + """IDA-style control-flow graph of one function, in character cells. + + The layout comes from ``idatui.graph`` (pure, offline-testable); this class + is only presentation, navigation and hit-testing. Nothing is pre-painted: a + big function is millions of cells, so each screen row is composed on demand + from the edge index plus whichever boxes cover that row -- the same + discipline as ``ListingView.render_line``. + + Box contents are the SAME ``Head`` rows the listing renders, so IDA's own + colour tags, names and operand text come along for free instead of this + growing a second disassembler renderer. + """ + + BINDINGS = [ + Binding("j,down", "cursor_down", "Down", show=False), + Binding("k,up", "cursor_up", "Up", show=False), + Binding("h,left", "cursor_left", "Left", show=False), + Binding("l,right", "cursor_right", "Right", show=False), + Binding("J", "succ_block", "Next block", show=False), + Binding("K", "pred_block", "Prev block", show=False), + Binding("w", "next_block", "Block →", show=False), + Binding("b", "prev_block", "Block ←", show=False), + Binding("0", "goto_entry", "Entry", show=False), + Binding("z", "zoom", "Zoom"), + Binding("m", "minimap", "Minimap", show=False), + Binding("f", "center", "Centre", show=False), + Binding("ctrl+d", "pan(12)", "½↓", show=False), + Binding("ctrl+u", "pan(-12)", "½↑", show=False), + Binding("pagedown", "pan(24)", "PgDn", show=False), + Binding("pageup", "pan(-24)", "PgUp", show=False), + Binding("home", "col_home", "bol", show=False), + Binding("end", "col_end", "eol", show=False), + *NavMixin.NAV_BINDINGS, + ] + + cursor_node = reactive(-1, repaint=False) + cursor_row = reactive(0, repaint=False) + cursor_x = reactive(0, repaint=False) + + ZOOMS = ("full", "compact", "collapsed") + + class CursorMoved(Message): + """Posted when the graph cursor lands on a new address.""" + + def __init__(self, ea: int | None, block: int) -> None: + super().__init__() + self.ea = ea + self.block = block + + def __init__(self) -> None: + super().__init__() + self.fc = None # domain.Flowchart + self.lay: graph.Layout | None = None + self.loaded_ea: int | None = None + self._blocks: dict[int, object] = {} + self._zoom = 0 + self._show_minimap = True + self._mini_cache: tuple | None = None + self._drag: tuple[int, int, float, float] | None = None + self._drag_map = False # the drag started on the minimap + self._hl_word = "" + self.trail: dict[int, str] | None = None + + # -- content ---------------------------------------------------------- # + def set_graph(self, fc, ea: int | None = None) -> None: + """Install a flowchart and lay it out at the current zoom.""" + self.fc = fc + self._blocks = {b.id: b for b in fc.blocks} if fc else {} + self.loaded_ea = fc.func_ea if fc else None + self._relayout() + if fc: + blk = fc.block_at(ea) if ea is not None else None + self.cursor_node = blk.id if blk else fc.entry + self.cursor_row = 0 + if blk is not None and ea is not None: + rows = self._rows(blk.id) + for i, h in enumerate(rows): + if h is not None and h.ea == ea: + self.cursor_row = i + break + self.cursor_x = 0 + self._center_cursor() + self.refresh(layout=True) + + def _relayout(self) -> None: + self._mini_cache = None + if not self.fc or not self.fc.blocks: + self.lay = None + self.virtual_size = Size(0, 0) + return + blocks = [graph.Block(id=b.id, start=b.start, end=b.end, + succs=list(b.succs)) for b in self.fc.blocks] + self.lay = graph.layout(blocks, self._sizer, entry=self.fc.entry) + self.virtual_size = Size(self.lay.width + 2, self.lay.height + 1) + + def _rows(self, nid: int): + """The Head rows shown inside block ``nid`` at the current zoom.""" + b = self._blocks.get(nid) + if b is None: + return [] + if self._zoom == 2: + return [None] # one synthetic summary row + return b.rows + + def _row_plain(self, nid: int, i: int) -> str: + b = self._blocks.get(nid) + if b is None: + return "" + if self._zoom == 2: + n = len(b.rows) + return f"{n} instruction{'s' if n != 1 else ''}" + rows = b.rows + if not (0 <= i < len(rows)): + return "" + h = rows[i] + if self._zoom == 0: + return f"{h.ea:08X} {self._head_text(h)}" + return self._head_text(h) + + @staticmethod + def _head_text(h) -> str: + return (f"{h.name} {h.text}" if h.name else h.text) + + def _sizer(self, b: graph.Block) -> tuple[int, int]: + nid = b.id + rows = self._rows(nid) + n = max(len(rows), 1) + label = f"loc_{b.start:X}" + widest = max([len(label) + 4] + + [len(self._row_plain(nid, i)) for i in range(n)]) + return (widest + 2 * _GPAD + 2, n + 2) + + # -- geometry --------------------------------------------------------- # + def _cur_node(self) -> graph.Node | None: + if self.lay is None: + return None + return self.lay.by_id.get(self.cursor_node) + + def cur_head(self): + rows = self._rows(self.cursor_node) + if self._zoom == 2 or not rows: + b = self._blocks.get(self.cursor_node) + return b.rows[0] if (b and b.rows) else None + i = max(0, min(self.cursor_row, len(rows) - 1)) + return rows[i] + + def _cursor_ea(self) -> int | None: + h = self.cur_head() + return h.ea if h is not None else None + + def _next_ea(self) -> int | None: + """The address after the cursor's instruction -- what ``follow`` uses to + skip the fall-through edge and land on a real branch target.""" + b = self._blocks.get(self.cursor_node) + if b is None: + return None + rows = b.rows + i = max(0, min(self.cursor_row, len(rows) - 1)) if rows else 0 + if rows and i + 1 < len(rows): + return rows[i + 1].ea + return b.end + + def _line_plain(self, _idx=None) -> str: + return self._row_plain(self.cursor_node, self.cursor_row) + + def word_under_cursor(self) -> str: + plain = self._line_plain() + if not plain: + return "" + a, b = _word_bounds(plain, self.cursor_x) + return plain[a:b] if b > a else "" + + def set_highlight(self, word: str) -> None: + if word != self._hl_word: + self._hl_word = word + self.refresh() + + def set_trail(self, trail: dict[int, str] | None) -> None: + self.trail = trail + self.refresh() + + # -- cursor motion ---------------------------------------------------- # + def _clamp_cursor(self) -> None: + if self.lay is None or not self.lay.nodes: + return + if self.cursor_node not in self.lay.by_id: + self.cursor_node = self.lay.nodes[0].id + rows = self._rows(self.cursor_node) + self.cursor_row = max(0, min(self.cursor_row, max(len(rows) - 1, 0))) + plain = self._line_plain() + self.cursor_x = max(0, min(self.cursor_x, max(len(plain) - 1, 0))) + + def _order(self) -> list[int]: + return [n.id for n in self.lay.nodes] if self.lay else [] + + def _moved(self) -> None: + self._clamp_cursor() + self._scroll_to_cursor() + self.refresh() + self.post_message(self.CursorMoved(self._cursor_ea(), self.cursor_node)) + + def action_cursor_down(self) -> None: + rows = self._rows(self.cursor_node) + if self.cursor_row + 1 < len(rows): + self.cursor_row += 1 + else: + order = self._order() + if self.cursor_node in order: + i = order.index(self.cursor_node) + if i + 1 < len(order): + self.cursor_node = order[i + 1] + self.cursor_row = 0 + self._moved() + + def action_cursor_up(self) -> None: + if self.cursor_row > 0: + self.cursor_row -= 1 + else: + order = self._order() + if self.cursor_node in order: + i = order.index(self.cursor_node) + if i > 0: + self.cursor_node = order[i - 1] + self.cursor_row = max(len(self._rows(self.cursor_node)) - 1, 0) + self._moved() + + def action_cursor_left(self) -> None: + self.cursor_x = max(0, self.cursor_x - 1) + self._moved() + + def action_cursor_right(self) -> None: + self.cursor_x = min(len(self._line_plain()), self.cursor_x + 1) + self._moved() + + def action_col_home(self) -> None: + self.cursor_x = 0 + self._moved() + + def action_col_end(self) -> None: + self.cursor_x = max(len(self._line_plain()) - 1, 0) + self._moved() + + def _hop(self, table: dict) -> None: + tgt = table.get(self.cursor_node) or [] + if not tgt: + self.app._status("no edge that way") + return + self.cursor_node = tgt[0][0] + self.cursor_row = 0 + self._moved() + + def action_succ_block(self) -> None: + if self.lay: + self._hop(self.lay.succ) + + def action_pred_block(self) -> None: + if self.lay: + self._hop(self.lay.pred) + + def _step_order(self, delta: int) -> None: + order = self._order() + if not order: + return + i = order.index(self.cursor_node) if self.cursor_node in order else 0 + self.cursor_node = order[max(0, min(len(order) - 1, i + delta))] + self.cursor_row = 0 + self._moved() + + def action_next_block(self) -> None: + self._step_order(1) + + def action_prev_block(self) -> None: + self._step_order(-1) + + def action_goto_entry(self) -> None: + if self.fc: + self.cursor_node = self.fc.entry + self.cursor_row = 0 + self._moved() + self._center_cursor() + + def action_pan(self, rows: int) -> None: + self.scroll_to(y=max(0, self.scroll_offset.y + rows), animate=False) + self._snap_into_view() + + def _viewport_has_block(self) -> bool: + if self.lay is None: + return False + y0 = int(self.scroll_offset.y) + x0 = int(self.scroll_offset.x) + y1, x1 = y0 + self.size.height, x0 + self.size.width + return any(n.y <= y1 and y0 <= n.bottom and n.x <= x1 and x0 <= n.right + for n in self.lay.nodes) + + def _snap_into_view(self) -> None: + """After a pan, if the viewport holds no block at all, ease to the + nearest one. + + Blocks cover a few percent of a laid-out graph -- 4.6% on an 87-block + function, under 1% on a 424-block one -- the rest being the padding that + keeps edges apart. Panning therefore lands in empty space more often + than not, and an empty screen gives you nothing to navigate back by. + Only fires when nothing is visible, so it never fights a deliberate pan. + """ + if self.lay is None or self._viewport_has_block(): + return + cy = self.scroll_offset.y + self.size.height / 2 + cx = self.scroll_offset.x + self.size.width / 2 + n = self._nearest_node(cy, cx) + if n is not None: + self._center_on(n, defer=False) + self.refresh() + + def action_zoom(self) -> None: + self._zoom = (self._zoom + 1) % len(self.ZOOMS) + self._relayout() + self._clamp_cursor() + self._center_cursor() + self.refresh(layout=True) + self.app._graph_status() # keeps the function name; names the zoom + + def action_minimap(self) -> None: + self._show_minimap = not self._show_minimap + self.refresh() + self.app._status(f"graph: minimap {'on' if self._show_minimap else 'off'}") + + def action_center(self) -> None: + self._center_cursor() + self.refresh() + + def action_copy_line(self) -> None: + plain = self._line_plain() + if not plain: + return + n = self.app._copy(plain) + self.app._status(f"copied line ({n} chars) to clipboard") + + def goto_ea(self, ea: int) -> bool: + """Put the cursor on ``ea`` if it lives in this graph.""" + if not self.fc: + return False + b = self.fc.block_at(ea) + if b is None: + return False + self.cursor_node = b.id + self.cursor_row = 0 + for i, h in enumerate(self._rows(b.id)): + if h is not None and h.ea == ea: + self.cursor_row = i + break + self.cursor_x = 0 + self._clamp_cursor() + self._center_cursor() + self.refresh() + return True + + # -- scrolling -------------------------------------------------------- # + def _cursor_cell(self) -> tuple[int, int] | None: + n = self._cur_node() + if n is None: + return None + row = n.y + 1 + max(0, min(self.cursor_row, n.h - 3)) + col = n.x + 1 + _GPAD + self.cursor_x + return (row, col) + + def _scroll_to_cursor(self) -> None: + cell = self._cursor_cell() + if cell is None: + return + row, col = cell + w, h = self.size.width, self.size.height + if w <= 0 or h <= 0: + return + y, x = self.scroll_offset.y, self.scroll_offset.x + if row < y: + y = row + elif row >= y + h: + y = row - h + 1 + if col < x + 2: + x = max(0, col - 2) + elif col >= x + w - 2: + x = col - w + 3 + if (y, x) != (self.scroll_offset.y, self.scroll_offset.x): + self.scroll_to(y=max(0, y), x=max(0, x), animate=False) + + def _center_on(self, n: graph.Node, defer: bool = True) -> None: + """Bring block ``n`` into the middle of the viewport. + + ``defer=False`` during a drag: layout is already valid then, and + queueing a callback per mouse-move makes the scrub lag behind. + """ + if n is None or self.size.width <= 0: + return + y = max(0, n.y - max(self.size.height // 2 - n.h // 2, 0)) + x = max(0, int(n.cx) - self.size.width // 2) + self.scroll_to(y=y, x=x, animate=False) + if not defer: + return + # Setting virtual_size then scrolling immediately clamps to 0 (max_scroll + # isn't recomputed until layout), so apply it again after the refresh. + def _again() -> None: + self.scroll_to(y=y, x=x, animate=False) + self.call_after_refresh(_again) + + def _center_cursor(self) -> None: + n = self._cur_node() + if n is not None: + self._center_on(n) + + # -- minimap hit-testing ----------------------------------------------- # + def _minimap_rect(self) -> tuple[int, int, int, int] | None: + """(left, top, w, h) of the minimap in CONTENT coordinates, or None. + + The minimap is pinned to the viewport, not the canvas, so these are + screen-relative and the scroll offset must NOT be added. Must agree with + _draw_minimap_row, which is why both take the inset from here. + """ + if not self._show_minimap or self.lay is None: + return None + w, h = self.size.width, self.size.height + if w < _MINI_W + 10 or h < _MINI_H + 2: + return None + return (w - _MINI_W - 2, 0, _MINI_W, _MINI_H) + + def _nearest_node(self, row: float, col: float) -> graph.Node | None: + """The block nearest a canvas point (distance 0 if the point is inside). + + Cells are about twice as tall as they are wide, so the column distance + is halved -- otherwise "nearest" means nearest in cells, which does not + look nearest on screen. + """ + if self.lay is None: + return None + best, best_d = None, None + for n in self.lay.nodes: + dx = 0.0 if n.x <= col <= n.right else min(abs(col - n.x), + abs(col - n.right)) + dy = 0.0 if n.y <= row <= n.bottom else min(abs(row - n.y), + abs(row - n.bottom)) + d = (dx * 0.5) ** 2 + dy ** 2 + if best_d is None or d < best_d: + best, best_d = n, d + return best + + def _minimap_seek(self, x: int, y: int, defer: bool = True) -> bool: + """Treat (x, y) as a point on the minimap and go to the block there. + + Deliberately snaps to the NEAREST BLOCK rather than scrolling to the raw + coordinate. Most of a laid-out graph is the padding that keeps edges + apart, so a coordinate-accurate jump usually parks the viewport in empty + space -- and the cursor, which only moved when the point landed exactly + on a block, stayed behind. Snapping means every click lands on something + and the keyboard carries on from there. + + Returns False if the point isn't on the minimap, so the caller can fall + through to ordinary canvas hit-testing. + """ + rect = self._minimap_rect() + if rect is None or self.lay is None: + return False + left, top, _w, _h = rect + gw, gh = _MINI_W - 2, _MINI_H - 2 + c, r = x - left - 1, y - top - 1 # inside the border + if not (0 <= c < gw and 0 <= r < gh): + return False + lay = self.lay + sx = max(lay.width / gw, 1e-9) + sy = max(lay.height / gh, 1e-9) + cx, cy = (c + 0.5) * sx, (r + 0.5) * sy # centre of that mini-cell + n = self._nearest_node(cy, cx) + if n is None: + self.scroll_to(x=max(0, int(cx - self.size.width / 2)), + y=max(0, int(cy - self.size.height / 2)), + animate=False) + return True + if n.id == self.cursor_node: + return True # already there; don't churn while dragging + self.cursor_node = n.id + self.cursor_row = 0 + self.cursor_x = 0 + self._clamp_cursor() + self._center_on(n, defer=defer) + self.refresh() + self.post_message(self.CursorMoved(self._cursor_ea(), self.cursor_node)) + return True + + # -- mouse ------------------------------------------------------------- # + def on_mouse_down(self, event) -> None: # type: ignore[no-untyped-def] + off = event.get_content_offset(self) + if off is None: + return + if self._minimap_seek(off.x, off.y): + self._drag = None + self._drag_map = True # keep scrubbing while the button is held + return + self._drag_map = False + self._drag = (off.x, off.y, self.scroll_offset.x, self.scroll_offset.y) + + def on_mouse_up(self, event) -> None: # type: ignore[no-untyped-def] + was_pan = self._drag is not None and not self._drag_map + self._drag = None + self._drag_map = False + if was_pan: + self._snap_into_view() # don't leave them adrift in the padding + + def on_mouse_move(self, event) -> None: # type: ignore[no-untyped-def] + if not event.button: + return + off = event.get_content_offset(self) + if off is None: + return + if self._drag_map: + # drag = scrub block to block through the overview + self._minimap_seek(off.x, off.y, defer=False) + return + if self._drag is None: + return + x0, y0, sx, sy = self._drag + self.scroll_to(x=max(0, sx + (x0 - off.x)), y=max(0, sy + (y0 - off.y)), + animate=False) + + def on_click(self, event) -> None: # type: ignore[no-untyped-def] + if self.lay is None: + return + off = event.get_content_offset(self) + if off is None: + return + # The minimap floats over the canvas, so it has to be tested FIRST -- + # otherwise a click on it is read as canvas coordinates and drops the + # cursor into whatever block happens to lie underneath. + if self._minimap_seek(off.x, off.y): + self.focus() + return + row = off.y + int(self.scroll_offset.y) + col = off.x + int(self.scroll_offset.x) + n = self.lay.node_at(row, col) + if n is None or n.block is None: + return + self.focus() + self.cursor_node = n.id + self.cursor_row = max(0, min(row - n.y - 1, + max(len(self._rows(n.id)) - 1, 0))) + self.cursor_x = max(0, col - n.x - 1 - _GPAD) + self._clamp_cursor() + self.refresh() + self.post_message(self.CursorMoved(self._cursor_ea(), self.cursor_node)) + if getattr(event, "chain", 1) >= 2: + self.post_message(FollowRequested(self)) + + # -- rendering -------------------------------------------------------- # + def _edge_styles(self) -> tuple[dict, set]: + hot = self.lay.incident.get(self.cursor_node, set()) if self.lay else set() + return (_S_EDGE, hot) + + def render_line(self, y: int) -> Strip: + width = self.size.width + if self.lay is None or not self.lay.nodes: + return Strip([Segment("".ljust(width), _S_GDIM)]) + row = int(self.scroll_offset.y) + y + col0 = int(self.scroll_offset.x) + out = _CellRow(width, _S_INSN) + base, hot = self._edge_styles() + + # 1. edge cells (an index query, never a painted canvas) + for col, (ch, kind, eid) in self.lay.painting.cells_at_row( + row, col0, col0 + width).items(): + st = (_S_EDGE_HOT if eid in hot else base).get(kind, _S_GDIM) + out.put(col - col0, ch, st) + + # 2. boxes covering this row (they win over edges: nothing routes inside) + for n in self.lay.nodes_at_row(row): + self._draw_node_row(out, n, row, col0) + + # 3. minimap, last, over everything + if self._show_minimap: + self._draw_minimap_row(out, y, width) + return out.strip().adjust_cell_length(width, _S_INSN) + + def _draw_node_row(self, out: _CellRow, n: graph.Node, row: int, + col0: int) -> None: + cur = n.id == self.cursor_node + bs = _S_GBORDER_CUR if cur else _S_GBORDER + left = n.x - col0 + w = n.w + b = self._blocks.get(n.id) + if row == n.y: + # top border carries the label: ┌─ loc_1234 ─────┐ + label = f"loc_{n.block.start:X}" if n.block else "" + if b is not None and b.rows and b.rows[0].name: + label = b.rows[0].name + out.text(left, graph.BOX["tl"] + graph.BOX["h"] * (w - 2) + + graph.BOX["tr"], bs) + tag = f" {label} " + if len(tag) <= w - 4: + st = _S_GENTRY if (self.fc and n.id == self.fc.entry) else ( + _S_GLABEL_CUR if cur else _S_GLABEL) + out.text(left + 2, tag, st) + if n.block is not None and n.block.selfloop: + out.put(left + w - 2, "↺", _S_EDGE[graph.E_BACK]) + return + if row == n.y + n.h - 1: + out.text(left, graph.BOX["bl"] + graph.BOX["h"] * (w - 2) + + graph.BOX["br"], bs) + return + out.put(left, graph.BOX["v"], bs) + out.put(left + w - 1, graph.BOX["v"], bs) + out.text(left + 1, " " * (w - 2), _S_INSN) + i = row - n.y - 1 + rows = self._rows(n.id) + if not (0 <= i < len(rows)): + return + text_col = left + 1 + _GPAD + h = rows[i] + plain = self._row_plain(n.id, i) + if h is None: # collapsed summary + out.text(text_col, plain, _S_GDIM) + else: + c = text_col + if self._zoom == 0: + out.text(c, f"{h.ea:08X} ", _S_ADDR) + c += 10 + if h.name: + out.text(c, f"{h.name} ", _S_LABEL) + c += len(h.name) + 2 + fallback = {"code": _S_INSN, "data": _S_DATA}.get(h.kind, _S_UNK) + if h.spans: + for kind, t in h.spans: + out.text(c, t, _S_SPAN.get(kind, fallback)) + c += len(t) + else: + out.text(c, h.text, fallback) + inner_a, inner_b = left + 1, left + w - 1 + # execution trail (a loaded Tenet trace), same palette as the listing + if self.trail is not None and h is not None: + k = self.trail.get(h.ea) + if k is not None: + out.restyle(inner_a, inner_b, + _S_TRAIL_NOW if k == "now" else + _S_TRAIL_PAST if k == "past" else _S_TRAIL_FUTURE) + if self._hl_word and plain: + for a, bb in _word_occurrences(plain, self._hl_word): + out.restyle(text_col + a, text_col + bb, _S_WORD) + if cur and i == self.cursor_row: + out.restyle(inner_a, inner_b, _S_CURSOR) + x = max(0, min(self.cursor_x, max(len(plain) - 1, 0))) + wa, wb = _word_bounds(plain, x) + if wb > wa: + out.restyle(text_col + wa, text_col + wb, _S_WORD) + if self.has_focus: + out.restyle(text_col + x, text_col + x + 1, _S_CELL) + + # -- minimap ---------------------------------------------------------- # + def _minimap(self) -> list[list[int]]: + """A coarse occupancy grid of the whole graph: 0 empty, 1 edge, 2 block, + 3 the cursor's block. Cached per (layout, zoom, cursor block).""" + key = (id(self.lay), self._zoom, self.cursor_node) + if self._mini_cache and self._mini_cache[0] == key: + return self._mini_cache[1] + gw, gh = _MINI_W - 2, _MINI_H - 2 + grid = [[0] * gw for _ in range(gh)] + lay = self.lay + if lay is not None and lay.width and lay.height: + sx = max(lay.width / gw, 1e-9) + sy = max(lay.height / gh, 1e-9) + for lo, hi, col, _kind, _eid in lay.painting.vruns: + c = min(int(col / sx), gw - 1) + for r in range(min(int(lo / sy), gh - 1), + min(int(hi / sy), gh - 1) + 1): + if not grid[r][c]: + grid[r][c] = 1 + for n in lay.nodes: + mark = 3 if n.id == self.cursor_node else 2 + r0, r1 = int(n.y / sy), int((n.y + n.h - 1) / sy) + c0, c1 = int(n.x / sx), int(n.right / sx) + for r in range(max(r0, 0), min(r1, gh - 1) + 1): + for c in range(max(c0, 0), min(c1, gw - 1) + 1): + if grid[r][c] < mark: + grid[r][c] = mark + self._mini_cache = (key, grid) + return grid + + def _draw_minimap_row(self, out: _CellRow, y: int, width: int) -> None: + # Inset by two: a ScrollView paints its vertical scrollbar over the last + # column, which otherwise eats the minimap's right border. + if self._minimap_rect() is None or not (0 <= y < _MINI_H): + return + left = self._minimap_rect()[0] # one source of truth with the hit-test + grid = self._minimap() + gw, gh = _MINI_W - 2, _MINI_H - 2 + lay = self.lay + sx = max(lay.width / gw, 1e-9) + sy = max(lay.height / gh, 1e-9) + vy0 = int(self.scroll_offset.y / sy) + vy1 = int((self.scroll_offset.y + self.size.height) / sy) + vx0 = int(self.scroll_offset.x / sx) + vx1 = int((self.scroll_offset.x + width) / sx) + + if y == 0: + out.put(left, graph.BOX["tl"], _S_MINI_BG) + for i in range(1, _MINI_W - 1): + out.put(left + i, graph.BOX["h"], _S_MINI_BG) + out.put(left + _MINI_W - 1, graph.BOX["tr"], _S_MINI_BG) + tag = f" {len(lay.nodes)} blocks " + out.text(left + 2, tag, _S_MINI_BG) + return + if y == _MINI_H - 1: + out.put(left, graph.BOX["bl"], _S_MINI_BG) + for i in range(1, _MINI_W - 1): + out.put(left + i, graph.BOX["h"], _S_MINI_BG) + out.put(left + _MINI_W - 1, graph.BOX["br"], _S_MINI_BG) + return + r = y - 1 + out.put(left, graph.BOX["v"], _S_MINI_BG) + out.put(left + _MINI_W - 1, graph.BOX["v"], _S_MINI_BG) + for c in range(gw): + v = grid[r][c] if r < len(grid) else 0 + inview = vy0 <= r <= vy1 and vx0 <= c <= vx1 + if v == 3: + ch, st = "█", _S_MINI_CUR + elif v == 2: + ch, st = "█", _S_MINI_NODE + elif v == 1: + ch, st = "·", _S_MINI_EDGE + else: + ch, st = " ", _S_MINI_BG + if inview and v != 3: + st = _S_MINI_VIEW if v == 0 else st + Style(bgcolor="#233044") + out.put(left + 1 + c, ch, st) + + +# --------------------------------------------------------------------------- # # Function list panel # --------------------------------------------------------------------------- # class FunctionsPanel(Vertical): @@ -2052,8 +3094,8 @@ class XrefsScreen(ModalScreen): self._preselect = preselect def compose(self) -> ComposeResult: - with Vertical(id="xref-box"): - yield Static(self._label, id="xref-title") + with Vertical(id="xref-box") as box: + box.border_title = Text(self._label.strip()) yield OptionList(*[Option(text) for _, text in self._items], id="xref-list") def on_mount(self) -> None: @@ -2133,8 +3175,8 @@ class SymbolPalette(ModalScreen): self._results: list[tuple] = [] def compose(self) -> ComposeResult: - with Vertical(id="pal-box"): - yield Static(" symbols", id="pal-title", markup=False) + with Vertical(id="pal-box") as box: + box.border_title = Text("symbols") yield Input(placeholder="fuzzy find symbol… ↑↓ select · Enter open · Esc close", id="pal-input") yield OptionList(id="pal-list") @@ -2213,8 +3255,8 @@ class SymbolPalette(ModalScreen): more = "+" if len(self._results) == cap else "" hint = " (F2: this binary)" if self._project_scope else ( " (F2: whole project)" if self._index is not None else "") - self.query_one("#pal-title", Static).update( - f" symbols [{scope}]: {len(self._results)}{more}{hint}") + self.query_one("#pal-box").border_title = Text( + f"symbols [{scope}]: {len(self._results)}{more}{hint}") def action_cursor_down(self) -> None: ol = self.query_one(OptionList) @@ -2251,6 +3293,182 @@ def _str_display(text: str, limit: int = 200) -> str: return out[:limit] + ("\u2026" if len(out) > limit else "") +class SearchPalette(ModalScreen): + """Ctrl+F: search the whole database, by text or by bytes. + + Unlike every other palette here this does NOT filter as you type: each + search walks the image or the listing in the database process, so it runs + on Enter. That gives Enter two jobs, which is fine as long as it is never + ambiguous: while the query differs from what was last searched, Enter + searches; once the results on screen belong to the query in the box, Enter + opens the highlighted one. The title says which it will do. + + Mode is guessed from the query (idatui/search.py) because asking first is a + tax on every search; F2 overrides the guess and a `hex:`/`text:` prefix + settles it outright. + """ + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + Binding("f2", "mode", "Text / bytes", show=False), + ] + LIMIT = 500 + + def __init__(self, program, initial: str = "") -> None: + super().__init__() + self._program = program + self._initial = initial + self._forced: str | None = None # F2: pin the mode + self._hits: list = [] + self._searched: tuple[str, str] | None = None # (mode, query) on screen + self._busy = False + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box") as box: + box.border_title = Text("search") + yield Input(placeholder="text, or bytes like 48 8b ?? c3 \u00b7 " + "Enter search \u00b7 F2 mode \u00b7 Esc close", + id="pal-input") + yield OptionList(id="pal-list") + + def on_mount(self) -> None: + inp = self.query_one("#pal-input", Input) + inp.value = self._initial + inp.focus() + self._retitle() + if self._initial: + self._run() + + # -- mode / title ------------------------------------------------------- # + def _query(self) -> str: + return self.query_one("#pal-input", Input).value.strip() + + def _mode_query(self) -> tuple[str, str]: + return search.classify(self._query(), self._forced) + + def _retitle(self, note: str = "") -> None: + mode, q = self._mode_query() + pinned = "" if self._forced is None else "*" + state = note + if not state: + if self._busy: + state = "searching\u2026" + elif self._searched == (mode, q) and q: + n = len(self._hits) + state = f"{n} hit{'' if n == 1 else 's'} \u2014 Enter opens" + elif q: + state = "Enter searches" + self.query_one("#pal-box").border_title = Text( + f"search [{mode}{pinned}]" + (f": {state}" if state else "")) + + def action_mode(self) -> None: + mode, _ = self._mode_query() + self._forced = search.TEXT if mode == search.BYTES else search.BYTES + self._searched = None # the results on screen are for the old mode + self._retitle() + + def on_input_changed(self, event: Input.Changed) -> None: + event.stop() # modal inputs bubble to the app's own #search handler + self._retitle() + + def on_input_submitted(self, event: Input.Submitted) -> None: + event.stop() + mode, q = self._mode_query() + if self._searched == (mode, q) and self._hits: + self.action_choose() + else: + self._run() + + # -- searching ---------------------------------------------------------- # + def _run(self) -> None: + mode, q = self._mode_query() + if not q: + return + if mode == search.BYTES: + problem = search.pattern_problem(q) + if problem: + # Refuse here rather than round-tripping: IDA's own message for + # a bad pattern is empty about half the time. + self._hits, self._searched = [], None + self.query_one(OptionList).clear_options() + self._retitle(problem) + return + q = search.normalise_pattern(q) + self._busy = True + self._retitle() + self._search(mode, q) + + @work(thread=True, exclusive=True, group="dbsearch") + def _search(self, mode: str, query: str) -> None: + try: + hits, err, truncated = self._program.search(query, mode, + limit=self.LIMIT) + except Exception as e: # noqa: BLE001 -- a search must not kill the app + hits, err, truncated = [], str(e), False + self.app.call_from_thread(self._present, mode, query, hits, err, truncated) + + def _present(self, mode: str, query: str, hits: list, err: str | None, + truncated: bool) -> None: + self._busy = False + self._hits = hits + # Remember what these results ARE, not what the box says now: the user + # may have typed on while the search ran, and then Enter must search + # again rather than open a hit from the previous query. + self._searched = (mode, query) if err is None else None + ol = self.query_one(OptionList) + ol.clear_options() + opts = [] + for h in hits: + label = Text() + label.append(f"{h.addr:08X} ", _S_ADDR) + label.append(f"{(h.func or h.seg or ''):<22.22} ", _S_DIM) + body = Text(h.line or "") + if mode == search.TEXT and query: + low, ql = (h.line or "").lower(), query.lower() + at = low.find(ql) + if at >= 0: + body.stylize(_S_NAME_MATCH, at, at + len(query)) + label.append_text(body) + opts.append(Option(label)) + ol.add_options(opts) + if hits: + ol.highlighted = 0 + if err: + self._retitle(err) + elif not hits: + self._retitle("no match") + else: + n = len(hits) + self._retitle(f"{n}{'+' if truncated else ''} " + f"hit{'' if n == 1 else 's'} \u2014 Enter opens") + + # -- moving / choosing --------------------------------------------------- # + def action_cursor_down(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = min((ol.highlighted or 0) + 1, ol.option_count - 1) + + def action_cursor_up(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = max((ol.highlighted or 0) - 1, 0) + + def action_choose(self) -> None: + ol = self.query_one(OptionList) + i = ol.highlighted + if i is not None and 0 <= i < len(self._hits): + self.dismiss(self._hits[i]) + + def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: + if 0 <= event.option_index < len(self._hits): + self.dismiss(self._hits[event.option_index]) + + def action_close(self) -> None: + self.dismiss(None) + + class StringsPalette(ModalScreen): """Every string in the binary (IDA's Shift+F12), filterable; Enter jumps to it in the unified listing.""" @@ -2279,8 +3497,8 @@ class StringsPalette(ModalScreen): self._results: list[tuple] = [] def compose(self) -> ComposeResult: - with Vertical(id="pal-box"): - yield Static(" strings", id="pal-title", markup=False) + with Vertical(id="pal-box") as box: + box.border_title = Text("strings") yield Input(placeholder="filter strings\u2026 \u2191\u2193 select \u00b7 " "Enter jump \u00b7 Esc close", id="pal-input") yield OptionList(id="pal-list") @@ -2353,8 +3571,8 @@ class StringsPalette(ModalScreen): more = "+" if len(rows) == cap else "" hint = " (F2: this binary)" if self._project_scope else ( " (F2: whole project)" if self._index is not None else "") - self.query_one("#pal-title", Static).update( - f" strings [{scope}]: {len(self._results)}{more} of {len(self._rows)}{hint}") + self.query_one("#pal-box").border_title = Text( + f"strings [{scope}]: {len(self._results)}{more} of {len(self._rows)}{hint}") def action_cursor_down(self) -> None: ol = self.query_one(OptionList) @@ -2397,12 +3615,15 @@ _HELP = ( )), ("Views", ( ("Tab / F5", "disassembly \u21c4 pseudocode"), + ("Space", "control-flow graph \u21c4 text"), ("s", "split view: listing + pseudocode"), ("Tab", "in split: switch the driving pane"), ("\\", "hex view"), - ("o", "cycle the opcode-bytes column"), + ("B", "cycle the opcode-bytes column"), ("Ctrl+B", "show/hide the names pane"), ("Ctrl+T", "structs / types editor"), + ("Ctrl+F", "search the database: text or bytes"), + ("Ctrl+E", "export findings as markdown"), ("Ctrl+P", "command palette"), )), ("Move", ( @@ -2424,6 +3645,7 @@ _HELP = ( ("d", "make data"), ("a", "make string"), ("u", "undefine"), + ("o / O", "literal format: hex/dec/bin/char/offset"), ("Ctrl+S", "save the database"), )), ("Search", ( @@ -2431,19 +3653,36 @@ _HELP = ( ("?", "search backward"), ("N", "previous match"), ("Ctrl+Y", "copy the current line"), - ("F1", "this cheatsheet"), + ("F1 / H", "this cheatsheet"), ("q", "quit"), )), + ("Graph (Space)", ( + ("j / k", "line up/down, crossing blocks"), + ("h / l", "column left / right"), + ("J / K", "follow an edge to a successor / predecessor"), + ("w / b", "next / previous block in layout order"), + ("0", "jump to the entry block"), + ("z", "zoom: full \u2192 compact \u2192 collapsed"), + ("m", "show/hide the minimap"), + ("f", "centre on the current block"), + ("Enter", "follow (stays in the graph if it lands here)"), + ("drag / click", "pan / put the cursor in a block"), + ("click minimap", "jump the view there (drag to scrub)"), + )), ) class QuitScreen(ModalScreen): - """Asked before exiting with unsaved database changes. Dismisses with - "save", "discard" or None (stay).""" + """Asked before exiting with unsaved database changes. + + Code Mode clients cannot roll a shared database back. The ``d`` choice means + "do not explicitly save": a GUI keeps the changes dirty, while a managed + idalib worker may persist them when its final lease closes. + """ BINDINGS = [ Binding("s", "save", "Save & quit"), - Binding("d", "discard", "Discard & quit"), + Binding("d", "discard", "Leave & quit"), Binding("escape,c", "cancel", "Cancel"), ] @@ -2454,13 +3693,13 @@ class QuitScreen(ModalScreen): def compose(self) -> ComposeResult: what = (f"{len(self._labels)} databases have unsaved changes" if len(self._labels) > 1 else "unsaved changes") - with Vertical(id="quit-box"): - yield Static(f"\u26a0 {what}", id="quit-title") + with Vertical(id="quit-box") as box: + box.border_title = Text(f"\u26a0 {what}") body = Text() for label in self._labels: body.append(f" \u2022 {label}\n", _S_LABEL) yield Static(body, id="quit-list") - yield Static("s save & quit d discard & quit Esc cancel", + yield Static("s save & quit d leave as-is & quit Esc cancel", id="quit-help") def action_save(self) -> None: @@ -2476,7 +3715,7 @@ class QuitScreen(ModalScreen): class HelpScreen(ModalScreen): """F1: the keyboard cheatsheet, replacing the permanent footer.""" - BINDINGS = [Binding("escape,f1,q,question_mark", "close", "Close")] + BINDINGS = [Binding("escape,f1,H,q,question_mark", "close", "Close")] #: widest cell content, +2 for the card's border, +2 for its padding _CARD_PAD = 4 @@ -2487,8 +3726,8 @@ class HelpScreen(ModalScreen): avail = max(self.app.size.width - 6, 20) cols = self._columns(avail) per = -(-len(_HELP) // cols) # ceil, so the columns stay balanced - with Vertical(id="help-box"): - yield Static(" keys", id="help-title", markup=False) + with Vertical(id="help-box") as box: + box.border_title = Text("keys") # Still inside a scroll container, so a genuinely tiny terminal # degrades to scrolling rather than clipping — but spread across the # width it shouldn't come to that. @@ -2504,7 +3743,7 @@ class HelpScreen(ModalScreen): classes="help-card", markup=False) card.border_title = title yield card - yield Static("Esc / F1 to close", id="help-foot") + yield Static("Esc · F1 · H to close", id="help-foot") @staticmethod def _section_widths() -> list[int]: @@ -2551,6 +3790,221 @@ class HelpScreen(ModalScreen): self.dismiss(None) +class RegWriteScreen(ModalScreen): + """Registers, and the instruction that set each one. + + "Which instruction set this register to its current value?" is the question + a trace exists to answer, and seeking backwards to it is a single keypress + here rather than a manual walk. Forward is offered too, but backward is what + people actually want — you notice a bad value after it has been used. + """ + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + Binding("enter", "choose", show=False, priority=True), + Binding("f", "choose_forward", show=False), + ] + + def __init__(self, rows, idx: int) -> None: + super().__init__() + self._rows = rows # (name, value, last_write, next_write) + self._idx = idx + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box") as box: + box.border_title = Text(f"registers at t={self._idx:,}") + box.border_subtitle = Text("Enter seeks to the write \u00b7 f seeks forward") + yield OptionList(id="pal-list") + + def on_mount(self) -> None: + ol = self.query_one(OptionList) + opts = [] + for name, val, last, nxt in self._rows: + label = Text() + label.append(f" {name:>4} ", _S_MNEM) + label.append(f"{val:#018x} " if val > 0xFFFFFFFF else f"{val:#010x} ", + _S_INSN) + if last is None: + label.append("never written in this trace", _S_DIM) + elif last == self._idx: + label.append("set by THIS instruction", _S_DATA) + else: + label.append(f"set at t={last:,}", _S_LABEL) + label.append(f" ({self._idx - last:,} steps back)", _S_DIM) + if nxt is not None: + label.append(f" next t={nxt:,}", _S_DIM) + opts.append(Option(label)) + ol.add_options(opts) + ol.highlighted = 0 + ol.focus() + + def action_cursor_down(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = min((ol.highlighted or 0) + 1, ol.option_count - 1) + + def action_cursor_up(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = max((ol.highlighted or 0) - 1, 0) + + def _pick(self, forward: bool) -> None: + i = self.query_one(OptionList).highlighted + if i is None or not (0 <= i < len(self._rows)): + self.dismiss(None) + return + _name, _val, last, nxt = self._rows[i] + self.dismiss(nxt if forward else last) + + def action_choose(self) -> None: + self._pick(False) + + def action_choose_forward(self) -> None: + self._pick(True) + + def on_option_list_option_selected(self, event) -> None: # type: ignore[no-untyped-def] + self._pick(False) + + def action_close(self) -> None: + self.dismiss(None) + + +class TraceDock(Vertical): + """Registers and a timeline for the loaded execution trace, docked right. + + Persistent rather than a modal: a trace turns every other view into "state + at time T", so the time and the registers are context you read WHILE looking + at code, not something you open and dismiss. + """ + + def __init__(self) -> None: + super().__init__(id="trace-dock") + self.trace = None + self.idx = 0 + + def compose(self) -> ComposeResult: + yield Static("", id="trace-head", markup=False) + yield Static("", id="trace-regs", markup=False) + yield Static("", id="trace-stack", markup=False) + yield TraceTimeline(id="trace-timeline") + + def show(self, trace, idx: int) -> None: + self.trace = trace + self.idx = idx + tl = self.query_one(TraceTimeline) + tl.trace, tl.idx = trace, idx + self.refresh_state() + + def refresh_state(self) -> None: + t = self.trace + if t is None: + return + n = max(t.length, 1) + pct = (self.idx + 1) * 100.0 / n + head = Text() + head.append(f" {self.idx:,}", _S_MNEM) + head.append(f" / {t.length - 1:,} ", _S_DIM) + head.append(f"{pct:5.1f}%\n", _S_ADDR) + # Register values are machine state and stay as the trace recorded them, + # but everything else on screen is in database addresses. Showing both + # here explains the relationship once, where it's read, instead of + # leaving "pc 0x2aed" next to "rip 0x7ffff6faaaed" to be puzzled over. + head.append(f" pc {t.ip(self.idx):#x}", _S_LABEL) + if t.slide: + head.append(f" (trace {t.raw_ip(self.idx):#x})", _S_DIM) + self.query_one("#trace-head", Static).update(head) + + # Registers, with the ones THIS instruction wrote called out: that + # difference is the entire reason a delta trace is readable. + changed = t.changed(self.idx) + body = Text() + pc = t.pc_name + for name in t.registers: + v = t.register(name, self.idx) + if v is None: + continue + hot = name in changed + body.append(f" {name:>4} ", _S_MNEM if hot else _S_DIM) + body.append(f"{v:#018x}\n" if v > 0xFFFFFFFF else f"{v:#010x}\n", + _S_DATA if hot else (_S_LABEL if name == pc else _S_INSN)) + self.query_one("#trace-regs", Static).update(body) + self._render_stack(t) + tl = self.query_one(TraceTimeline) + tl.idx = self.idx + tl.refresh() + + + STACK_WORDS = 8 + + def _render_stack(self, t) -> None: # type: ignore[no-untyped-def] + """The stack as of this instant, read out of the trace. + + This is where a trace's memory actually is: on two real traces, NONE of + the accesses fell inside the image — every one was stack or heap. A + memory view that could only address the image would have nothing to show. + + Bytes the trace never saw are printed as '??' rather than zeros. A trace + knows what it observed and nothing else, and quietly rendering unseen + memory as zero would invent facts. + """ + sp_name = next((r for r in ("rsp", "esp", "sp") if r in t.reg_at), "") + sp = t.register(sp_name, self.idx) if sp_name else None + out = Text() + if sp is None: + self.query_one("#trace-stack", Static).update(out) + return + width = 8 if (t.info and "64" in (t.info.arch or "")) else 4 + out.append(f" stack ({sp_name})\n", _S_DIM) + for k in range(self.STACK_WORDS): + a = sp + k * width + data, known = t.memory_raw(a, width, self.idx) + out.append(" \u25b8" if k == 0 else " ", _S_MNEM) + out.append(f"{a:012x} ", _S_ADDR) + if all(known): + v = int.from_bytes(data, "little") + out.append(f"{v:0{width * 2}x}\n", _S_DATA if k == 0 else _S_INSN) + elif any(known): + out.append("".join(f"{b:02x}" if known[i] else "??" + for i, b in enumerate(data)) + "\n", _S_INSN) + else: + out.append("?" * (width * 2) + "\n", _S_SEP) + self.query_one("#trace-stack", Static).update(out) + + +class TraceTimeline(Static): + """The trace as a vertical bar: where you are, and where you've been. + + Tenet's timeline is a Qt widget you scroll and drag to zoom. A terminal + column can't do that, but it can do the part that matters — show the shape + of the trace and your position in it — with one row per N timestamps. + """ + + def __init__(self, **kw) -> None: + super().__init__("", **kw) + self.trace = None + self.idx = 0 + + def render(self) -> Text: + t = self.trace + out = Text() + h = max(self.size.height - 1, 1) + if t is None or not t.length: + return out + out.append(" timeline\n", _S_DIM) + h = max(h - 1, 1) + per = max(t.length / h, 1.0) + here = int(self.idx / per) + for row in range(h): + if row == here: + out.append(" \u25b6", _S_MNEM) + out.append(f" {int(row * per):>10,}\n", _S_ADDR) + else: + out.append(" \u2502\n", _S_SEP if row % 5 else _S_ADDR) + return out + + class LoadOptionsScreen(ModalScreen): """Ask how to load a file no loader recognised. @@ -2583,9 +4037,8 @@ class LoadOptionsScreen(ModalScreen): def compose(self) -> ComposeResult: from .formats import PROCESSORS self._all = list(PROCESSORS) - with Vertical(id="pal-box"): - yield Static(" unrecognised file \u2014 how should IDA load it?", - id="pal-title", markup=False) + with Vertical(id="pal-box") as box: + box.border_title = Text("unrecognised file \u2014 how should IDA load it?") yield Static(f" {os.path.basename(self._path)} ({self._nbytes:,} bytes) " f"\u2014 no loader matched; without a processor IDA " f"assumes x86 at 0", id="load-note", markup=False) @@ -2648,8 +4101,8 @@ class LoadOptionsScreen(ModalScreen): ol.add_options(opts) if rows: ol.highlighted = 0 - self.query_one("#pal-title", Static).update( - f" unrecognised file \u2014 processor? ({len(rows)})") + self.query_one("#pal-box").border_title = Text( + f"unrecognised file \u2014 processor? ({len(rows)})") def action_cursor_down(self) -> None: ol = self.query_one(OptionList) @@ -2706,8 +4159,8 @@ class ProjectPalette(ModalScreen): self._results: list[dict] = [] def compose(self) -> ComposeResult: - with Vertical(id="pal-box"): - yield Static(" binaries", id="pal-title", markup=False) + with Vertical(id="pal-box") as box: + box.border_title = Text("binaries") yield Input(placeholder="filter binaries\u2026 \u2191\u2193 select \u00b7 " "Enter switch \u00b7 Esc close", id="pal-input") yield OptionList(id="pal-list") @@ -2754,8 +4207,8 @@ class ProjectPalette(ModalScreen): # you are rather than at whatever sorts first. active = next((i for i, e in enumerate(rows) if e["active"]), 0) ol.highlighted = active - self.query_one("#pal-title", Static).update( - f" binaries: {len(rows)} of {len(self._entries)}") + self.query_one("#pal-box").border_title = Text( + f"binaries: {len(rows)} of {len(self._entries)}") def action_cursor_down(self) -> None: ol = self.query_one(OptionList) @@ -2797,7 +4250,8 @@ class ConfirmScreen(ModalScreen): self._note = note def compose(self) -> ComposeResult: - with Vertical(id="confirm-box"): + with Vertical(id="confirm-box") as box: + box.border_title = Text("\u26a0 confirm") yield Static(self._message, id="confirm-msg", markup=False) if self._note: yield Static(self._note, id="confirm-note", markup=False) @@ -2810,8 +4264,43 @@ class ConfirmScreen(ModalScreen): self.dismiss(False) -_LOGO_PATH = os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "logo.ans") +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +_LOGO_PATH = os.path.join(_REPO_ROOT, "logo.ans") +#: The same artwork as a real image, for terminals that can draw one. logo.ans +#: is half-blocks (two pixels per cell); this is a transparent PNG at 768px. +LOGO_PNG = os.path.join(_REPO_ROOT, "logo.png") +_LOGO_BOX = (60, 33) # the most room the splash will give the art +#: Rows the loading box spends on everything that is not the artwork: border 2, +#: padding 2, the art's margin 1, title 1, note 1 + margin 1, help 1 + margin 1. +LOGO_CHROME_ROWS = 10 +#: Below this the image is a postage stamp; show the text splash instead. +LOGO_MIN_ROWS = 8 +_logo_cells: tuple[int, int] | None = None + + +def logo_cells(max_rows: int | None = None) -> tuple[int, int]: + """Cell footprint for the image, derived from the artwork and the terminal's + real cell size rather than hardcoded. + + Cells are nowhere near square (9x22 px here, 1:2.44), so a fixed box picked + for one aspect ratio stretches any other. Recomputing means the art can be + replaced without anyone remembering to edit a constant. + + ``max_rows`` shrinks it to the room actually available. The terminal scales + the image into whatever cell box we place it in, so there is no reason for + the splash to be all-or-nothing -- and it WAS all-or-nothing: a 31-row pane + is one row short of the natural size, so the logo silently disappeared + rather than being drawn a little smaller. + """ + global _logo_cells + if _logo_cells is None: + px = kittygfx.png_size(LOGO_PNG) + _logo_cells = kittygfx.fit(px, *_LOGO_BOX) if px else _LOGO_BOX + if max_rows is None or max_rows >= _logo_cells[1]: + return _logo_cells + px = kittygfx.png_size(LOGO_PNG) + return (kittygfx.fit(px, _LOGO_BOX[0], max(max_rows, 1)) if px + else (_LOGO_BOX[0], max(max_rows, 1))) _logo_cache: object = False # False == not yet loaded (None == absent/unreadable) @@ -2842,17 +4331,45 @@ class LoadingScreen(ModalScreen): super().__init__() self._title = title self._note = note + self._image = False # drawing the real image, not the block art + self._cells: tuple[int, int] | None = None # image size, in cells + self._last_place = 0.0 # throttles re-anchoring after a repaint + + def _room(self) -> int: + """Rows left for artwork once the box's own furniture is paid for.""" + return self.app.size.height - LOGO_CHROME_ROWS + + def _fits(self, rows: int) -> bool: + """Room for art of exactly ``rows`` (the block art cannot be resized).""" + return self._room() >= rows and self.app.size.width >= 64 def compose(self) -> ComposeResult: with Vertical(id="loading-box"): - logo = _load_logo() - # Only show the splash art when the terminal can fit it plus the - # title/note/help + box chrome; otherwise fall back to a text-only - # overlay so nothing important is clipped off a small screen. - if logo is not None: - sz = self.app.size - n = len(logo.split("\n")) - if sz.height >= n + 9 and sz.width >= 64: + # A terminal that can draw a real image gets one: same artwork, + # same cell footprint, ~10x the linear resolution of the block art. + # The image is anchored to screen cells rather than composited by + # Textual (no unicode-placeholder support here), so the widget is + # only reserved blank space -- see _place_logo. + # Scale the image to the room there is, rather than demanding its + # natural size and vanishing when one row is missing. + room = self._room() + cols, rows = logo_cells(room) + self._cells = (cols, rows) + kittygfx.log(f"compose: supported={kittygfx.supported()} " + f"app.size={self.app.size} room={room} " + f"cells={cols}x{rows} natural={logo_cells()}") + if (kittygfx.supported() and self.app.size.width >= 64 + and room >= LOGO_MIN_ROWS): + self._image = True + blank = Static("\n" * (rows - 1), id="loading-image") + blank.styles.height = rows + yield blank + else: + logo = _load_logo() + # Only show the splash art when the terminal can fit it plus + # the title/note/help + box chrome; otherwise fall back to a + # text-only overlay so nothing important is clipped. + if logo is not None and self._fits(len(logo.split("\n"))): # Align.center, not the box's align-horizontal: the 1fr # title/note siblings make the child group span the full # width, so container alignment has nothing left to centre. @@ -2867,6 +4384,61 @@ class LoadingScreen(ModalScreen): self.query_one("#loading-note", Static).update(text) except Exception: # noqa: BLE001 -- not mounted yet / already gone pass + # Textual doesn't know the image is there, so a repaint can drop it. + # Re-anchoring is one short escape with no image data; throttled so a + # chatty progress callback can't turn it into a flicker. + if self._image: + now = time.monotonic() + if now - self._last_place > 0.2: + self._last_place = now + self._place_logo() + + # -- the image, which Textual knows nothing about ---------------------- # + def _place_logo(self) -> None: + """Anchor the image over the blank cells reserved for it. + + Deferred to after a refresh because a widget has no screen region until + it has been laid out, and re-run on resize because the region moves. + """ + if not self._image: + return + try: + region = self.query_one("#loading-image", Static).region + except Exception as e: # noqa: BLE001 -- gone already + kittygfx.log(f"place_logo: no widget ({e})") + return + kittygfx.log(f"place_logo: region={region}") + if not region.width or not region.height: + return + # The reserved region is the truth about how much room there is; the + # image is scaled into exactly it, so a resize needs no relayout. + cols, rows = self._cells or logo_cells() + rows = min(rows, region.height) + col = region.x + max((region.width - cols) // 2, 0) # centre it + kittygfx.place(region.y, col, min(cols, region.width), rows) + + def on_mount(self) -> None: + if not self._image: + return + # Upload HERE, not from the launcher: Textual is on the alternate screen + # by now, and an image uploaded to the primary screen cannot be placed + # from the alternate one -- the placement reports success and draws + # nothing at all. + if not kittygfx.upload(LOGO_PNG): + self._image = False + return + self.call_after_refresh(self._place_logo) + + def on_resize(self) -> None: + if self._image: + kittygfx.clear() + self.call_after_refresh(self._place_logo) + + def on_unmount(self) -> None: + # The image is anchored to the screen, not owned by the compositor, so + # it would sit there over the disassembly forever if we didn't say so. + if self._image: + kittygfx.clear() def action_hide(self) -> None: self.dismiss() @@ -2902,14 +4474,24 @@ class StructEditor(ModalScreen): selected struct; Ctrl+S declares (creates or updates) it; Ctrl+N starts a new one; Delete removes the highlighted struct; Esc returns to the list then closes. + + ``/`` from the list opens a fuzzy filter over the struct names -- the same + key and the same matcher as the code views' search and the symbol palette, + rather than a third way to find something by typing. """ BINDINGS = [ Binding("ctrl+s", "save", "Save", priority=True), Binding("ctrl+n", "new", "New", priority=True), Binding("ctrl+y", "copy", "Copy", priority=True), + Binding("slash", "filter", "Filter", show=False), Binding("delete,d", "delete", "Delete", show=False), Binding("escape", "close", "Close"), + # Only ever reached while the FILTER has focus: a focused OptionList + # consumes up/down itself, so these move its highlight from the prompt + # (type to narrow, arrow to pick, exactly like the symbol palette). + Binding("up", "cursor_up", "Up", show=False), + Binding("down", "cursor_down", "Down", show=False), ] NEW_TEMPLATE = "struct NewStruct\n{\n int field;\n};\n" @@ -2919,7 +4501,9 @@ class StructEditor(ModalScreen): def __init__(self, program: Program) -> None: super().__init__() self._program = program - self._structs: list[Struct] = [] + self._all: list[Struct] = [] # every struct the database has + self._structs: list[Struct] = [] # the VISIBLE rows (== _all when unfiltered) + self._filter = "" self._loaded: str | None = None # name currently in the editor self._loaded_src: str | None = None # its source, to detect unsaved edits @@ -2942,16 +4526,23 @@ class StructEditor(ModalScreen): then() def compose(self) -> ComposeResult: - with Vertical(id="se-box"): + with Vertical(id="se-box") as box: + box.border_title = Text("structs / types") + box.border_subtitle = Text("Esc close") with Horizontal(id="se-panes"): with Vertical(id="se-left"): - yield Static(" structs", id="se-title") + yield Static("structs", id="se-title") + yield Input(placeholder="fuzzy filter\u2026 \u2191\u2193 pick \u00b7 " + "Enter edit \u00b7 Esc clear", + id="se-filter") yield OptionList(id="se-list") with Vertical(id="se-right"): - yield Static(" C definition", id="se-hint") - yield TextArea("", id="se-edit") + yield Static("C definition", id="se-hint") + yield CTextArea("", id="se-edit") + # Esc is on the border subtitle; repeating it here cost the row. yield Static( - "Enter edit · Ctrl+S save · Ctrl+Y copy · Ctrl+N new · d/Del delete · Esc close", + "Enter edit · / filter · Ctrl+S save · Ctrl+Y copy · Ctrl+N new · " + "d/Del delete", id="se-status") def on_mount(self) -> None: @@ -2969,26 +4560,134 @@ class StructEditor(ModalScreen): self.app.call_from_thread(self._populate, structs, select) def _populate(self, structs: list[Struct], select: str | None) -> None: - self._structs = structs + self._all = structs + self._show(select) + + def _show(self, select: str | None = None) -> None: + """Rebuild the list from ``_all`` through the current filter. + + ``_structs`` stays the VISIBLE rows, because every action here indexes + it by the list's highlighted row -- load, delete and the confirm dialogs + would all address the wrong struct if the two ever drifted apart. + """ + q = self._filter.strip() + # A struct we were told to select but the filter would hide (a rename, or + # one just created under a stale filter) beats the filter -- otherwise a + # successful save looks like the struct vanished. + if q and select is not None and _fuzzy(select, q) is None: + inp = self.query_one("#se-filter", Input) + inp.value = "" + inp.display = False + self._filter = q = "" + rows: list[tuple[Struct, tuple[int, ...]]] = [] + if q: + scored = [] + for s in self._all: + m = _fuzzy(s.name, q) + if m is not None: + scored.append((m[0], m[1], s)) + scored.sort(key=lambda t: (-t[0], t[2].name)) + rows = [(s, pos) for _, pos, s in scored] + else: + rows = [(s, ()) for s in self._all] + self._structs = [s for s, _ in rows] + ol = self.query_one("#se-list", OptionList) ol.clear_options() - for s in structs: + # Ragged rows read as noise, so the name column is padded to the widest + # name actually present (capped, so one monstrous C++ mangling can't + # push the size/count columns off the pane). + width = min(max((len(s.name) for s, _ in rows), default=0), 22) + opts = [] + for s, pos in rows: kw = "union" if s.is_union else "struct" + name = s.name if len(s.name) <= width else s.name[:width - 1] + "\u2026" label = Text() - label.append(s.name, _S_LABEL) - label.append(f" {s.size:#x} {s.members}f {kw}", _S_DIM) - ol.add_option(Option(label)) - if structs: + nm = Text(f"{name:<{width}}", style=_S_LABEL) + for p in pos: + if p < len(name): + nm.stylize(_S_NAME_MATCH, p, p + 1) + label.append_text(nm) + label.append(f"{s.size:>#7x}{s.members:>4}f {kw}", _S_DIM) + opts.append(Option(label)) + ol.add_options(opts) + if rows: idx = 0 if select is not None: - idx = next((i for i, s in enumerate(structs) if s.name == select), 0) + idx = next((i for i, s in enumerate(self._structs) + if s.name == select), 0) ol.highlighted = idx + cap = "structs" + if q: + cap = f"structs {len(rows)}/{len(self._all)}" + self.query_one("#se-title", Static).update(cap) def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: i = event.option_index if 0 <= i < len(self._structs): self._confirm_discard(lambda n=self._structs[i].name: self._load(n)) + # -- filter ("/" from the list) ----------------------------------------- # + def action_filter(self) -> None: + """'/': fuzzy-filter the struct list, as you type.""" + # A '/' typed into the definition is a division operator, and one typed + # into the filter is the Input's business; only the list opens this. + if self.focused is not self.query_one("#se-list", OptionList): + return + inp = self.query_one("#se-filter", Input) + inp.display = True + inp.value = self._filter + inp.focus() + + def _clear_filter(self, focus_list: bool = True) -> None: + inp = self.query_one("#se-filter", Input) + inp.value = "" + inp.display = False + self._filter = "" + self._show(select=self._loaded) + if focus_list: + self.query_one("#se-list", OptionList).focus() + + def on_input_changed(self, event: Input.Changed) -> None: + # Modal Input messages bubble to the App (which owns #func-filter and + # the search prompt) -- stop them here or they drive the main screen. + if event.input.id != "se-filter": + return + event.stop() + self._filter = event.value + self._show() + + def on_input_submitted(self, event: Input.Submitted) -> None: + if event.input.id != "se-filter": + return + event.stop() + # Enter on a filtered list loads the highlighted struct, like Enter on + # the list itself -- typing a name and pressing Enter should open it. + ol = self.query_one("#se-list", OptionList) + i = ol.highlighted + if i is not None and 0 <= i < len(self._structs): + self._confirm_discard(lambda n=self._structs[i].name: self._load(n)) + else: + ol.focus() + + def _filter_focused(self) -> bool: + return self.focused is self.query_one("#se-filter", Input) + + def action_cursor_up(self) -> None: + self._move_highlight(-1) + + def action_cursor_down(self) -> None: + self._move_highlight(1) + + def _move_highlight(self, delta: int) -> None: + if not self._filter_focused(): + return # the list has focus and moves itself + ol = self.query_one("#se-list", OptionList) + if not ol.option_count: + return + cur = ol.highlighted or 0 + ol.highlighted = max(0, min(cur + delta, ol.option_count - 1)) + @work(thread=True, exclusive=True, group="se-load") def _load(self, name: str) -> None: try: @@ -3047,6 +4746,8 @@ class StructEditor(ModalScreen): self._set_status(f"save failed — {msg}", error=True) return self._loaded = name + if name: + self.app.journal.record("type", None, name) if formatted: # Reformat the editor to IDA's canonical layout (keeps cursor at top). ta = self.query_one("#se-edit", TextArea) @@ -3098,6 +4799,7 @@ class StructEditor(ModalScreen): return if getattr(self.app, "_dirty", None) is not None: self.app._dirty = True + self.app.journal.record("del_type", None, name) self._refresh() self._set_status(f"deleted {name}") @@ -3113,10 +4815,15 @@ class StructEditor(ModalScreen): self._set_status(f"copied {what} ({n} chars) to clipboard") def action_close(self) -> None: - # A stray Esc while editing returns to the list instead of discarding. + # Esc backs out one level at a time: out of the definition, then out of + # the filter, and only then out of the dialog. Closing the editor from + # under a half-typed filter is the kind of thing you only do once. if self.focused is self.query_one("#se-edit", TextArea): self.query_one("#se-list", OptionList).focus() return + if self._filter_focused() or self._filter: + self._clear_filter() + return self._confirm_discard(lambda: self.dismiss(None), "Discard unsaved changes and close?") @@ -3167,7 +4874,12 @@ class IdaCommands(Provider): app.action_strings), ("Switch binary…", "another binary in the project (Ctrl+O)", app.action_switch_binary), - ("Keyboard shortcuts", "the key cheatsheet (F1)", app.action_help), + ("Search database…", "disassembly text or a byte pattern with " + "wildcards (Ctrl+F)", app.action_find), + ("Export findings…", "your comments, names and types as markdown " + "(Ctrl+E)", app.action_export), + ("Keyboard shortcuts", "the key cheatsheet (F1 or H)", + app.action_help), ("Follow symbol under cursor", "jump to the referenced symbol (Enter)", lambda: va("follow")), ("Show xrefs to symbol", "cross-references to the cursor symbol (x)", @@ -3180,6 +4892,14 @@ class IdaCommands(Provider): ("Hex view", "raw bytes at the cursor (\\)", app.action_hex), ("Split view (listing ⇄ pseudocode)", "side-by-side synced views (s)", app.action_toggle_split), + ("Graph view (control flow)", + "the function's basic blocks as a graph (Space)", + app.action_toggle_graph), + ("Graph: cycle zoom", + "full → compact → collapsed (z, in the graph)", + lambda: va("zoom")), + ("Graph: toggle minimap", + "the overview box (m, in the graph)", lambda: va("minimap")), ("Rename symbol…", "rename the symbol under the cursor (n)", lambda: va("rename")), ("Set type / prototype…", "retype the symbol under the cursor (y)", @@ -3194,7 +4914,19 @@ class IdaCommands(Provider): lambda: va("make_string")), ("Undefine", "undefine the item at the cursor (u)", lambda: va("undefine")), - ("Toggle opcode bytes", "cycle the opcode-bytes column (o)", + ("Literal format: next", "cycle the literal under the cursor (o)", + lambda: va("op_format", "cycle")), + ("Literal format: previous", "the other way round (O)", + lambda: va("op_format", "back")), + *((f"Literal format: {label}", f"show the literal as {label} ({fmt})", + (lambda f=fmt: va("op_format", f))) + for fmt, label in (("hex", "hexadecimal"), ("dec", "decimal"), + ("oct", "octal"), ("bin", "binary"), + ("char", "a character"), + ("offset", "an offset (reference)"), + ("stack", "a stack variable"), + ("default", "IDA's own choice"))), + ("Toggle opcode bytes", "cycle the opcode-bytes column (B)", lambda: va("toggle_opcodes")), ("Structs / types editor", "view + edit local types (Ctrl+T)", app.action_structs), @@ -3227,11 +4959,11 @@ class IdaTui(App): #left { width: 30%; min-width: 42; max-width: 44; border-right: solid $panel; } #func-table { height: 1fr; } #func-filter { dock: top; } - DisasmView { width: 1fr; padding: 0 1; } DecompView { width: 1fr; } ListingView { width: 1fr; padding: 0 1; } #panes.split ListingView { border-right: tall $panel-lighten-2; } HexView { width: 1fr; padding: 0 1; } + GraphView { width: 1fr; padding: 0 1; } #search { height: 1; border: none; padding: 0 1; background: $primary-darken-2; color: $text; @@ -3256,42 +4988,77 @@ class IdaTui(App): height: 1; border: none; padding: 0 1; background: $primary-darken-3; color: $text; } + #export { + height: 1; border: none; padding: 0 1; + background: $success-darken-3; color: $text; + } #status { height: 1; background: $panel; color: $text; padding: 0 1; } .decomp-loading { width: 100%; height: 100%; content-align: center middle; background: $panel-darken-1; } - QuitScreen { align: center middle; } - #quit-box { width: 64; height: auto; border: thick $warning; background: $panel; } - #quit-title { dock: top; height: 1; background: $warning; color: $background; text-style: bold; padding: 0 1; } - #quit-list { height: auto; padding: 1 2 0 2; } - #quit-help { height: 1; color: $text-muted; padding: 0 2; margin-top: 1; } - HelpScreen { align: center middle; } - #help-box { width: auto; max-width: 98%; height: auto; max-height: 90%; - border: thick $accent; background: $panel; } - #help-title { dock: top; height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } + /* ---- modal chrome --------------------------------------------------- + One shape for every dialog: a THIN round border in a colour that + recedes, and the title set into the top border rather than painted as a + solid bar across the full width. The old `thick` frame plus an inverse + title bar was two heavy rectangles around content that is itself the + only thing worth looking at. Size stays per-box; everything else is + here, once. */ + #quit-box, #help-box, #xref-box, #pal-box, #se-box, #confirm-box, + #loading-box, #busy-box { + background: $panel; + border: round $panel-lighten-3; + border-title-color: $accent; + border-title-style: bold; + border-title-align: left; + border-subtitle-color: $text-muted; + border-subtitle-align: right; + } + /* The two that ask a question you can answer wrongly keep the warning hue. */ + #quit-box, #confirm-box { border: round $warning; + border-title-color: $warning; } + /* Lists inside a dialog are content, not another framed panel: drop the + stock OptionList border/background so the dialog has ONE edge. Focus + still reads from the highlight bar (bright when focused, dim when not). */ + #quit-box OptionList, #help-box OptionList, #xref-box OptionList, + #pal-box OptionList, #se-box OptionList { + border: none; background: transparent; padding: 0 1; + } + #quit-box { width: 64; height: auto; padding: 1 1 0 1; } + #quit-list { height: auto; padding: 0 1; } + #quit-help { height: 1; color: $text-muted; padding: 0 1; margin-top: 1; } + #help-box { width: auto; max-width: 98%; height: auto; max-height: 90%; } #help-body { height: auto; max-height: 100%; width: auto; padding: 1 1; } #help-cols { height: auto; width: auto; } .help-col { height: auto; width: auto; margin-right: 1; } .help-card { height: auto; width: auto; padding: 0 1; border: round $panel-lighten-2; } #help-foot { dock: bottom; height: 1; color: $text-muted; padding: 0 2; } - XrefsScreen { align: center middle; } - #xref-box { width: 84; max-height: 70%; height: auto; border: thick $accent; background: $panel; } - #xref-title { dock: top; height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } + #xref-box { width: 84; max-height: 70%; height: auto; padding: 0 0; } #xref-list { height: auto; max-height: 100%; } - /* every #pal-box palette centres, not just the symbol one */ - SymbolPalette, StringsPalette, ProjectPalette, - LoadOptionsScreen { align: center middle; } + /* A dialog centres because it is a dialog — not because someone remembered + to add it to a list. That list is how SearchPalette shipped pinned to the + top of the screen, and the comment that used to sit here ("every #pal-box + palette centres, not just the symbol one") was itself the second time. + ModalScreen matches subclasses, so every modal below inherits this and + the next one gets it for free; `modal_centering` in the pilot suite + fails if one ever opts out by accident. Textual's own Ctrl+P palette is a + ModalScreen too and wants its stock top alignment, so it opts out here, + deliberately and visibly. */ + ModalScreen { align: center middle; } + CommandPalette { align: center top; } /* Give the stock Ctrl+P command palette side padding instead of full width; the input + results inherit this width (results is an overlay, so pin it). */ CommandPalette > Vertical { width: 80%; max-width: 120; } CommandPalette #--results { width: 100%; } - #pal-box { width: 96; max-width: 92%; height: auto; max-height: 80%; - border: thick $accent; background: $panel; } - #pal-title { dock: top; height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } + #pal-box { width: 96; max-width: 92%; height: auto; max-height: 80%; } #pal-input { border: none; height: 1; margin: 0 1; background: $panel; color: $text; } #pal-list { height: auto; max-height: 24; } + #trace-dock { dock: right; width: 34; background: $surface; border-left: solid $panel; } + #trace-head { height: 2; padding: 0 1; background: $panel; } + #trace-regs { height: auto; padding: 1 0 0 0; } + #trace-stack { height: auto; padding: 1 0 0 0; } + #trace-timeline { height: 1fr; padding: 1 0 0 0; } #load-note { height: 2; padding: 1 1 0 1; color: $text-muted; } /* Cap the processor list so the ADDRESS FIELD is always on screen: with the palette default (24) the box outgrew the terminal and the field you need @@ -3299,31 +5066,33 @@ class IdaTui(App): LoadOptionsScreen #pal-list { max-height: 12; } #load-base { border: none; height: 1; margin: 1 1 0 1; background: $panel; color: $text; } #load-help { height: 1; padding: 0 1; color: $text-muted; } - #confirm-note { height: auto; padding: 0 1; color: $text-muted; } - StructEditor { align: center middle; } - #se-box { width: 90%; height: 84%; border: thick $accent; background: $panel; } + #confirm-note { height: auto; color: $text-muted; } + #se-box { width: 90%; height: 84%; } #se-panes { height: 1fr; } - #se-left { width: 38; border-right: solid $accent; } + #se-left { width: 38; border-right: solid $panel-lighten-3; } #se-right { width: 1fr; } - #se-title, #se-hint { height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } + /* Column captions, not title bars: the dialog already has one title. */ + #se-title, #se-hint { height: 1; color: $text-muted; text-style: bold; + padding: 0 1; } + /* Hidden until '/'; same one-row prompt shape as the app's own prompts. */ + #se-filter { display: none; height: 1; border: none; padding: 0 1; + background: $primary-darken-2; color: $text; } #se-list { height: 1fr; } - #se-edit { height: 1fr; border: none; } - #se-status { height: 1; background: $panel-darken-2; color: $text-muted; padding: 0 1; } - ConfirmScreen { align: center middle; } - #confirm-box { width: 60; height: auto; border: thick $warning; - background: $panel; padding: 1 2; } + #se-edit { height: 1fr; border: none; padding: 0 1; } + /* A footer, but a quiet one: a shade of the dialog's own background + instead of the old near-black bar. */ + #se-status { height: 1; background: $panel-darken-1; color: $text-muted; + padding: 0 1; } + #confirm-box { width: 60; height: auto; padding: 1 2; } #confirm-msg { height: auto; } #confirm-help { height: 1; color: $text-muted; margin-top: 1; } - LoadingScreen { align: center middle; } - #loading-box { width: 72; height: auto; border: thick $accent; - background: $panel; padding: 1 2; } + #loading-box { width: 72; height: auto; padding: 1 2; } #loading-logo { width: 100%; height: auto; margin-bottom: 1; } + #loading-image { width: 100%; margin-bottom: 1; } #loading-title { width: 1fr; height: 1; text-style: bold; } #loading-note { height: auto; color: $text-muted; margin-top: 1; } #loading-help { height: auto; color: $text-muted; margin-top: 1; } - BusyScreen { align: center middle; } - #busy-box { width: auto; min-width: 26; height: auto; border: thick $accent; - background: $panel; padding: 1 2; } + #busy-box { width: auto; min-width: 26; height: auto; padding: 1 2; } #busy-msg { height: 1; text-style: bold; } #busy-help { height: 1; color: $text-muted; margin-top: 1; } """ @@ -3332,12 +5101,32 @@ class IdaTui(App): Binding("q", "quit", "Quit"), Binding("ctrl+n", "symbols", "Symbols"), Binding("ctrl+t", "structs", "Structs"), + Binding("ctrl+f", "find", "Find"), + Binding("ctrl+e", "export", "Export", show=False), Binding("backslash", "hex", "Hex"), Binding("s", "toggle_split", "Split", show=False), + # IDA's own key for text/graph. Graph mode is opt-in and self-contained: + # with it off nothing else in the app does any extra work. + Binding("space", "toggle_graph", "Graph", show=False), Binding("quotation_mark,shift+f12", "strings", "Strings", show=False), Binding("ctrl+o", "switch_binary", "Binaries", show=False), Binding("ctrl+l", "load_options", "Reload as…", show=False), - Binding("f1", "help", "Keys", show=False), + # Trace stepping. ] / [ move one instruction, } / { step over a + # call by following the stack pointer. + # Seeking, as opposed to stepping: jump to the next/previous time THIS + # thing was touched, where "this thing" is whatever the focused view + # addresses — an instruction in the code views, a byte in hex. + Binding("greater_than_sign", "seek_next_hit", "Next hit", show=False), + Binding("less_than_sign", "seek_prev_hit", "Prev hit", show=False), + Binding("W", "seek_reg_write", "Reg writes", show=False), + Binding("right_square_bracket", "step_fwd", "Step", show=False), + Binding("left_square_bracket", "step_back", "Step back", show=False), + Binding("right_curly_bracket", "step_over_fwd", "Step over", show=False), + Binding("left_curly_bracket", "step_over_back", "Step over back", show=False), + # H as well as F1: terminals and multiplexers swallow function keys all + # the time (and the one that does it is upstream of us, so there is + # nothing to fix on this side), which left the cheatsheet unreachable. + Binding("f1,H", "help", "Keys", show=False), Binding("g", "goto", "Goto"), Binding("slash", "filter", "Filter", show=False), Binding("ctrl+b", "toggle_functions", "Names", show=False), @@ -3348,7 +5137,7 @@ class IdaTui(App): def __init__(self, open_path: str | None = None, keepalive: bool = True, rpc_path: str | None = None, ttl: int = 1800, - project=None, load_args: str = "") -> None: + project=None, load_args: str = "", trace_path: str = "") -> None: super().__init__() # Project mode is additive: with no project this is the plain # single-binary app, unchanged. @@ -3362,27 +5151,37 @@ class IdaTui(App): self._load_for_label = None # project binary the dialog is for self._no_functions = False # analysis produced nothing at all self._flash: str | None = None # message a pending reload must keep + self._flash_until = 0.0 # ...until this monotonic time self._pending_switch = None # switch waiting on that answer self._nav_seq = 0 # bumped per navigation; drops stale ones + #: Literal positions for the decompilation being loaded (worker thread + #: -> the view, handed over when the pseudocode is applied). + self._pending_nums: dict = {} # None = teardown wasn't an explicit quit (crash/kill): save defensively. # False = the user chose discard, or we already saved on the way out. self._save_on_exit: bool | None = None self._index = None # project-wide symbol/string index if project is not None: from .index import ProjectIndex - from .pool import WorkerPool - self._pool = WorkerPool(project, ttl=ttl) + from .pool import DatabasePool + self._pool = DatabasePool(project, ttl=ttl) self._index = ProjectIndex( os.path.join(project.index_dir, "project.db")) self._binary = project.refs[0].label open_path = project.refs[0].staged self._open_path = open_path self._ttl = ttl - self._load_args = load_args or "" # IDA switches for a headerless blob + self._load_args = load_args or "" # first-open options for a headerless blob + self._new_database = False # Ctrl+L asks Code Mode for a fresh IDB + self._title = (os.path.basename(open_path) if open_path else "") + #: Where we are in the execution trace, and everything that moves us. + #: Owns the trace state; the _trace/_t/_trail_* properties below + #: forward to it. + self.trace_ctl = TraceController(self, trace_path or "") self._do_keepalive = keepalive self._rpc_path = rpc_path self._rpc = None - self.client: WorkerClient | None = None + self.client: CodeModeClient | None = None self.program: Program | None = None self._loading_screen: LoadingScreen | None = None self._ka = None @@ -3397,8 +5196,9 @@ class IdaTui(App): # ONE notion of "which pane you're in": _active, kept in step with focus # (on_descendant_focus does that while split). There used to be a second, # _pref, but it was only ever assigned "listing" — see _code_mode(). - self._active = "listing" # currently shown view (in split: the focused pane) + self._active = ViewMode.LISTING # currently shown view (in split: the focused pane) self._split = False # side-by-side listing + pseudocode + self._graph_sticky = False # stay in graph mode across navigations self._split_eamap: list[list[int]] = [] # split: decomp line -> instr EAs self._split_ea2line: dict[int, int] = {} # split: instr EA -> decomp line self._split_range: tuple[int, int] | None = None # decomp'd fn ea span @@ -3411,11 +5211,15 @@ class IdaTui(App): self._conn_screen: LoadingScreen | None = None self._reconnecting = False # a reconnect attempt is in flight self._search_ctx: tuple[object | None, int] = (None, 1) - self._rename_ctx: tuple[object | None, str] = (None, "") - self._rename_addr: int | None = None # set for address-based (listing) naming - self._comment_ctx: tuple[object | None, int, str] = (None, 0, "") - self._retype_ctx: tuple[object | None, str, int, str] = (None, "", 0, "") - self._makedata_ctx: tuple[object | None, int] = (None, 0) + #: The one-line prompts above the footer. Each holds its own context + #: for exactly as long as it is on screen; see idatui/prompt.py. + self.prompts = PromptBar(self, "search", "rename", "comment", + "retype", "makedata", "goto", "export") + #: Everything that writes to the database (idatui/edit_ctl.py). + self.edits = EditController(self) + #: What those writes were, so the findings export can say which + #: comments and names are YOURS -- the database cannot (journal.py). + self.journal = Journal() self._xref_focus_name: str | None = None self._dirty = False @@ -3425,14 +5229,22 @@ class IdaTui(App): fp = FunctionsPanel(id="left") fp.display = False # overlay-first: reveal the docked pane with Ctrl+B yield fp - # The unified continuous listing is the one code view; DisasmView is - # deprecated (kept only for DisasmModel, still used by the domain). + # The unified continuous listing is the one code view. (DisasmModel + # is still used by the domain to index a function's instructions.) lst = ListingView() yield lst yield DecompView() hx = HexView() hx.display = False yield hx + gv = GraphView() + gv.display = False + yield gv + # Docked right and only shown once a trace is loaded, so a normal + # session looks exactly as it did. + td = TraceDock() + td.display = False + yield td si = Input(id="search") si.display = False si.can_focus = False @@ -3457,6 +5269,10 @@ class IdaTui(App): gi.display = False gi.can_focus = False yield gi + xi = Input(id="export") + xi.display = False + xi.can_focus = False + yield xi # markup=False: the status is plain text full of [listing]/[split]/[label] # markers and symbol names that may contain brackets. With Textual markup # on, a single-word marker parses as a style tag and is silently eaten — @@ -3475,8 +5291,8 @@ class IdaTui(App): if self._rpc_path: self._start_rpc() # A file no loader recognises has to be described before it can be - # opened, so ask BEFORE the worker starts — once IDA has made a database - # the answer is baked in and changing it means deleting the .i64. + # opened, so ask BEFORE Code Mode creates it — once IDA has made a database + # the answer is baked in and changing it requires a fresh-IDB reopen. if self._project is not None: ref = self._pending_load_ref() if ref is not None: @@ -3507,6 +5323,11 @@ class IdaTui(App): if os.path.exists(self._open_path + ".i64") or os.path.exists( os.path.splitext(self._open_path)[0] + ".i64"): return False + try: + if registered_database(self._open_path): + return False + except Exception: + pass # connect() will surface registry failures with full diagnostics return needs_load_options(self._open_path) def action_load_options(self) -> None: @@ -3520,13 +5341,18 @@ class IdaTui(App): forward. """ if not self._can_reload(): - self._status("nothing to reload") + if self.client is not None and self.client.backend == "gui": + self._status( + "reload unavailable for a GUI-owned database — reopen it in IDA") + else: + self._status("nothing to reload") return n = len(self._func_index) if self._func_index else 0 note = ("this image has no functions, so nothing is lost" if n == 0 else - f"discards the database for this binary \u2014 {n} functions, " - f"plus any names and comments you've added") + f"discards the database for this binary \u2014 {n} " + f"function{'s' if n != 1 else ''}, plus any names and comments " + f"you've added") self.push_screen(ConfirmScreen("Reload with different options?", note), self._on_reload_confirmed) @@ -3538,10 +5364,13 @@ class IdaTui(App): ref = self._project.by_label(self._binary) if ref is not None: path, label = ref.source, ref.label - # Drop the worker first: it holds the database open, and the .i64 can't - # be removed (or rebuilt) underneath a live one. - self._release_worker() - self._drop_database() + # Release our lease first. Code Mode waits for a managed worker's final + # lease grace, then creates the replacement IDB atomically. A GUI-backed + # database is rejected by _can_reload(): the TUI must never close it. + self._release_database() + self._new_database = True + if label is not None and self._pool is not None: + self._pool.recreate_on_next_open(label) self._reset_for_reload() self._load_args = "" if label is not None and self._project is not None: @@ -3553,7 +5382,9 @@ class IdaTui(App): self._pending_switch = None self._ask_load_options(path, label=label) - def _release_worker(self) -> None: + def _release_database(self) -> None: + if self.program is not None: + self.program.close() if self._pool is not None and self._binary is not None: try: self._pool.evict(self._binary, save=False) @@ -3567,23 +5398,6 @@ class IdaTui(App): self.client = None self.program = None - def _drop_database(self) -> None: - """Remove the .i64 (and any unpacked scratch) so the next open re-reads - the raw image with new options.""" - base = self._open_path - if self._project is not None and self._binary is not None: - ref = self._project.by_label(self._binary) - if ref is not None: - base = ref.staged - if not base: - return - for suffix in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - for cand in (base + suffix, os.path.splitext(base)[0] + suffix): - try: - os.remove(cand) - except OSError: - pass - def _reset_for_reload(self) -> None: self._no_functions = False self._func_index = None @@ -3627,6 +5441,11 @@ class IdaTui(App): if os.path.exists(ref.db) or os.path.exists( os.path.splitext(ref.staged)[0] + ".i64"): return None # already analysed: the .i64 records how + try: + if registered_database(ref.staged, output_database=ref.db): + return None + except Exception: + pass from .formats import needs_load_options return ref if needs_load_options(ref.source) else None @@ -3680,17 +5499,38 @@ class IdaTui(App): asyncio.get_running_loop().create_task(_serve()) - async def on_unmount(self) -> None: - if self._rpc is not None: - await self._rpc.stop() - # -- status helper ----------------------------------------------------- # - def _status(self, text: str) -> None: - if self._binary: # project mode: always say which binary you're in - text = f"[{self._binary}] {text}" + def _status(self, text: str, priority: bool = False) -> None: + """Write the status bar. ``priority`` marks the RESULT of something the + user did. + + An action's result is written, and then the reload it triggered writes + its own idle status on top — cursor moved, filter re-applied, functions + re-counted. I patched that at five separate call sites before admitting + it's one problem: routine chatter must not outrank an answer. A priority + message holds the bar briefly and is cleared by the next keypress, i.e. + when the user has read it and moved on. + """ + import time as _time + if priority: + self._flash = text + self._flash_until = _time.monotonic() + 8.0 + elif self._flash and _time.monotonic() < self._flash_until: + text = self._flash + # Always say WHICH file this is. In project mode that's the binary's + # label; otherwise the filename we opened. Cheap on purpose — _module() + # asks the worker, and this runs on every status write. + tag = self._binary or self._title + if tag: + text = f"[{tag}] {text}" # An image with no functions at all is nearly always a blob described # wrongly, and that stays true as you scroll around — so it belongs in # the status bar, not in a one-off message the next write clobbers. + # It stops being true the moment a function exists, though: latching it + # meant the warning survived defining one with `p` and kept telling you + # the load was wrong when it no longer was. + if self._no_functions and self._func_index is not None and len(self._func_index): + self._no_functions = False if self._no_functions: text += " \u2014 no functions: wrong processor/base? Ctrl+L to reload" try: @@ -3723,9 +5563,10 @@ class IdaTui(App): # -- connection loss / recovery --------------------------------------- # def _handle_exception(self, error: BaseException) -> None: - """Intercept a lost-connection error from any worker so the whole app - doesn't die when the analysis server goes away (it can idle out, be - killed, or the box can sleep). Everything else crashes as usual.""" + """Intercept a lost Code Mode lease so the app can rediscover the DB. + + Everything unrelated to database connectivity crashes as usual. + """ from textual.worker import WorkerFailed orig = error.error if isinstance(error, WorkerFailed) else error if isinstance(orig, IDAConnectionError): @@ -3757,15 +5598,15 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="reconnect") def _reconnect(self) -> None: - # The worker died (segfault -> dropped socket). Respawn it: it re-opens - # and re-analyzes the binary in a fresh process, then we rebuild. + # The registered instance disappeared. Rediscover it; Code Mode may find + # a GUI/replacement worker, then we rebuild caches against the new handle. try: if self._open_path is None: self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args) client.connect(progress=lambda m: self.app.call_from_thread( self._conn_note, m)) except Exception as e: # noqa: BLE001 @@ -3773,7 +5614,7 @@ class IdaTui(App): return self.app.call_from_thread(self._after_reconnect, client, Program(client)) - def _after_reconnect(self, client: "WorkerClient", program: "Program") -> None: + def _after_reconnect(self, client: "CodeModeClient", program: "Program") -> None: self.client = client self.program = program self._reconnecting = False @@ -3793,13 +5634,13 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="connect") def _connect(self) -> None: try: - client = self._open_worker_client() + client = self._open_database_client() if client is None: return # the opener already reported + dismissed the overlay module = client.health().get("module", "?") if self._do_keepalive: - # Keep the session warm while we run; don't make it immortal, so - # it's reclaimed after the TUI closes. (No-op for the worker.) + # Compatibility shim: DatabaseHandle's SSE lease already owns + # liveness and heartbeat behavior. self._ka = client.keepalive(interval=120.0).start() program = Program(client) except Exception as e: # noqa: BLE001 @@ -3814,31 +5655,33 @@ class IdaTui(App): return self.client = client self.program = program - self.app.call_from_thread(self._status, f"{module} — loading functions…") + self._new_database = False + self.app.call_from_thread( + self._status, f"{module} [{client.backend}] — loading functions…") self._load_functions() - def _open_worker_client(self): # type: ignore[no-untyped-def] - """Our idalib-worker path: spawn the worker (it opens + analyzes the - binary in its own process) and connect. Returns the client, or None.""" - from .worker_client import WorkerClient - if self._pool is not None: # project mode: the pool owns the workers + def _open_database_client(self): # type: ignore[no-untyped-def] + """Attach through Code Mode, reusing a GUI or managed idalib database.""" + if self._pool is not None: # project mode: the pool owns the leases label = self._binary or self._project.refs[0].label client = self._pool.get(label, progress=lambda m: self.app.call_from_thread(self._status, m)) self._binary = label self._pool.set_active(label) self._open_path = self._project.by_label(label).staged + self._title = os.path.basename(self._open_path) return client if not self._open_path: self.app.call_from_thread( - self._status, "the worker backend needs a binary path") + self._status, "Code Mode needs a database or executable path") self.app.call_from_thread(self._dismiss_loading) return None base = os.path.basename(self._open_path) self.app.call_from_thread( - self._status, f"starting worker — initial auto-analysis of {base}…") - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + self._status, f"discovering Code Mode database for {base}…") + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args, + new_database=self._new_database) client.connect(progress=lambda m: self.app.call_from_thread( self._status, m)) return client @@ -3850,7 +5693,6 @@ class IdaTui(App): self._func_index = idx self.app.call_from_thread(lambda: self.query_one("#func-table", DataTable).clear()) last = 0 - module = self._module() while not idx.complete: idx.load_next_page() rows = idx.window(last, len(idx) - last) @@ -3858,18 +5700,20 @@ class IdaTui(App): if rows: self.app.call_from_thread(self._append_rows, rows) self.app.call_from_thread( - self._status, f"{module} — {last} functions…" + self._status, f"{last} functions…" ) # If a filter is active (typed during load), re-apply it over the full set. if self._filter_term: self.app.call_from_thread(self._apply_filter, self._filter_term) else: self.app.call_from_thread( - self._status, f"{module} — {len(idx)} functions (Ctrl+N: find symbol)") + self._status, f"{len(idx)} functions (Ctrl+N: find symbol)") # Land somewhere useful instead of an empty pane: main() if present, # otherwise pop the fuzzy symbol picker. self.app.call_from_thread(self._auto_land) self._index_binary() # project mode: keep the cross-binary index fresh + if self.trace_ctl.armed: + self._load_trace() # needs the index above: rebasing reads it @work(thread=True, exclusive=True, group="prewarm") def _prewarm_provider(self) -> None: @@ -3909,7 +5753,7 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="index") def _index_binary(self) -> None: """Fold this binary's symbols + strings into the project index, so it can - be searched later even when its worker is gone.""" + be searched later even when its Code Mode lease is gone.""" if self._index is None or self._project is None or self._binary is None: return ref = self._project.by_label(self._binary) @@ -3927,7 +5771,7 @@ class IdaTui(App): imps, exps = self.program.linkage() entries += [(KIND_IMPORT, i.addr, i.name) for i in imps] entries += [(KIND_EXPORT, e.addr, e.name) for e in exps] - except Exception: # noqa: BLE001 -- an old worker has no list_linkage + except Exception: # noqa: BLE001 -- indexing is best-effort pass try: n = self._index.reindex(self._binary, entries, source=ref.source) @@ -3969,7 +5813,19 @@ class IdaTui(App): if fn is not None: self._open_function(fn.addr, fn.name) elif len(self._func_index): - self.action_symbols() + # No main(): land on the first function rather than pushing the + # symbol palette. A modal as the *startup* state leaves a human + # staring at a picker over an empty pane, and silently swallows + # every keystroke an RPC driver injects while `ping` still says + # ready:true. Landing somewhere real is better for both; Ctrl+N is + # one keypress away. + first = self._func_index.get(0) + if first is not None: + self._open_function(first.addr, first.name) + self._status(f"no entry function — opened {first.name} " + "(Ctrl+N: find symbol)") + else: + self.action_symbols() else: self._land_without_functions() @@ -4000,7 +5856,13 @@ class IdaTui(App): cursor=0, push=True, is_region=True) def _can_reload(self) -> bool: - """Whether we're able to re-open this binary with different options.""" + """Whether Code Mode can replace this IDB with different options. + + A GUI database is owned by the user and has no remote close/rollback + route. Managed idalib databases can be released and reopened fresh. + """ + if self.client is not None and self.client.backend == "gui": + return False if self._project is not None and self._binary is not None: return True return bool(self._open_path) @@ -4080,7 +5942,7 @@ class IdaTui(App): if term: self._status(f"filter '{term}': {len(matched)}/{total}") else: - self._status(f"{self._module()} — {total} functions") + self._status(f"{total} functions") def _apply_pending_filter(self) -> None: self._filter_timer = None @@ -4115,8 +5977,7 @@ class IdaTui(App): left = self.query_one("#left", FunctionsPanel) left.display = not left.display if not left.display: - self.query_one(DecompView if self._active == "decomp" - else ListingView).focus() + self._focus_code_view() else: self.query_one("#func-table", DataTable).focus() @@ -4190,6 +6051,9 @@ class IdaTui(App): def _on_quit_choice(self, choice: str | None) -> None: if choice == "discard": + # Code Mode has no rollback/close-without-save operation. For GUI + # sessions this leaves changes dirty in IDA; a managed worker owns + # its final save policy and may persist them on final lease release. self._save_on_exit = False self.exit() elif choice == "save": @@ -4206,7 +6070,7 @@ class IdaTui(App): if self._pool is not None: self._pool.close_all(save=True) # saves each resident worker elif self.program is not None: - self.program.client.call("idb_save", timeout=600.0) + self.program.client.save_database() except Exception as e: # noqa: BLE001 -- still exit, but say so self.app.call_from_thread(self._status, f"save failed: {e}") self.app.call_from_thread(self._finish_exit) @@ -4246,7 +6110,7 @@ class IdaTui(App): self._ask_load_options(ref.source, label=label) return # Snapshot what we're leaving so coming back restores the view, then let - # the pool hand us a worker (spawning + evicting as the budget dictates). + # the pool hand us a lease (attaching + evicting as the budget dictates). if self._binary is not None: self._states[self._binary] = BinaryState( label=self._binary, program=self.program, @@ -4267,8 +6131,8 @@ class IdaTui(App): self.app.call_from_thread(self._switch_failed, label, str(e)) return st = self._states.get(label) - # The Program (and its caches) only survive while that worker does; a - # binary that was evicted comes back with a fresh one. Either way the nav + # The Program (and its caches) only survive while that lease does; an + # evicted binary reattaches. Either way the nav # history is just addresses, so it always survives. reuse = (st is not None and st.program is not None and getattr(st.program, "client", None) is client) @@ -4282,7 +6146,8 @@ class IdaTui(App): self._binary = label self._pool.set_active(label) self._open_path = self._project.by_label(label).staged - self._active = st.active if st else "listing" + self._title = os.path.basename(self._open_path) + self._active = st.active if st else ViewMode.LISTING self._split = st.split if st else False self._filter_term = st.filter_term if st else "" self._dirty = st.dirty if st else False @@ -4304,7 +6169,7 @@ class IdaTui(App): self._did_auto_land = False self._auto_land() return - # Cold (first visit, or the worker was evicted): rebuild the index, then + # Cold (first visit, or the lease was evicted): rebuild the index, then # land back where we were via _pending_restore. self._cur = None self._func_index = None @@ -4352,6 +6217,35 @@ class IdaTui(App): binary=self._binary), self._on_string_chosen) + def action_find(self) -> None: + """Ctrl+F: search the whole database — disassembly text, or bytes.""" + if self.program is None: + self._status("not connected yet") + return + if self._prompt_active(): + return + # Seed it with the word under the cursor: the search you want is + # usually about the thing you are looking at. + seed = "" + view = self._active_code_view() + if view is not None: + try: + seed = view.word_under_cursor() or "" + except Exception: # noqa: BLE001 -- a seed is a nicety, never a + seed = "" # reason not to open the search + self.push_screen(SearchPalette(self.program, seed), self._on_hit_chosen) + + def _on_hit_chosen(self, hit) -> None: # type: ignore[no-untyped-def] + if hit is None: + return + # Navigate to the ITEM, not the matched byte: a pattern can start in + # the middle of an instruction, and there is nothing to put a cursor on + # there. The status names the exact address so it isn't lost. + self._goto_ea(hit.head, push=True) + if hit.addr != hit.head: + self._status(f"match at {hit.addr:#x} (inside {hit.head:#x})", + priority=True) + def _on_string_chosen(self, choice) -> None: # type: ignore[no-untyped-def] if choice is None: return @@ -4361,28 +6255,6 @@ class IdaTui(App): return self._goto_ea(addr, push=True) # land on the literal in the listing - def on_descendant_focus(self, event) -> None: # type: ignore[no-untyped-def] - """Keep ``_active`` in step with focus while split. - - Tab moves both together, but focus also moves on its own — a click, or a - pane focusing itself after a load — and then ``_active`` still names the - pane you're NOT in. Everything downstream trusts ``_active``: follow - resolves the word under that pane's cursor and pushes history for it, so - Enter in the pseudocode would follow something from the listing and the - next Esc got spent undoing it. - """ - if not self._split: - return - w = self.focused - mode = ("decomp" if isinstance(w, DecompView) - else "listing" if isinstance(w, ListingView) else None) - if mode is None or mode == self._active: - return - self._active = mode - self._sync_split(mode) # re-link the band from the new driver - if not self.query_one(DecompView).loading: - self._status_for_cur("split") # never clobber "decompiling…" - def action_toggle_view(self) -> None: """Tab: switch the code pane between disassembly and pseudocode (or leave the hex view back to the preferred code view).""" @@ -4400,17 +6272,35 @@ class IdaTui(App): return if self._split: # In split mode Tab/F5 just moves focus between the two panes. - self._active = "decomp" if self._active == "listing" else "listing" - (self.query_one(DecompView) if self._active == "decomp" + self._active = ViewMode.DECOMP if self.is_listing else ViewMode.LISTING + (self.query_one(DecompView) if self.is_decomp else self.query_one(ListingView)).focus() self._sync_split(self._active) # re-link from the new driver self._status_for_cur("split") return - if self._active == "hex": + if self.is_hex: self._active = self._code_mode() self._show_active() return - if self._active == "listing": + if self.is_graph: + # F5/Tab out of the graph lands in the pseudocode at the block the + # cursor was on (Space is the key that returns to the listing). + gv = self.query_one(GraphView) + ea = gv._cursor_ea() + self._graph_sticky = False + if ea is None: + self._active = ViewMode.LISTING + self._show_active() + return + gv.display = False + dec = self.query_one(DecompView) + dec.display = True + dec.loading = True + dec.focus() + self._status("decompiling…") + self._decomp_from_listing(ea) + return + if self.is_listing: # F5/Tab in the continuous listing: decompile the function under the # cursor (IDA-style), if the cursor is inside a defined routine. ea = self.query_one(ListingView)._cursor_ea() @@ -4435,7 +6325,7 @@ class IdaTui(App): self._decomp_return = None if ret is not None: self._cur = ret - self._active = "listing" + self._active = ViewMode.LISTING self._open_entry(ret, push=False) elif self._cur is not None: # No F5 snapshot (we arrived via a decomp navigation): show THIS entry @@ -4446,7 +6336,7 @@ class IdaTui(App): ea = dec._line_ea(dec.cursor) self._toggle_to_listing(ea if ea is not None else self._cur.ea) else: - self._active = "listing" + self._active = ViewMode.LISTING self._show_active() @work(thread=True, group="nav") @@ -4459,14 +6349,14 @@ class IdaTui(App): def _apply_toggle_listing(self, idx: int) -> None: cur = self._cur if cur is None: - self._active = "listing" + self._active = ViewMode.LISTING self._show_active() return cur.view = "listing" cur.cursor = idx cur.cursor_x = 0 cur.scroll_y = -1 # derive a viewport (keeps the target in context) - self._active = "listing" + self._active = ViewMode.LISTING self._open_entry(cur, push=False) @work(thread=True, group="nav") @@ -4485,7 +6375,7 @@ class IdaTui(App): # We optimistically raised the pseudocode overlay; drop back to the # listing since there's nothing to decompile here. self.query_one(DecompView).loading = False - self._active = "listing" + self._active = ViewMode.LISTING self._show_active() self._status(msg) @@ -4501,7 +6391,7 @@ class IdaTui(App): entry = NavEntry(ea=fn_addr, name=fn_name, is_region=False, dec_cursor=max(dec_idx, 0)) self._cur = entry - self._active = "decomp" + self._active = ViewMode.DECOMP self._show_active() def action_toggle_split(self) -> None: @@ -4517,12 +6407,114 @@ class IdaTui(App): return self._split = not self._split if self._active not in ("listing", "decomp"): - self._active = "listing" + self._active = ViewMode.LISTING if self._split: self._enter_split(self._cur.ea, self._cur.name) else: self._show_active() + # -- graph mode -------------------------------------------------------- # + #: Above this, a CFG graph stops being a picture and becomes a hairball -- + #: IDA's own is unreadable at this size too. Refusing beats rendering soup. + GRAPH_MAX_BLOCKS = 400 + + def action_toggle_graph(self) -> None: + """Space: swap the code view for the function's control-flow graph.""" + if self._prompt_active(): + return + if self.is_graph: + self._graph_sticky = False + self._active = self._code_mode() + self._show_active() + return + if self._cur is None or self.program is None: + self._status("open a function first") + return + self._graph_sticky = True + gv = self.query_one(GraphView) + ea = self._graph_target_ea() + if gv.loaded_ea is not None and gv.fc is not None \ + and gv.fc.func_ea == self._cur.ea: + self._active = ViewMode.GRAPH + self._split = False + self._show_active() + if ea is not None: + gv.goto_ea(ea) + self._graph_status() + return + self._status(f"{self._cur.name} — building graph…") + self._load_graph(self._cur.ea, ea) + + def _graph_target_ea(self) -> int | None: + """The address the graph should land on: wherever the code view's cursor + is, so Space doesn't lose your place.""" + try: + if self.is_listing: + return self.query_one(ListingView)._cursor_ea() + if self.is_decomp: + dv = self.query_one(DecompView) + return dv._line_ea(dv.cursor) + except Exception: # noqa: BLE001 + pass + return self._cur.ea if self._cur else None + + @work(thread=True, exclusive=True, group="graph") + def _load_graph(self, func_ea: int, want_ea: int | None) -> None: + assert self.program is not None + err = "" + fc = None + try: + fc = self.program.flowchart(func_ea) + except IDAConnectionError: + raise + except Exception as e: # noqa: BLE001 + err = f"{type(e).__name__}: {e}" + self.app.call_from_thread(self._apply_graph, func_ea, want_ea, fc, err) + + def _apply_graph(self, func_ea: int, want_ea: int | None, fc, # type: ignore[no-untyped-def] + err: str) -> None: + if self._cur is None or self._cur.ea != func_ea: + return # a newer navigation won + if not self._graph_sticky and not self.is_graph: + # The load lost its race: the user has since left graph mode (or a + # rename's reload queued one behind their back). Forcing the view + # here drags them back into a graph they already dismissed. + return + if fc is None: + self._status(err or "no control-flow graph for this function " + "(is it a thunk or an import?)") + return + if len(fc.blocks) > self.GRAPH_MAX_BLOCKS: + self._status( + f"{fc.name}: {len(fc.blocks)} blocks — too many to graph " + f"(limit {self.GRAPH_MAX_BLOCKS}); staying in the listing") + return + gv = self.query_one(GraphView) + gv.set_graph(fc, want_ea) + self._active = ViewMode.GRAPH + self._split = False + self._show_active() + self._graph_status() + + def _graph_status(self) -> None: + gv = self.query_one(GraphView) + if gv.lay is None or gv.fc is None: + return + s = gv.lay.stats + loops = f", {s['back']} loop{'s' if s['back'] != 1 else ''}" if s["back"] else "" + self._status( + f"{gv.fc.name} @ {gv.fc.func_ea:#x} [graph: {s['blocks']} blocks, " + f"{s['edges']} edges{loops}] " + f"z=zoom({gv.ZOOMS[gv._zoom]}) m=map J/K=edge space=text") + + def on_graph_view_cursor_moved(self, msg: "GraphView.CursorMoved") -> None: + gv = self.query_one(GraphView) + gv.set_highlight(gv.word_under_cursor()) + if msg.ea is not None and gv.fc is not None: + b = gv.fc.block_at(msg.ea) + extra = f" block {b.start:#x}" if b else "" + self._status(f"{gv.fc.name} @ {msg.ea:#x}{extra} [graph]") + @work(thread=True, group="split") def _enter_split(self, ea: int, name: str) -> None: # Load the listing for the current function (bg) then reveal both panes. @@ -4546,17 +6538,19 @@ class IdaTui(App): cursor; press again (or Tab/Esc) to return to the code view.""" if self._cur is None or self.program is None: return - if self._active == "hex": + if self.is_hex: self._active = self._code_mode() self._show_active() return - if self._active in ("listing", "disasm"): + if self.is_listing: ea = self.query_one(ListingView)._cursor_ea() + elif self.is_graph: + ea = self.query_one(GraphView)._cursor_ea() else: dec = self.query_one(DecompView) ea = dec._line_ea(dec.cursor) self._hex_pending_ea = ea if ea is not None else self._cur.ea - self._active = "hex" + self._active = ViewMode.HEX self._show_active() def action_filter(self) -> None: @@ -4567,9 +6561,53 @@ class IdaTui(App): inp.value = self._filter_term inp.focus() + def action_export(self) -> None: + """Ctrl+E: write what you have worked out to a markdown report. + + Prefilled with a path beside the binary, because the common case is + "just write it" and the rare case is worth one edit. + """ + if self.program is None: + self._status("no database open", priority=True) + return + inp = self.query_one("#export", Input) + inp.placeholder = "export findings to… (markdown) — Enter" + inp.can_focus = True + inp.display = True + inp.value = findings.default_path(self._open_path or "findings") + self.query_one("#status", Static).display = False + inp.focus() + + def _end_export(self) -> None: + inp = self.query_one("#export", Input) + inp.display = False + inp.can_focus = False + self.query_one("#status", Static).display = True + + def export_findings(self, path: str | None = None) -> None: + """Gather + write the report off the UI thread (it walks the database).""" + self._status("exporting findings…", priority=True) + self._export_worker(path) + + @work(thread=True, exclusive=True, group="export") + def _export_worker(self, path: str | None) -> None: + try: + self.journal.load(self.program) + self.journal.flush(self.program) + out, f = findings.export(self.program, self._open_path or "", path, + journal=self.journal) + except Exception as e: # noqa: BLE001 -- a bad path is a message, not a crash + self.call_from_thread(self._status, f"export failed: {e}", True) + return + n_named = len(findings._user_names(f)) + self.call_from_thread( + self._status, + f"exported {len(f.comments)} comments, {n_named} names, " + f"{len(f.types)} types → {out}", True) + def action_goto(self) -> None: inp = self.query_one("#goto", Input) - inp.placeholder = ("hex goto: 0xADDR or name — Enter" if self._active == "hex" + inp.placeholder = ("hex goto: 0xADDR or name — Enter" if self.is_hex else "goto: name or 0xADDR — Enter") inp.can_focus = True inp.display = True @@ -4605,8 +6643,7 @@ class IdaTui(App): elif self.query_one("#left", FunctionsPanel).display: table.focus() else: - self.query_one(DecompView if self._active == "decomp" - else ListingView).focus() + self._focus_code_view() # -- input submit (filter / goto) ------------------------------------- # def on_search_requested(self, msg: SearchRequested) -> None: @@ -4624,38 +6661,68 @@ class IdaTui(App): def on_follow_requested(self, msg: FollowRequested) -> None: view = msg.view word = view.word_under_cursor() - if isinstance(view, DisasmView): + if isinstance(view, GraphView): ea = view._cursor_ea() - # The instruction's ordinary fall-through edge (to the next line) is - # an indistinguishable 'code' xref; pass it so follow can skip it and - # land on a call/jump's real target instead of the next instruction. - nxt = view.model.cached_line(view.cursor + 1) if view.model else None - if ea is not None: - self._follow_disasm(ea, word, nxt.ea if nxt else None) - elif isinstance(view, ListingView): + if ea is None: + return + # Inside the graph, a jump to a block of THIS function should move + # the cursor, not navigate away and rebuild the whole picture. + tgt = self._graph_local_target(view, ea, word) + if tgt is not None: + view.goto_ea(tgt) + self.post_message(GraphView.CursorMoved(tgt, view.cursor_node)) + return + self._follow_disasm(ea, word, view._next_ea()) + return + if isinstance(view, ListingView): ea = view._cursor_ea() if ea is not None: + # _next_ea() is the following item: follow uses it to skip the + # ordinary fall-through edge, which is an indistinguishable + # 'code' xref, and land on a call/jump's real target instead. self._follow_disasm(ea, word, view._next_ea()) elif isinstance(view, DecompView) and view._texts: self._follow_decomp(view._texts[view.cursor], word, view._line_ea(view.cursor)) + def _graph_local_target(self, view: "GraphView", ea: int, + word: str) -> int | None: + """If the cursor's instruction branches somewhere inside this same + graph, return that address.""" + if view.fc is None or self.program is None: + return None + try: + if word and self._looks_like_symbol(word): + t = self.program.resolve(word) + if t and view.fc.block_at(t) is not None: + return t + except Exception: # noqa: BLE001 + pass + try: + for xr in self.program.xrefs_from(ea): + t = xr.to + if t and t != view._next_ea() and view.fc.block_at(t) is not None: + return t + except Exception: # noqa: BLE001 + pass + return None + def on_xrefs_requested(self, msg: XrefsRequested) -> None: if self._xref_active: # one gather at a time; ignore a second 'x' return view = msg.view word = view.word_under_cursor() - if isinstance(view, DisasmView): + if isinstance(view, GraphView): ea = view._cursor_ea() if ea is not None: - # span = this instruction .. the next, to pre-select the dialog - # entry for the site we invoked xrefs from. - nxt = view.model.cached_line(view.cursor + 1) if view.model else None - self._push_busy("finding xrefs\u2026") - self._xrefs_disasm(ea, word, ea, nxt.ea if nxt else None) - elif isinstance(view, ListingView): + self._push_busy("finding xrefs…") + self._xrefs_disasm(ea, word, ea, view._next_ea()) + return + if isinstance(view, ListingView): ea = view._cursor_ea() if ea is not None: + # span = this instruction .. the next, to pre-select the dialog + # entry for the site we invoked xrefs from. self._push_busy("finding xrefs\u2026") self._xrefs_disasm(ea, word, ea, view._next_ea()) elif isinstance(view, DecompView) and view._texts: @@ -4740,7 +6807,7 @@ class IdaTui(App): def _cross_binary_impl(self, name: str) -> tuple[str, int] | None: """``(binary, addr)`` of a project binary that EXPORTS ``name``. - Reads the on-disk index, so a provider resolves even when its worker was + Reads the on-disk index, so a provider resolves even when its lease was evicted — the whole reason the index exists. """ if self._index is None or self._project is None or not name: @@ -4913,7 +6980,7 @@ class IdaTui(App): def _foreign_importers(self, subj: int, subj_name, fn): # type: ignore[no-untyped-def] """Project binaries that IMPORT the symbol at ``subj`` — the other half of the phase-3 join, read from the on-disk index so a caller shows up - whether or not its worker is resident. + whether or not its database lease is resident. Only for a symbol this binary actually exports: a local name that happens to collide with another binary's import isn't a caller of ours. @@ -4963,556 +7030,171 @@ class IdaTui(App): # If xrefs was invoked from the decompiler, land the jump back in the # decompiler (when the target is decompilable) rather than the listing. self._goto_ea(addr, push=True, focus_name=self._xref_focus_name, - prefer_decomp=(self._active == "decomp")) + prefer_decomp=(self.is_decomp)) - # -- rename ------------------------------------------------------------ # - @staticmethod - def _is_pseudocode_label(view, name: str) -> bool: - """True if ``name`` is a Hex-Rays goto label in ``view``. The rename tool - has no label category (only func/global/local/stack), so renaming one - fails with a misleading 'local variable not found'; detect it up front - and explain instead. A label is the default ``LABEL_n`` or any token used - as a ``goto`` target.""" - if not isinstance(view, DecompView): - return False - if re.fullmatch(r"LABEL_\d+", name): - return True - body = "\n".join(getattr(view, "_texts", []) or []) - return re.search(rf"\bgoto\s+{re.escape(name)}\b", body) is not None + # -- database edits ---------------------------------------------------- # + # The bodies live in EditController (idatui/edit_ctl.py). What stays here is + # what Textual insists on owning: on_<Message> handlers, which it dispatches + # by name on the DOMNode, and @work entry points, whose worker machinery + # wants a DOMNode host. Both are one-line delegates. def on_rename_requested(self, msg: RenameRequested) -> None: - # In the flat listing, 'n' names the ADDRESS under the cursor (create a - # label), not a symbol-by-name. This is what lets you name a bare/ - # undefined byte — e.g. the free byte at addr+1 after shrinking a u16 to - # a u8 — which the word-under-cursor path can't do (no symbol to rename). - if isinstance(msg.view, ListingView): - ea = msg.view._cursor_ea() - if ea is None: - self._status("no address on this line to name") - return - head = msg.view.cur_head() - word = msg.view.word_under_cursor() - mnem = head.text.split(" ", 1)[0] if (head and head.text) else "" - # If the cursor is on a symbol token (a call/branch target, a data - # reference, or this head's own label) rename THAT symbol; otherwise - # create/rename a label at the head's address (bare/undefined bytes). - if (word and self._looks_like_symbol(word) and word != mnem - and word.lower() not in _ASM_KEYWORDS): - self._rename_ctx = (msg.view, word) - self._rename_addr = None - placeholder = f"rename '{word}' — Enter=apply Esc=cancel" - prefill = word - else: - cur = head.name if (head is not None and head.name) else "" - self._rename_ctx = (msg.view, cur) - self._rename_addr = ea - placeholder = f"name @ {ea:#x} — Enter=apply Esc=cancel" - prefill = cur - self.query_one("#status", Static).display = False - inp = self.query_one("#rename", Input) - inp.placeholder = placeholder - inp.can_focus = True - inp.display = True - inp.value = prefill - inp.focus() - return - if not msg.name: - self._status("nothing to rename under the cursor") - return - if self._is_pseudocode_label(msg.view, msg.name): - self._status( - f"can't rename pseudocode label '{msg.name}' " - "(Hex-Rays goto labels aren't renamable via the API)") + self.edits.request_rename(msg) + + def on_comment_requested(self, msg: CommentRequested) -> None: + self.edits.request_comment(msg) + + def on_retype_requested(self, msg: RetypeRequested) -> None: + if self._cur is None or self.program is None: return - self._rename_ctx = (msg.view, msg.name) - self._rename_addr = None - self.query_one("#status", Static).display = False - inp = self.query_one("#rename", Input) - inp.placeholder = f"rename '{msg.name}' — Enter=apply Esc=cancel" - inp.can_focus = True - inp.display = True - inp.value = msg.name - inp.focus() + self._prepare_retype(msg.view, msg.name) - def _end_rename(self) -> None: - inp = self.query_one("#rename", Input) - inp.display = False - inp.can_focus = False - self._rename_addr = None - self.query_one("#status", Static).display = True - view, _ = self._rename_ctx - if view is not None: - view.focus() + def on_make_data_requested(self, msg: MakeDataRequested) -> None: + self.edits.request_make_data(msg) - # -- comments ---------------------------------------------------------- # - def _line_ea_for(self, view) -> int | None: # type: ignore[no-untyped-def] - """Address of the line under the cursor in either code view.""" - if isinstance(view, (DisasmView, ListingView)): - return view._cursor_ea() - if isinstance(view, DecompView): - return view._line_ea(view.cursor) - return None + def on_op_format_requested(self, msg: OpFormatRequested) -> None: + self.edits.request_op_format(msg) - @staticmethod - def _existing_comment(view) -> str: - """Current line comment (for prefill), parsed from the rendered text. In - pseudocode a comment is `// text` before the trailing /*0xEA*/ markers; - C has no `//` operator, so the last `//` is unambiguously the comment.""" - if isinstance(view, DecompView) and 0 <= view.cursor < len(view._texts): - s = re.sub(r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", view._texts[view.cursor]) - i = s.rfind("//") - return s[i + 2:].strip() if i >= 0 else "" - return "" + def on_edit_item_requested(self, msg: EditItemRequested) -> None: + self.edits.request_edit_item(msg) - def on_comment_requested(self, msg: CommentRequested) -> None: - ea = self._line_ea_for(msg.view) - # Signature / local-declaration lines carry no address; fall back to the - # function's entry ea so commenting the header annotates the function. - func_level = ea is None - if func_level: - ea = self._cur.ea if self._cur else None - if ea is None: - self._status("no address on this line to comment") - return - existing = "" if func_level else self._existing_comment(msg.view) - self._comment_ctx = (msg.view, ea, existing) - self.query_one("#status", Static).display = False - inp = self.query_one("#comment", Input) - inp.placeholder = ( - f"function comment @ {ea:#x} — Enter=apply (empty=clear) Esc=cancel" - if func_level else - f"comment @ {ea:#x} — Enter=apply (empty=clear) Esc=cancel") - inp.can_focus = True - inp.display = True - inp.value = existing - inp.focus() + @work(thread=True, exclusive=True, group="rename") + def _do_rename(self, view, old: str, new: str) -> None: # type: ignore[no-untyped-def] + self.edits.do_rename(view, old, new) - def _end_comment(self) -> None: - inp = self.query_one("#comment", Input) - inp.display = False - inp.can_focus = False - self.query_one("#status", Static).display = True - view, _, _ = self._comment_ctx - if view is not None: - view.focus() + @work(thread=True, exclusive=True, group="rename") + def _do_name_addr(self, addr: int, name: str) -> None: + self.edits.do_name_addr(addr, name) @work(thread=True, exclusive=True, group="comment") def _do_comment(self, ea: int, text: str) -> None: - assert self.program is not None - # The prompt is single-line, so a literal '\n' (backslash-n) means a real - # newline — Hex-Rays renders each as its own '//' line. Lets long notes - # wrap instead of running off the right edge and clipping. - text = text.replace("\\n", "\n") - try: - res = self.program.set_comment(ea, text) - except IDAToolError as e: - self.app.call_from_thread(self._status, f"comment failed: {e.message}") - return - data = res.get("result") if isinstance(res, dict) else None - if isinstance(data, list) and data and isinstance(data[0], dict) and data[0].get("error"): - self.app.call_from_thread(self._status, f"comment failed: {data[0]['error']}") - return - self.app.call_from_thread(self._after_comment, ea, text) + self.edits.do_comment(ea, text) + + @work(thread=True, exclusive=True, group="retype") + def _prepare_retype(self, view, word: str | None) -> None: # type: ignore[no-untyped-def] + self.edits.prepare_retype(view, word) + + @work(thread=True, exclusive=True, group="retype-apply") + def _do_retype(self, kind: str, subject: int, word: str, new: str) -> None: + self.edits.do_retype(kind, subject, word, new) + + @work(thread=True, exclusive=True, group="makedata") + def _do_make_data(self, ea: int, type_decl: str, + anchor: ViewAnchor | None = None) -> None: + self.edits.do_make_data(ea, type_decl, anchor) + + @work(thread=True, exclusive=True, group="opformat") + def _do_op_format(self, mode: str, where: str, ea: int, col: int, + line: int = -1) -> None: + self.edits.do_op_format(mode, where, ea, col, line) + + @work(thread=True, exclusive=True, group="edititem") + def _do_edit_item(self, kind: str, ea: int, + anchor: ViewAnchor | None = None) -> None: + self.edits.do_edit_item(kind, ea, anchor) def _reload_active_code(self) -> None: - """Refresh whichever code view is showing after an edit (comment/rename/ - retype), in place: re-decompile if in the decompiler, else reload the - listing.""" - cur = self._cur - if cur is None: - return - if self._active == "decomp": - # Snapshot the LIVE pseudocode position before forcing a recompile. - # dec_scroll_y isn't tracked on every move, so without this the reload - # falls into show()'s derive path (a bare scroll_to) and leaves a - # stale frame until the next cursor move; capturing the real scroll - # makes show() take the robust _apply_scroll path and repaint now. - dec = self.query_one(DecompView) - if cur.ea == dec.loaded_ea: - cur.dec_cursor = dec.cursor - cur.dec_cursor_x = dec.cursor_x - cur.dec_scroll_y = round(dec.scroll_offset.y) - cur.dec_scroll_x = round(dec.scroll_offset.x) - dec.loaded_ea = None # force re-decompile - self._show_active() - else: - # Capture the LIVE position from the widget (the source of truth) - # rather than trusting nav-entry tracking, which goes stale. Capture - # it as ADDRESSES via the anchor: bump_names() discards the segment - # model so the reload rebuilds it, and an edit that changes how many - # rows an item takes makes the old indices point somewhere else. - # Index capture is CORRECT here and an anchor is not: a rename or - # comment doesn't change how many rows anything takes, and the model - # this rebuilds is constructed empty — index_of_ea on it returns -1 - # until pages load, so an anchor would resolve to nothing while - # costing an extra model build on the UI thread. Address anchoring is - # for the edit paths that DO change row structure (see _do_edit_item). - lst = self.query_one(ListingView) - cur.view = "listing" - if lst.model is not None: - cur.cursor = lst.cursor - cur.cursor_x = lst.cursor_x - cur.scroll_y = round(lst.scroll_offset.y) - self._open_entry(cur, push=False) + self.edits.reload_active_code() - def _after_comment(self, ea: int, text: str) -> None: - # A comment shows in both views but only after Hex-Rays recompiles, so - # reuse the name-generation invalidation (bumps gen -> decompile is - # force_recompiled lazily; disasm/listing caches are cleared). - self.program.bump_names() - self._reload_active_code() - self._dirty = True - verb = "cleared comment" if not text else "commented" - self._status(f"{verb} @ {ea:#x} (Ctrl+S to save)") + # -- trace ------------------------------------------------------------- # + # Everything below delegates to TraceController (idatui/trace_ctl.py). The + # keys stay declared here because Textual only merges BINDINGS from DOMNode + # subclasses, and the @work entry points stay here because Textual's worker + # machinery wants a DOMNode host. - # -- retype (set type, IDA 'y') --------------------------------------- # - def on_retype_requested(self, msg: RetypeRequested) -> None: - if self._cur is None or self.program is None: - return - self._prepare_retype(msg.view, msg.name) + # Read-only views onto the controller's state. The pilot suite and the RPC + # layer read the trace position by these names; they are properties rather + # than attributes so there is exactly one owner and no copy to drift. + @property + def _trace(self): # type: ignore[no-untyped-def] + return self.trace_ctl.trace - @work(thread=True, exclusive=True, group="retype") - def _prepare_retype(self, view, word: str | None) -> None: # type: ignore[no-untyped-def] - """Work out whether the cursor is on a local variable or a function, and - fetch the current type/prototype to prefill the prompt.""" - assert self.program is not None and self._cur is not None - ft = self.program.func_types(self._cur.ea) - kind: str | None = None - subject: int = self._cur.ea - prefill = "" - # 1) a local variable (or arg) of the current function - if word and ft is not None: - lv = next((v for v in ft.lvars if v.name == word), None) - if lv is not None: - kind, prefill = "lvar", lv.type - # 2) a symbol under the cursor: a function (retype its prototype) or a - # global/data item (retype the variable). Without the data case a - # global fell through to (3) and silently retyped the ENCLOSING - # function's prototype instead. - if kind is None and self._looks_like_symbol(word): - try: - tgt = self.program.resolve(word) - except Exception: # noqa: BLE001 - tgt = None - if tgt is not None: - tft = self.program.func_types(tgt) - if tft is not None: - kind, subject, prefill = "func", tgt, tft.prototype - else: - dt = self.program.data_type(tgt) - if dt is not None and not dt.get("is_func"): - kind, subject = "data", tgt - prefill = dt.get("type") or self._guess_data_type( - dt.get("size") or 0) - # 3) fall back to the current function itself - if kind is None and ft is not None: - kind, subject, prefill = "func", self._cur.ea, ft.prototype - if kind is None: - self.app.call_from_thread(self._status, "nothing to retype under the cursor") - return - self.app.call_from_thread(self._open_retype, view, kind, subject, word or "", prefill) + @property + def _t(self) -> int: + return self.trace_ctl.t - @staticmethod - def _guess_data_type(size: int) -> str: - """A sensible prefill when a global carries no type yet.""" - return {1: "unsigned __int8", 2: "unsigned __int16", - 4: "unsigned __int32", 8: "unsigned __int64"}.get( - size, f"char[{size}]" if size > 0 else "void *") + @property + def _trail_map(self) -> list: + return self.trace_ctl.trail_map - def _open_retype(self, view, kind: str, subject: int, word: str, # type: ignore[no-untyped-def] - prefill: str) -> None: - self._retype_ctx = (view, kind, subject, word) - self.query_one("#status", Static).display = False - inp = self.query_one("#retype", Input) - label = "prototype" if kind == "func" else f"type for '{word}'" - inp.placeholder = f"{label} — Enter=apply Esc=cancel" - inp.can_focus = True - inp.display = True - inp.value = prefill - inp.focus() + @property + def _trail_map_ea(self): # type: ignore[no-untyped-def] + return self.trace_ctl.trail_map_ea - def _end_retype(self) -> None: - inp = self.query_one("#retype", Input) - inp.display = False - inp.can_focus = False - self.query_one("#status", Static).display = True - view = self._retype_ctx[0] - if view is not None: - view.focus() + @property + def _trail_line_of(self) -> dict: + return self.trace_ctl.trail_line_of - @work(thread=True, exclusive=True, group="retype-apply") - def _do_retype(self, kind: str, subject: int, word: str, new: str) -> None: - assert self.program is not None - if kind == "func": - err = self.program.set_function_type(subject, new) - elif kind == "data": # a global / data item referenced in the body - err = self.program.set_data_type(subject, new) - else: # lvar of the current function - err = self.program.set_lvar_type(self._cur.ea, word, new) - if err: - self.app.call_from_thread(self._status, f"retype failed: {err}") - return - self.app.call_from_thread(self._after_retype, kind, word) + @work(thread=True, exclusive=True, group="load-trace") + def _load_trace(self) -> None: + self.trace_ctl.load() - def _after_retype(self, kind: str, word: str) -> None: - # A type change alters the pseudocode (and disasm operand types), so - # recompile via the name-generation invalidation and reopen in place. - self.program.bump_names() - self._reload_active_code() - self._dirty = True - what = "prototype" if kind == "func" else f"'{word}'" - self._status(f"retyped {what} (Ctrl+S to save)") + def _seek(self, idx: int, follow: bool = True) -> None: + self.trace_ctl.seek(idx, follow) - # -- typed data definition (make_data, IDA 'd') ----------------------- # - @staticmethod - def _default_data_type(head) -> str: # type: ignore[no-untyped-def] - """A sensible prefill C type for defining data over ``head``.""" - sz = getattr(head, "size", 0) or 0 - return {1: "unsigned __int8", 2: "unsigned __int16", - 4: "unsigned __int32", 8: "unsigned __int64"}.get( - sz, f"char[{sz}]" if sz > 0 else "unsigned __int8") + def _step(self, delta: int) -> None: + self.trace_ctl.step(delta) - def on_make_data_requested(self, msg: MakeDataRequested) -> None: - view = msg.view - ea = view._cursor_ea() if isinstance(view, ListingView) else None - if ea is None: - self._status("no address on this line to define data") - return - self._makedata_ctx = (view, ea) - self.query_one("#status", Static).display = False - inp = self.query_one("#makedata", Input) - inp.placeholder = (f"data type @ {ea:#x} (e.g. int, char[16], my_struct)" - " — Enter=apply Esc=cancel") - inp.can_focus = True - inp.display = True - head = view.cur_head() if isinstance(view, ListingView) else None - inp.value = self._default_data_type(head) if head is not None else "int" - inp.focus() + def _step_over(self, direction: int) -> None: + self.trace_ctl.step_over(direction) - def _end_makedata(self) -> None: - inp = self.query_one("#makedata", Input) - inp.display = False - inp.can_focus = False - self.query_one("#status", Static).display = True - view = self._makedata_ctx[0] - if view is not None: - view.focus() + def _paint_trail(self) -> None: + self.trace_ctl.paint_trail() - @work(thread=True, exclusive=True, group="makedata") - def _do_make_data(self, ea: int, type_decl: str, - anchor: ViewAnchor | None = None) -> None: # worker context - assert self.program is not None - try: - self.program.make_data(ea, type_decl) - except Exception as e: # noqa: BLE001 - self.app.call_from_thread(self._status, f"make data: {e}") - return - self.program.bump_items() - anchor = anchor or ViewAnchor() - anchor.flash = f"data ({type_decl}) @ {ea:#x} (Ctrl+S to save)" - name = self.program.region_label(ea) - lm = self.program.listing(ea) - idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 - _cur, top = self._anchor_rows(anchor, lm, ea) - self.app.call_from_thread( - self._open_at, ea, name, idx, False, -1, 0, True, None, top) - self.app.call_from_thread(self._edit_done, anchor) + @work(thread=True, exclusive=True, group="split-resync") + def _resync_decomp_async(self, ea: int) -> None: + self._resync_decomp(ea) - @work(thread=True, exclusive=True, group="rename") - def _do_rename(self, view, old: str, new: str) -> None: # type: ignore[no-untyped-def] - assert self.program is not None - prog, cur = self.program, self._cur - kind = "data" - addr: int | None = None - batch: dict = {"data": {"old": old, "new": new}} - resolved: int | None = None - try: - resolved = prog.resolve(old) - except Exception: # noqa: BLE001 - resolved = None - if resolved is not None: - fn = prog.function_of(resolved) - if fn is not None and fn.addr == resolved: - kind, addr = "func", resolved - batch = {"func": {"addr": hex(resolved), "name": new}} - else: - kind, batch = "data", {"data": {"old": old, "new": new}} - elif isinstance(view, DecompView) and cur is not None: - dec = prog.decompile(cur.ea) - ref = next((r for r in dec.refs if r.name == old), None) - if ref is not None: - fn = prog.function_of(ref.addr) - if fn is not None and fn.addr == ref.addr: - kind, addr = "func", ref.addr - batch = {"func": {"addr": hex(ref.addr), "name": new}} - else: - kind, batch = "data", {"data": {"old": old, "new": new}} - else: - kind = "local" - batch = {"local": {"func_addr": hex(cur.ea), "old": old, "new": new}} - elif cur is not None: # disasm view - if old.startswith(("var_", "arg_")): - kind = "stack" - batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} - try: - res = prog.client.call("rename", batch=batch) - except IDAToolError as e: - self.app.call_from_thread(self._status, f"rename failed: {e.message}") - return - summary = res.get("summary", {}) if isinstance(res, dict) else {} - if not (summary.get("ok", 0) > 0 and summary.get("failed", 0) == 0): - msg = "rename failed" - for catk in ("func", "data", "local", "stack"): - items = res.get(catk) if isinstance(res, dict) else None - if isinstance(items, list) and items and items[0].get("error"): - msg = f"rename failed: {items[0]['error']}" - self.app.call_from_thread(self._status, msg) - return - self.app.call_from_thread(self._after_rename, kind, addr, old, new) + def action_seek_next_hit(self) -> None: + self.trace_ctl.seek_hit(1) - def _after_rename(self, kind: str, addr: int | None, old: str, new: str) -> None: - cur = self._cur - # A renamed symbol can appear in many functions, so invalidate globally; - # each function refreshes its names the next time it's viewed. - self.program.bump_names() - self._reload_active_code() - if kind == "func" and addr is not None and self._func_index is not None: - self._func_index.update_name(addr, new) - for e in self._nav: - if e.ea == addr: - e.name = new - # Update the one cell in place (a full rebuild would race the - # initial streaming load and duplicate row keys). - table = self.query_one("#func-table", DataTable) - try: - name_col = list(table.columns.keys())[1] - table.update_cell(str(addr), name_col, new) - except Exception: # noqa: BLE001 -- row filtered out / not yet streamed - pass - self._dirty = True - self._status(f"renamed {old} → {new} (Ctrl+S to save)") + def action_seek_prev_hit(self) -> None: + self.trace_ctl.seek_hit(-1) - @work(thread=True, exclusive=True, group="rename") - def _do_name_addr(self, addr: int, name: str) -> None: # worker context - """Set a label at ``addr`` (listing 'n'). Works on a bare/undefined byte - — unlike the symbol-by-name path, this names the address directly.""" - assert self.program is not None - try: - res = self.program.client.call( - "rename", batch={"data": {"addr": hex(addr), "new": name}}) - except IDAToolError as e: - self.app.call_from_thread(self._status, f"name failed: {e.message}") - return - summary = res.get("summary", {}) if isinstance(res, dict) else {} - if not (summary.get("ok", 0) > 0 and summary.get("failed", 0) == 0): - err = "name failed" - items = res.get("data") if isinstance(res, dict) else None - if isinstance(items, list) and items and items[0].get("error"): - err = f"name failed: {items[0]['error']}" - self.app.call_from_thread(self._status, err) - return - # The label shows in the listing's head rows -> invalidate + reopen. - self.program.bump_items() - lm = self.program.listing(addr) - label = self.program.region_label(addr) - idx = max(lm.ensure_ea(addr), 0) if lm is not None else 0 - self.app.call_from_thread(self._open_at_named, label, addr, idx, name) + def action_seek_reg_write(self) -> None: + self.trace_ctl.seek_reg_write() - def _open_at_named(self, label: str, addr: int, idx: int, name: str) -> None: - self._open_at(addr, label, idx, False, -1, 0, True) - self._dirty = True - self._status(f"named {addr:#x} → {name} (Ctrl+S to save)") + def action_step_fwd(self) -> None: + self.trace_ctl.step(1) - # -- item structure edits (IDA c/p/u) --------------------------------- # - def on_edit_item_requested(self, msg: EditItemRequested) -> None: - view = msg.view - ea = (view._cursor_ea() - if isinstance(view, (DisasmView, ListingView)) else None) - if ea is None: - self._status("no address on this line to (re)define") - return - self._do_edit_item(msg.kind, ea, self._anchor()) + def action_step_back(self) -> None: + self.trace_ctl.step(-1) - @work(thread=True, exclusive=True, group="edititem") - def _do_edit_item(self, kind: str, ea: int, - anchor: ViewAnchor | None = None) -> None: # worker context - assert self.program is not None - verb = {"code": "defined code", "func": "created function", - "undef": "undefined", "string": "made string"}[kind] - try: - if kind == "code": - # Keep going until something stops it: one instruction is rarely - # what you want, and on a raw image it means pressing `c` once - # per opcode for the length of a function. - r = self.program.define_code_run(ea) - n, why = int(r.get("count", 0)), r.get("stopped", "") - if n == 0 and why == "defined": - # Already code/data here — a no-op, not a failure. Saying - # "failed to create instruction" for it would be a lie. - self.app.call_from_thread( - self._status, f"already defined @ {ea:#x}") - return - if n == 0: - raise IDAToolError("define_code", - f"@ {ea:#x}: Failed to create instruction") - end = int(str(r.get("end", hex(ea))), 0) - reason = {"undecodable": "hit bytes that don't decode", - "flow": "control flow ends here", - "defined": "ran into existing code/data", - "segment": "end of segment", - "limit": "instruction limit"}.get(why, why) - verb = (f"defined {n} instruction{'s' if n != 1 else ''} " - f"({ea:#x}\u2013{end:#x}) \u2014 {reason}") - elif kind == "func": - self.program.define_func(ea) - elif kind == "string": - s = self.program.make_string(ea) - verb = f"made string ({s[:24]!r})" if s else verb - else: - self.program.undefine(ea) - except Exception as e: # noqa: BLE001 -- surface soft/hard tool errors - self.app.call_from_thread(self._status, f"{kind}: {e}") - return - # Structure changed everywhere: drop all item/function/decomp caches. - self.program.bump_items() - # Re-resolve: a define_func upgrades the region to a real function view; - # anything else re-reads the (still function-less) listing in place. - anchor = anchor or ViewAnchor() - anchor.flash = f"{verb} @ {ea:#x} (Ctrl+S to save)" - fn = self.program.function_of(ea) - if fn is not None: - model = self.program.disasm(fn.addr, fn.name) - idx = 0 if ea == fn.addr else model.index_of_ea(ea) - _cur, top = self._anchor_rows(anchor, model, ea) - self.app.call_from_thread( - self._open_at, fn.addr, fn.name, idx, False, -1, 0, False, - None, top) - else: - name = self.program.region_label(ea) - lm = self.program.listing(ea) - idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 - _cur, top = self._anchor_rows(anchor, lm, ea) - self.app.call_from_thread( - self._open_at, ea, name, idx, False, -1, 0, True, None, top) - self.app.call_from_thread(self._edit_done, anchor) + def action_step_over_fwd(self) -> None: + self.trace_ctl.step_over(1) + + def action_step_over_back(self) -> None: + self.trace_ctl.step_over(-1) - def _edit_done(self, anchor: ViewAnchor) -> None: - """One place where an edit's aftermath is settled. + @work(thread=True, exclusive=True, group="load-funcs") + def _reindex_functions(self) -> None: + """Rebuild the function index in place after an edit changed it. - The reload this edit triggered will write its own status when it lands — - after this — so the message is handed over as a flash rather than - written and lost. + Deliberately not _load_functions(): that one is the BOOT path — it + clears the table, streams progress and then auto-lands, which would + yank the view away from the function you just made. """ - self._dirty = True - self._flash = anchor.flash - if anchor.flash: - self._status(anchor.flash) + if self.program is None: + return + idx = self.program.functions() + idx.load_all() + self._func_index = idx + self.app.call_from_thread(self._after_reindex) + + def _after_reindex(self) -> None: + idx = self._func_index + if idx is None: + return + if len(idx): + self._no_functions = False + self._apply_filter(self._filter_term) # repopulate the names pane @work(thread=True, exclusive=True, group="save") def _save(self) -> None: assert self.program is not None try: - self.program.client.call("idb_save", timeout=300.0) + self.journal.flush(self.program) # ride along into the .i64 + self.program.client.save_database() except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"save failed: {e}") return @@ -5567,7 +7249,8 @@ class IdaTui(App): idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 name = fn.name if fn is not None else self.program.region_label(ea) self.app.call_from_thread( - self._open_at, ea, name, idx, push, -1, 0, fn is None, focus_name) + self._open_at_if_current, seq, ea, name, idx, push, fn is None, + focus_name) def _open_decomp_entry(self, fn_addr: int, fn_name: str, dec_idx: int, dec_cursor_x: int, push: bool, @@ -5585,7 +7268,7 @@ class IdaTui(App): # Snapshot where we jumped from so 'back' returns there. If that was # the pseudocode (F5 makes a transient _cur not yet on the stack), # record its position and push it as a decomp entry. - if self._active == "decomp" and self._cur is not None: + if self.is_decomp and self._cur is not None: dv = self.query_one(DecompView) src = self._cur src.view = "decomp" @@ -5740,6 +7423,18 @@ class IdaTui(App): cur = row_of(fallback_ea) return (cur, row_of(a.top_ea)) + def _open_at_if_current(self, seq: int, ea: int, name: str, cursor: int, + push: bool, is_region: bool, + focus_name: str | None) -> None: + """Apply a navigation result only if it's still the one being awaited. + + The decompiler path has had this since 756589a; the listing path hadn't, + so a slow navigation could still land on top of a newer one. + """ + if seq != self._nav_seq: + return + self._open_at(ea, name, cursor, push, -1, 0, is_region, focus_name) + def _open_at(self, ea: int, name: str, cursor: int, push: bool, dec_cursor: int = -1, dec_cursor_x: int = 0, is_region: bool = False, focus_name: str | None = None, @@ -5784,39 +7479,26 @@ class IdaTui(App): view.focus() def on_key(self, event) -> None: # type: ignore[no-untyped-def] + # Any keypress means the result of the last edit has been read. + self._flash = None if event.key != "escape": return - if self.query_one("#search", Input).display: - event.stop() - event.prevent_default() - self._end_search(cancel=True) - return - if self.query_one("#rename", Input).display: - event.stop() - event.prevent_default() - self._end_rename() - return - if self.query_one("#comment", Input).display: - event.stop() - event.prevent_default() - self._end_comment() - return - if self.query_one("#retype", Input).display: - event.stop() - event.prevent_default() - self._end_retype() - return - if self.query_one("#makedata", Input).display: + # Esc closes whichever prompt is up. This used to be one copy of these + # four lines per prompt, which is how they drifted apart -- search + # cancelled its highlight, goto restored focus, the edit prompts did + # neither consistently. + prompt = self.prompts.active() + if prompt is not None: event.stop() event.prevent_default() - self._end_makedata() - return - if self.query_one("#goto", Input).display: - event.stop() - event.prevent_default() - self._end_goto() - (self.query_one(HexView) if self._active == "hex" - else self._code_view()).focus() + if prompt.id == "search": + self._end_search(cancel=True) + elif prompt.id in ("goto", "export"): + self._end_goto() if prompt.id == "goto" else self._end_export() + (self.query_one(HexView) if self.is_hex + else (self._code_view() or self.query_one(ListingView))).focus() + else: + prompt.close() return fi = self.query_one("#func-filter", Input) if fi.display: @@ -5841,42 +7523,32 @@ class IdaTui(App): view.repeat_last(direction) self._end_search() return - if inp.id == "rename": - view, old = self._rename_ctx - addr = self._rename_addr # capture before _end_rename clears it - self._end_rename() - if addr is not None: # listing: name this address (create a label) - if value and value != old: - self._do_name_addr(addr, value) - return - if view is not None and value and value != old: - self._do_rename(view, old, value) - return - if inp.id == "comment": - view, ea, existing = self._comment_ctx - self._end_comment() - if view is not None and value != existing: # empty value clears it - self._do_comment(ea, value) - return - if inp.id == "retype": - view, kind, subject, word = self._retype_ctx - self._end_retype() - if view is not None and value: - self._do_retype(kind, subject, word, value) - return - if inp.id == "makedata": - view, ea = self._makedata_ctx - self._end_makedata() - if view is not None and value: - self._do_make_data(ea, value, self._anchor()) + # The edit prompts all submit the same way: take the context the prompt + # was holding (close() hands it over, so it can't be read twice or go + # stale) and let the controller decide what to do with it. + submit = {"rename": self.edits.submit_rename, + "comment": self.edits.submit_comment, + "retype": self.edits.submit_retype, + "makedata": self.edits.submit_make_data}.get(inp.id or "") + if submit is not None: + ctx = self.prompts[inp.id].close() + if ctx is not None: + submit(ctx, value) return if inp.id == "goto": self._end_goto() - (self.query_one(HexView) if self._active == "hex" - else self._code_view()).focus() + (self.query_one(HexView) if self.is_hex + else (self._code_view() or self.query_one(ListingView))).focus() if value: self._goto(value) return + if inp.id == "export": + self._end_export() + (self.query_one(HexView) if self.is_hex + else (self._code_view() or self.query_one(ListingView))).focus() + if value: + self.export_findings(value) + return # filter mode: already applied incrementally; Enter just confirms + closes. self._apply_filter(value) self.query_one("#func-table", DataTable).focus() @@ -5889,15 +7561,63 @@ class IdaTui(App): """ return self._active_code_view() + def _line_ea_for(self, view) -> int | None: # type: ignore[no-untyped-def] + """Address of the line under the cursor in either code view. + + Shared, not an edit helper: the comment prompt, the goto prompt's + readback and rpc.py's `where` all ask the same question. + """ + if isinstance(view, ListingView): + return view._cursor_ea() + if isinstance(view, DecompView): + return view._line_ea(view.cursor) + return None + + # Read _active through these rather than comparing strings. The bare + # comparisons are what let the old "disasm" value take the wrong branch in + # five places, and they are what the next new mode would have to hunt down. + @property + def is_listing(self) -> bool: + return self._active == ViewMode.LISTING + + @property + def is_decomp(self) -> bool: + return self._active == ViewMode.DECOMP + + @property + def is_hex(self) -> bool: + return self._active == ViewMode.HEX + + @property + def is_graph(self) -> bool: + return self._active == ViewMode.GRAPH + + @property + def in_code(self) -> bool: + """A code view over a NavEntry -- where follow/xrefs/rename mean something.""" + return self._active in ViewMode.code_modes() + def _active_code_view(self): # type: ignore[no-untyped-def] - """The currently-shown code widget (for reading the cursor address).""" - if self._active in ("listing", "disasm"): + """The currently-shown code widget (for reading the cursor address). + + Everything downstream trusts ``_active``, so a new mode that forgets to + appear here doesn't degrade -- it returns None and the first caller that + dereferences it crashes the app (which is exactly how graph mode + announced itself the first time it was driven). + """ + if self.is_listing: return self.query_one(ListingView) - if self._active == "decomp": + if self.is_decomp: return self.query_one(DecompView) + if self.is_graph: + return self.query_one(GraphView) return None - def _palette_action(self, name: str) -> None: + def _focus_code_view(self) -> None: + """Put focus back on whichever code view is showing.""" + (self._active_code_view() or self.query_one(ListingView)).focus() + + def _palette_action(self, name: str, *args) -> None: """Run a cursor-scoped code-view action (rename/xrefs/follow/…) picked from the command palette against the active code view.""" view = self._active_code_view() @@ -5909,7 +7629,7 @@ class IdaTui(App): if fn is None: self._status(f"'{name}' isn't available in this view") return - fn() + fn(*args) def action_continuous_here(self) -> None: """'L': open the continuous segment listing at the cursor — one long @@ -5945,9 +7665,16 @@ class IdaTui(App): except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"goto: {e}") return - if self._active == "hex": + if self.is_hex: self.app.call_from_thread(self._hex_goto, ea) return + # A goto that lands inside the graph you're already looking at should + # move the cursor, not tear the picture down and build the same one. + if self.is_graph: + gv = self.query_one(GraphView) + if gv.fc is not None and gv.fc.block_at(ea) is not None: + self.app.call_from_thread(gv.goto_ea, ea) + return # Navigate to the containing function at the right line (handles both a # function name and a mid-function address). self._do_navigate(ea, push=True) @@ -5990,9 +7717,9 @@ class IdaTui(App): self._do_navigate(ea, push) - def _code_mode(self) -> str: + def _code_mode(self) -> ViewMode: """The code view to return to from hex — always the unified listing.""" - return "listing" + return ViewMode.LISTING def _open_entry(self, entry: NavEntry, push: bool) -> None: if self.program is None: @@ -6003,7 +7730,7 @@ class IdaTui(App): if entry.view == "decomp": # This entry was viewed in the decompiler (a jump from pseudocode, or # a back/forward to one) — restore it there instead of the listing. - self._active = "decomp" + self._active = ViewMode.DECOMP dec = self.query_one(DecompView) if dec.loaded_ea == entry.ea: # already decompiled: reposition without a recompile @@ -6032,13 +7759,19 @@ class IdaTui(App): lst.load( lm, entry.name, cursor=entry.cursor, cursor_x=entry.cursor_x, scroll_y=sy, focus=focus) - self._active = "listing" + self._active = ViewMode.LISTING self._show_active() + # Graph mode is sticky: following a call from the graph should land in + # the callee's graph, not dump you back into the listing. The rebuild is + # async and re-checks _cur, so a fast second jump just drops the stale one. + if self._graph_sticky and entry.ea: + self._load_graph(entry.ea, entry.ea) # Prompt overlays that own the keyboard while visible; a background # navigation must not yank focus out from under them (else typed keys leak # into a code view as destructive verbs — e.g. 'u' = undefine). - _PROMPT_IDS = ("search", "rename", "comment", "retype", "goto", "func-filter") + _PROMPT_IDS = ("search", "rename", "comment", "retype", "goto", "export", + "func-filter") def _prompt_active(self) -> bool: for iid in self._PROMPT_IDS: @@ -6053,12 +7786,13 @@ class IdaTui(App): dec = self.query_one(DecompView) lst = self.query_one(ListingView) hx = self.query_one(HexView) + gv = self.query_one(GraphView) # Don't steal focus from an open prompt (search/rename/…) — a late async # navigation completing here would otherwise pull it into the code view. grab = not self._prompt_active() - if self._split and self._active in ("listing", "decomp"): + if self._split and self._active in (ViewMode.LISTING, ViewMode.DECOMP): # Side-by-side: listing (left) + pseudocode (right), one focused. - hx.display = False + hx.display = gv.display = False lst.display = dec.display = True self.query_one("#panes").set_class(True, "split") busy = self._cur is not None and dec.loaded_ea != self._cur.ea @@ -6068,7 +7802,7 @@ class IdaTui(App): else: dec.loading = False if grab: - (dec if self._active == "decomp" else lst).focus() + (dec if self.is_decomp else lst).focus() if busy and self._cur is not None: # Keep the in-flight message: the idle status used to overwrite # it, so travelling history in split showed nothing at all while @@ -6084,13 +7818,19 @@ class IdaTui(App): self._split_eamap = [] self._split_ea2line = {} self._split_range = None - dec.display = lst.display = hx.display = False - if self._active in ("listing", "disasm"): + dec.display = lst.display = hx.display = gv.display = False + if self.is_graph: + gv.display = True + if grab: + gv.focus() + self._graph_status() + return + if self.is_listing: lst.display = True if grab: lst.focus() self._status_for_cur("listing") - elif self._active == "hex": + elif self.is_hex: hx.display = True if grab: hx.focus() @@ -6131,7 +7871,7 @@ class IdaTui(App): self._status("hex: no loaded segments") return hx.load(model, ea) - if self._active == "hex": + if self.is_hex: hx.focus() def _hex_status(self, va: int) -> None: @@ -6153,10 +7893,6 @@ class IdaTui(App): self._goto_ea(msg.va, push=True) def _status_for_cur(self, mode: str) -> None: - flash, self._flash = self._flash, None - if flash: - self._status(flash) # the edit that caused this reload wins - return if self._cur is not None: self._status(f"{self._cur.name} @ {self._cur.ea:#x} [{mode}]") @@ -6164,31 +7900,54 @@ class IdaTui(App): def _load_decomp(self, ea: int, name: str) -> None: assert self.program is not None dec = self.program.decompile(ea) - self.app.call_from_thread(self._apply_decomp, ea, name, dec) + why = "" + if not dec.failed: + # Where the number literals are, fetched in the same worker as the + # decompile (it is one call and the answer is cached with it) so the + # view can mark the one under the cursor without a round trip per + # keypress. + self._pending_nums = self.program.pc_nums(ea) + if dec.failed: + # Ask Hex-Rays why, in the same worker: the plain tool reports + # "Decompilation failed at 0x0" and drops the only useful part. + # "Decompile failed" with no reason is indistinguishable from a bug + # in this app, and for the common cause (a 32-bit function in a + # 64-bit database) the user cannot even guess the fix. + why = self.program.decomp_error(ea) + self.app.call_from_thread(self._apply_decomp, ea, name, dec, why) - def _apply_decomp(self, ea: int, name: str, dec) -> None: # type: ignore[no-untyped-def] + def _apply_decomp(self, ea: int, name: str, dec, # type: ignore[no-untyped-def] + why: str = "") -> None: view = self.query_one(DecompView) view.loading = False if dec.failed: + detail = f" \u2014 {why}" if why else "" # No pseudocode for this function: fall back to the code view rather # than an error panel. If we came from the continuous listing (F5), - # return there; otherwise show the disassembly. + # return to exactly where we left; otherwise just show the listing. if self._cur is None or self._cur.ea != ea: return # navigated away; stale result if self._decomp_return is not None: ret = self._decomp_return self._decomp_return = None self._cur = ret - self._active = "listing" + self._active = ViewMode.LISTING + # Hand the reason over as a flash BEFORE reopening: going back + # to the listing reloads it, and the reload writes its own + # status afterwards — which is precisely how "F5 does nothing" + # looked like nothing at all. + msg = f"{name}: cannot decompile{detail}" + self._status(msg, priority=True) self._open_entry(ret, push=False) - self._status(f"{name} — decompile failed; back to the listing") return - self._active = "disasm" + # LISTING, not the old "disasm": it is the same widget, and a value + # only this path produced meant half the app took the wrong branch + # for it (see ViewMode). + self._active = ViewMode.LISTING + self._status(f"{name}: cannot decompile{detail}", priority=True) self._show_active() - self._status( - f"{name} — no pseudocode (decompile failed); showing disassembly") return - if self._active == "decomp": + if self.is_decomp: view.focus() # loading cover had blurred it; restore focus # Restore the saved pseudocode position when returning to this function. cur = self._cur @@ -6199,6 +7958,7 @@ class IdaTui(App): sx = cur.dec_scroll_x if same else 0 note = " (truncated)" if dec.truncated else "" view.show(ea, dec.code or "", cursor=c, cursor_x=cx, scroll_y=sy, scroll_x=sx) + view.set_nums(self._pending_nums) if self._split: self._sync_split(self._active) # crude link now self._load_split_map(ea) # then upgrade to the region map @@ -6207,11 +7967,22 @@ class IdaTui(App): self._status( f"{name} @ {ea:#x} [pseudocode {len(dec.code or '')} chars]{note}") - def _sync_split(self, source: str) -> None: + def _sync_split(self, source: str, resync: bool = True) -> None: """Split view: highlight (+ scroll into view) the companion pane's location for the focused pane's cursor. The companion only gets a band + scroll (its cursor never moves), so there is no echo/ping-pong. Uses the - rich per-line ea map (decomp_map) when loaded, else the single marker.""" + rich per-line ea map (decomp_map) when loaded, else the single marker. + + ``resync=False`` means "the decompiler has just been re-pointed for this + anchor, don't ask again". Without it this recursed forever: _split_range + is the min/max of the decomp_map's addresses, which does NOT cover every + address in the function (Hex-Rays doesn't attribute them all), so an + anchor inside the loaded function but outside that span asked for a + resync, got back the function already loaded, re-entered here, and asked + again -- one worker and one lookup_funcs round trip per iteration, for as + long as the cursor sat there. 21,156 calls for a single address in one + scenario, and in the live app an idle split view pegging the worker. + """ if not self._split or self.program is None: return lst = self.query_one(ListingView) @@ -6244,7 +8015,7 @@ class IdaTui(App): dec.set_link(None) return rng = self._split_range - if rng is not None and not (rng[0] <= ea <= rng[1]): + if resync and rng is not None and not (rng[0] <= ea <= rng[1]): # left the decompiled function — follow to whatever function is # under the anchor (the unified view spans many functions). self._resync_decomp(ea) @@ -6270,11 +8041,11 @@ class IdaTui(App): return top, 0 def on_listing_view_scrolled(self, msg: "ListingView.Scrolled") -> None: - if self._split and self._active == "listing": + if self._split and self.is_listing: self._sync_split("listing") def on_decomp_view_scrolled(self, msg: "DecompView.Scrolled") -> None: - if self._split and self._active == "decomp": + if self._split and self.is_decomp: self._sync_split("decomp") @work(thread=True, group="split-resync", exclusive=True) @@ -6294,7 +8065,10 @@ class IdaTui(App): return self._cur.ea, self._cur.name = fn.addr, fn.name if dec.loaded_ea == fn.addr: # already decompiled (scrolled back): just relink - self._sync_split("listing") + # resync=False: we ARE the resync. Re-entering the resync branch is + # how this looped forever for an address the decomp_map's span + # doesn't cover. + self._sync_split("listing", resync=False) return dec.loading = True self._load_decomp(fn.addr, fn.name) # -> _apply_decomp -> map -> re-sync @@ -6303,7 +8077,7 @@ class IdaTui(App): """A split-aware status line reflecting the focused pane + the link.""" if self._cur is None: return - if self._active == "decomp": + if self.is_decomp: dec = self.query_one(DecompView) ea = dec._line_ea(dec.cursor) n = (len(self._split_eamap[dec.cursor]) @@ -6341,8 +8115,18 @@ class IdaTui(App): self.app.call_from_thread(self._apply_split_map, ea, m) def _apply_split_map(self, ea: int, m: list) -> None: - if not self._split or self._cur is None or self._cur.ea != ea: - return # left split / navigated away + """Index the per-line instruction map for the decompiled function. + + Keyed to what the DECOMPILER holds, not to _cur, and not conditional on + split being on. The old guard dropped the result whenever _cur had moved + while the fetch was in flight — during trace stepping that is almost + always — leaving the split view working from the map of the function you + just left. _cur follows the cursor; this map describes the pseudocode on + screen, and those are different things. + """ + dec = self._try_view(DecompView) + if dec is not None and dec.loaded_ea is not None and ea != dec.loaded_ea: + return # a stale fetch for a function we no longer show self._split_eamap = m self._split_ea2line = {} alleas = [] @@ -6353,7 +8137,12 @@ class IdaTui(App): # ea span of the decompiled function: when the listing cursor leaves it, # _sync_split re-points the decomp to the function under the cursor. self._split_range = (min(alleas), max(alleas)) if alleas else None - self._sync_split(self._active) # re-link with the region map + # ONE index, shared with the trace path: it used to keep a parallel copy + # of exactly this, fetched separately and keyed differently, which is how + # the two ended up describing different functions. + self.trace_ctl.adopt_map(ea, m, self._split_ea2line, self._split_range) + if self._split: + self._sync_split(self._active) # re-link with the region map def on_decomp_view_cursor_moved(self, msg: DecompView.CursorMoved) -> None: dv = self._try_view(DecompView) @@ -6361,7 +8150,7 @@ class IdaTui(App): self._nav[-1].dec_cursor = msg.index self._nav[-1].dec_cursor_x = dv.cursor_x if self._split: - if self._active == "decomp": + if self.is_decomp: self._sync_split("decomp") self._split_status() return @@ -6387,25 +8176,20 @@ class IdaTui(App): if msg.index >= 0: self._nav[-1].scroll_y = round(lst.scroll_offset.y) if self._split: - if self._active == "listing": + if self.is_listing: self._sync_split("listing") self._split_status() return ea = msg.ea if ea is not None: - # A pending flash (the result of an edit that caused this reload) - # outranks the idle hint: the cursor lands here as part of the - # reload, so this handler would otherwise always have the last word. - flash, self._flash = self._flash, None - if flash: - self._status(flash) - return sec = self.program.section_of(ea) if self.program else None self._status(f"{sec or '?'} @ {ea:#x} [listing] " "(c code · p func · u undefine · Enter follow)") # -- teardown ---------------------------------------------------------- # - def on_unmount(self) -> None: + async def on_unmount(self) -> None: + if self._rpc is not None: + await self._rpc.stop() if self._ka is not None: self._ka.stop() if self.program is not None: @@ -6417,7 +8201,7 @@ class IdaTui(App): elif self.client is not None: if self._save_on_exit is None and self._dirty: try: # unexpected teardown with edits: don't drop them - self.client.call("idb_save", timeout=600.0) + self.client.save_database() except Exception: # noqa: BLE001 pass self.client.close() diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py new file mode 100644 index 0000000..1ac995e --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1550 @@ +"""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. + +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. +""" +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import threading +import time +from pathlib import Path +from textwrap import dedent, indent +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` (257 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.client import ( + ClientError, + DatabaseHandle, + InstanceDisconnectedError, + RemoteError, + ) + from ida_codemode.registry import ( + REGISTRY_DIR, + FileLock, + RegistryEntry, + canonical_path, + idb_key, + scan_instances, + ) + from ida_codemode.resolver import IdbBusy, expected_idb_path +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. + ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc] + DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc] + REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment] + IdbBusy = expected_idb_path = 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-mcp is not installed in this environment " + f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " + "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a " + "database.") from _CODEMODE_ERROR + + +def database_owner(idb_path: str, staged_path: str | None = None): + """The registry entry 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. Registry errors that + happen WITH the library installed still propagate -- those mean "we could + not determine ownership", which is not the same as "nobody owns it". + """ + if DatabaseHandle is None: + return None + expected_key = idb_key(idb_path) + staged = canonical_path(staged_path) if staged_path else None + for item in scan_instances(timeout=0.5): + entry = item.entry + if entry.idb_key == expected_key: + return entry + if staged and entry.exe_path and canonical_path(entry.exe_path) == staged: + return entry + return None + + +def registered_database(path: str, output_database: str | None = None) -> bool: + """Whether a live/lock-held Code Mode instance owns this target.""" + _require_codemode() + source = canonical_path(path) + expected = canonical_path(output_database) if output_database else expected_idb_path(source) + expected_key = idb_key(expected) + for instance in scan_instances(timeout=0.5): + entry = instance.entry + if entry.idb_key == expected_key: + return True + if not output_database and entry.backend == "gui" and entry.exe_path: + if canonical_path(entry.exe_path) == source: + return True + return False + + +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 + + +#: Key of the pre-serialised payload envelope. See _script(). +_PACKED = "__idatui_json__" + +#: Serialise the answer INSIDE the database process and hand back one string. +#: +#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the +#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and +#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms +#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very +#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client +#: parses it, which it was going to do at the transport layer anyway. +_PACK_EPILOGUE = ( + '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' +) + + +#: Keep Code Mode's per-line trace hook installed while our snippet runs. +#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour. +_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0") + + +def _script(args: dict[str, Any], body: str) -> str: + """Bind JSON arguments without interpolating user text into Python code. + + Also runs the body with Code Mode's trace hook detached, which is worth an + order of magnitude. The runtime wraps every execute_python in + sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which + turns on line tracing in every frame it sees -- so every line of every + function we call pays a Python-level callback. Measured on this box: + ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process) + and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms + traced. That single hook was the whole residual gap against the old worker. + + What this gives up: the deadline is no longer enforced for a pure-Python + loop inside our snippet. The runtime's OTHER cancellation path -- a + threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of + the trace and still fires, so a long IDA operation is still interruptible; + and every operation here is bounded by its own count/limit argument. The + trace is restored in a finally, so a raising snippet cannot leak the change. + """ + encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) + head = f"import json\na = json.loads({encoded!r})\n" + if _KEEP_TRACE: + return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" + return ( + f"{head}" + "import sys\n" + "_idatui_trace = sys.gettrace()\n" + "sys.settrace(None)\n" + "try:\n" + f"{indent(dedent(body).strip(), ' ')}\n" + ' _idatui_packed = {"' + _PACKED + '": json.dumps(' + 'result, separators=(",", ":"), default=str)}\n' + "finally:\n" + " sys.settrace(_idatui_trace)\n" + "_idatui_packed\n" + ) + + +_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: + 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: + 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 + 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. +_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", ""))') + + +# 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 +''' + + +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_entry: RegistryEntry | None = None + self._connect_lock = threading.Lock() + + def _database_exists(self) -> bool: + """Whether the IDB this open would target is already on disk. + + Its loader switches are baked in, so they must not be sent again. + """ + try: + target = self._output_database or expected_idb_path(self._path) + except Exception: # noqa: BLE001 -- resolver unavailable: assume fresh + return False + return bool(target) and os.path.exists(target) + + 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 + 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: + # Loader switches describe how to IMPORT a raw file and + # are recorded in the database it produces. Sending them + # again for a database that already exists is a FATAL + # error in IDA itself ("Switch '-b400' can be used only + # when loading a new file"), which kills the worker + # before it can report anything useful. So: describe the + # import only when there is an import to describe. + fresh = self._new_database or not self._database_exists() + handle = DatabaseHandle.open( + self._path, + spawn=self._spawn, + timeout=max(0.1, timeout), + output_database=self._output_database, + processor=self._processor if fresh else None, + # DatabaseHandle calls this image_base and wants the + # natural (16-byte aligned) address; it does the + # conversion to IDA's paragraph-based -b itself. + image_base=self._loading_address if fresh else None, + file_type=self._file_type if fresh else None, + new_database=self._new_database, + ) + break + except IdbBusy: + if not self._new_database or time.monotonic() >= deadline: + raise + if progress: + progress("waiting for the previous Code Mode lease to close…") + # Remember the record before managed shutdown withdraws + # its JSON. The lifetime lock remains held until IDA has + # actually closed the IDB; waiting on it avoids racing a + # replacement worker into the old process's file lock. + expected = canonical_path( + self._output_database or expected_idb_path(self._path) + ) + owners = [item.entry for item in scan_instances(timeout=0.5) + if item.entry.idb_key == idb_key(expected)] + if owners: + self._wait_for_entry_release( + owners[0], max(0.0, deadline - time.monotonic()) + ) + else: + time.sleep(0.2) + if progress: + backend = handle.entry.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_entry = handle.entry + 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.entry.pid if self._handle is not None else None + + @property + def backend(self) -> str | None: + return self._handle.entry.backend if self._handle is not None else None + + def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + 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) + except RemoteError as exc: + details = exc.details or {} + message = str(exc) + if details.get("traceback"): + message += f"\n{details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError("execute_python", message) from exc + except (InstanceDisconnectedError, ClientError) 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: + 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 (InstanceDisconnectedError, ClientError) 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.entry + 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.entry.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.entry + 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_entry = handle.entry + handle.close() # release our lease; never close a GUI/other client's DB + + @staticmethod + def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool: + _require_codemode() + path = REGISTRY_DIR / f"{entry.record_id}.lock" + deadline = time.monotonic() + max(0.0, timeout) + while True: + lock = FileLock(path) + try: + if lock.try_acquire(): + return True + except OSError: + pass + finally: + lock.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.1, deadline - time.monotonic())) + + 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. + """ + entry = self._last_entry + if entry is None or entry.backend != "idalib": + return False + return self._wait_for_entry_release(entry, timeout) + + def __enter__(self) -> "CodeModeClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() diff --git a/idatui/diag.py b/idatui/diag.py new file mode 100644 index 0000000..b45a72d --- /dev/null +++ b/idatui/diag.py @@ -0,0 +1,121 @@ +"""Where swallowed errors go. + +A TUI must not die because one background load failed, so this codebase catches +broadly -- around fifty ``except Exception`` sites, two dozen of which resolve +to ``pass``. That is the right policy and it has one bad consequence: with +forty-odd ``@work(thread=True)`` workers, a failure in a background load leaves +no trace at all. The view just stays empty, and there is nothing to read +afterwards because the app owns the screen. + +``kittygfx`` already solved this for itself with ``$IDATUI_KITTY_LOG``. This is +the same idea for everything else: + +* ``$IDATUI_LOG=/tmp/x.log`` writes every swallowed error to a file. Unset (the + default) it costs an ``os.environ`` lookup and nothing else. +* The last few are kept in memory regardless, so ``drive diag`` can ask a live + app "what went wrong recently?" -- which is the question you actually have + when a driver reports success and the pane shows nothing. + +Use it where an exception would otherwise vanish:: + + with swallow("decomp_map(%#x)" % ea): + self._apply_split_map(ea, self.program.decomp_map(ea)) + +NOT for expected control flow. ``query_one`` raising because a modal owns the +screen is normal and happens constantly; wrapping that would bury the real +entries in noise. The test for whether it belongs here is "would I want to see +this after the fact?". +""" +from __future__ import annotations + +import contextlib +import os +import threading +import time +import traceback +from collections import deque + +#: Bounded on purpose: this is a debugging aid inside a long-running TUI, not an +#: audit log. Old entries are worth less than the memory. +_MAX = 50 +_ring: deque[dict] = deque(maxlen=_MAX) +_lock = threading.Lock() + + +def _logfile() -> str | None: + """Read the env var per call, not once at import. + + The pilot suite and the RPC tests set it after importing the app, and a + cached value would silently disable the thing being tested. + """ + return os.environ.get("IDATUI_LOG") or None + + +def log(msg: str) -> None: + """Append a line to ``$IDATUI_LOG``. No-op when it isn't set.""" + path = _logfile() + if not path: + return + try: + with open(path, "a", encoding="utf-8") as fh: + fh.write(f"{time.strftime('%H:%M:%S')} {msg}\n") + except OSError: + pass # a broken log path must never break the app + + +def note(what: str, exc: BaseException) -> None: + """Record a swallowed exception: in the ring always, in the log if enabled.""" + entry = { + "when": time.time(), + "what": what, + "error": f"{type(exc).__name__}: {exc}", + "where": _origin(exc), + "thread": threading.current_thread().name, + } + with _lock: + _ring.append(entry) + log(f"[swallowed] {what}: {entry['error']} ({entry['where']})") + if _logfile(): + log("".join(traceback.format_exception( + type(exc), exc, exc.__traceback__)).rstrip()) + + +def _origin(exc: BaseException) -> str: + """file:line where it was actually raised (the deepest frame we have).""" + tb = exc.__traceback__ + last = None + while tb is not None: + last = tb + tb = tb.tb_next + if last is None: + return "?" + f = last.tb_frame + return f"{os.path.basename(f.f_code.co_filename)}:{last.tb_lineno}" + + +@contextlib.contextmanager +def swallow(what: str, *, reraise: tuple = ()): + """Run a block, record anything it raises, and carry on. + + ``reraise`` lets a caller keep the exceptions it genuinely handles -- most + usefully ``IDAConnectionError``, which the app turns into a reconnect and + must not have eaten here. + """ + try: + yield + except reraise: + raise + except Exception as e: # noqa: BLE001 -- the whole point + note(what, e) + + +def recent(n: int = 10) -> list[dict]: + """The last ``n`` swallowed errors, newest last.""" + with _lock: + items = list(_ring) + return items[-n:] if n > 0 else items + + +def clear() -> None: + with _lock: + _ring.clear() diff --git a/idatui/domain.py b/idatui/domain.py index e4fbec9..d65c473 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -1,19 +1,15 @@ -"""Domain / paging layer: address-centric models over the raw MCP client. +"""Domain / paging layer: address-centric models over IDA Code Mode. This is where the "millions of lines" problem is solved, so the TUI widgets only ever see a viewport-sized slice. Every hard-won constraint from ``docs/PAGING_FINDINGS.md`` is encoded here: -* Per-call caps are silent (over the cap the server returns 10, not a clamp), so - we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. -* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. -* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly - is **block-cached** (revisits are free) and **prefetches** the next block on a - background thread (the client is concurrency-safe). -* ``include_total`` scans the whole function (~200ms on monsters); totals are - fetched once and cached. -* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is - null); that is surfaced as data, not an exception. +* Page sizes remain bounded so remote execution returns viewport-scale JSON. +* Pagination advances by the number of rows actually returned. +* Deep head walks are block-cached (revisits are free) and neighboring blocks + prefetch through the thread-safe Code Mode client. +* Expensive function totals are fetched once and cached. +* Decompilation failures are surfaced as data, not application crashes. Everything here is synchronous and thread-safe. The TUI runs these calls from Textual worker threads; the internal prefetch pool is separate and small. @@ -22,18 +18,19 @@ Textual worker threads; the internal prefetch pool is separate and small. from __future__ import annotations import bisect -import json import re import threading -import urllib.request from concurrent.futures import ThreadPoolExecutor +from collections.abc import Sequence from dataclasses import dataclass, field, replace +from typing import NamedTuple from typing import Callable, TYPE_CHECKING +from . import diag from .errors import IDAToolError if TYPE_CHECKING: # type hint only - from .worker_client import WorkerClient # noqa: F401 + from .codemode_client import CodeModeClient # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 @@ -63,7 +60,7 @@ class Func: def from_raw(cls, d: dict) -> "Func": addr = _as_int(d["addr"]) name = d.get("name") - # An unnamed function (server returns null/empty) must still have a + # An unnamed function must still have a # usable string name — synthesize IDA's sub_ADDR so every consumer # (palette, sort, rename prefill) can treat name as a str. if not name: @@ -89,10 +86,20 @@ class Line: ) -@dataclass(frozen=True) -class Head: - """One flat-listing item (from the ``heads`` server tool): a code - instruction, a data item, or an undefined byte run.""" +class Head(NamedTuple): + """One flat-listing item (from the Code Mode ``heads`` operation): a code + instruction, a data item, or an undefined byte run. + + A ``NamedTuple`` rather than a dataclass because this is by far the + most-constructed object in the codebase -- a jump to an address near the end + of a big binary builds one per listing row it walks past, a quarter of a + million of them -- and ``tuple.__new__`` costs 1.9us where a frozen + dataclass's ``__init__`` costs 2.9us. Attribute reads are marginally slower + (10ns vs 20ns), which is the right trade: rows are built far more often than + they are read, and a viewport only ever reads forty of them. + + Immutable, like the frozen dataclass it replaced. + """ ea: int kind: str # 'code' | 'data' | 'unknown' | 'member' @@ -100,23 +107,78 @@ class Head: text: str name: str | None = None raw: bytes | None = None # opcode/item bytes (filled in for code by the model) + #: [(kind, text)] from IDA's own colour tags — mnem/reg/num/name/str/punct/… + #: None when Code Mode didn't provide them (or the spans + #: disagreed with the plain text, in which case the text wins). + #: + #: Held exactly as it came off the wire, and **read-only**. The worker + #: memoises its per-line render, so one list is shared by every row that + #: says the same thing — pickle preserves that, and 228 000 rows of bash + #: reference about 53 000 lists. Copying each row's into a fresh tuple threw + #: the sharing away and cost 0.9 µs a row for nothing. + spans: Sequence | None = None + #: [(start, end, n)] — where each operand sits in ``text``, from IDA's own + #: COLOR_OPND markers. Lets the view show which operand the cursor is on, + #: and is the same information the worker maps a column through, so the + #: highlight and the edit can't disagree. Read-only, as ``spans`` is. + ops: Sequence | None = None @property def label(self) -> str | None: # Line-compatible alias return self.name + def op_at(self, col: int) -> tuple[int, int, int] | None: + """The operand whose text contains ``col``, or None.""" + for lo, hi, n in self.ops or (): + if lo <= col < hi: + return (lo, hi, n) + return None + @classmethod - def from_raw(cls, d: dict) -> "Head": + def from_raw(cls, d: dict, raw: bytes | None = None) -> "Head": + # Spans and operand extents are stored as they arrive: the worker's own + # tool emits [str, str] and [int, int, int], so re-coercing them was + # re-proving that once per listing row -- and copying them into tuples + # destroyed the sharing the worker's line cache had just created. return cls( ea=_as_int(d["ea"]), kind=d.get("kind", "unknown"), size=int(d.get("size", 0) or 0), text=d.get("text", ""), name=d.get("name"), + raw=raw, + spans=d.get("spans") or None, + ops=d.get("ops") or None, ) @dataclass +class BasicBlock: + """One node of a function's control-flow graph, with the listing rows that + make up its body (filled in by ``Program.flowchart``).""" + + id: int + start: int + end: int + succs: list[tuple[int, str]] = field(default_factory=list) + rows: list[Head] = field(default_factory=list) + + +@dataclass +class Flowchart: + func_ea: int + name: str + entry: int + blocks: list[BasicBlock] + + def block_at(self, ea: int) -> BasicBlock | None: + for b in self.blocks: + if b.start <= ea < b.end: + return b + return None + + +@dataclass class Ref: addr: int name: str @@ -190,6 +252,53 @@ def link_name(raw: str) -> str: @dataclass(frozen=True) +class SearchHit: + """One database-wide search result (Ctrl+F). + + ``addr`` is where the match starts -- for a byte pattern that can be inside + an instruction, so ``head`` is the item to navigate to and ``line`` is what + that item renders as. + """ + addr: int + head: int + line: str = "" + func: str | None = None + func_addr: int | None = None + seg: str = "" + + +@dataclass(frozen=True) +class Comment: + """One comment somebody wrote into the database. + + ``line`` is the disassembly the comment is attached to, carried along so a + report can show what was being commented ON without a second round trip. + ``whole_func`` marks a function comment rather than an instruction one. + """ + addr: int + text: str + repeatable: bool = False + whole_func: bool = False + line: str = "" + seg: str = "" + func: str | None = None + func_addr: int | None = None + + +@dataclass(frozen=True) +class NamedItem: + """An address carrying a real name -- one you typed, or one the file's own + symbols supplied. IDA records both as "user" names and does not remember + which was which, so a report must say so rather than claim authorship.""" + addr: int + name: str + is_func: bool = False + size: int = 0 + proto: str | None = None + seg: str = "" + + +@dataclass(frozen=True) class Linkage: """One import or export: a name this binary takes from, or offers to, other modules. ``module`` is set for imports (the library IDA attributes it to), @@ -233,7 +342,7 @@ class FunctionIndex: """A lazily-paginated, cached view of the function list. Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by - ``next_offset``). A single index instance corresponds to one server-side + ``next_offset``). A single index instance corresponds to one remote ``filter`` glob (``None`` = all functions). """ @@ -253,7 +362,7 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.call("list_funcs", queries=[query])) + data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) added = 0 with self._lock: for d in data: @@ -363,7 +472,7 @@ class DisasmModel: code function this equals the heads row count that backs the lines.""" if self._total is not None: return self._total - payload = self._prog.client.call( + payload = self._prog.client.invoke( "disasm", addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") @@ -424,7 +533,7 @@ class DisasmModel: # The function disasm view is a listing filtered to the function: fetch a # block of heads (one per instruction for code). Over-fetch one row so # the block knows where its last instruction ends (opcode-byte sizing). - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(self.ea), offset=b * self.BLOCK, count=self.BLOCK + 1, **self._end_kw(), ) @@ -565,15 +674,15 @@ class ListingModel: """A flat, IDA-style disassembly *listing* over one segment: code, data and undefined heads interleaved, unlike ``DisasmModel`` (one function, code only). - Backed by the injected ``heads`` server tool, which walks item heads and - renders each via ``generate_disasm_line``. The segment is walked lazily in + Backed by the Code Mode adapter's ``heads`` operation, which walks item heads + and renders each via ``generate_disasm_line``. The segment is walked lazily in forward pages (``FunctionIndex`` style); line index == position in the walked head list. Random access to an address is O(distance-from-seg-start) the first time (then cached) — the same tradeoff as ``disasm offset=N``. Grows on demand as the viewport scrolls. Synchronous + thread-safe. """ - PAGE = 500 # heads per server call (well under the tool's 2000 cap) + PAGE = 500 # viewport-scale heads per Code Mode execution def __init__(self, program: "Program", seg_start: int, seg_end: int, name: str | None = None): @@ -591,6 +700,27 @@ class ListingModel: # demand. _row_at[i] is the logical row where physical head i starts. self._row_at: list[int] = [] self._head_eas: list[int] = [] # parallel to _heads, for bisect + #: Which name generation each head's TEXT was rendered at, parallel to + #: _heads. A rename bumps :attr:`_text_gen`; the rows themselves stay + #: (their addresses and row numbers are unchanged) and are re-rendered a + #: block at a time when something asks for them. See invalidate_text. + self._head_gen: list[int] = [] + self._text_gen = 0 + #: Whether a rename has ever staled this model. Until one has, every + #: read takes exactly the path it always did. + self._renamed = False + #: One entry per loaded PAGE: where its heads start, the address it was + #: fetched from, the digest it came back with, and how many rows it + #: held. A stale-text refresh re-asks for exactly that page, so it can + #: be told "still identical" for the price of the render alone. + self._page_head: list[int] = [] + self._page_addr: list[int] = [] + self._page_digest: list[object] = [] + self._page_rows: list[int] = [] + #: Set if a text refresh came back with a different head sequence, which + #: means something DID move the walk. Program.listing() throws the model + #: away when it sees this, so the next read rebuilds from scratch. + self.stale_structure = False self._rows = 0 # total logical rows loaded self._ubytes: dict[int, bytes] = {} # lazily-read bytes for those rows self._next: int | None = seg_start # next address to fetch from @@ -605,30 +735,48 @@ class ListingModel: # containing a huge coalesced undefined run doesn't pull megabytes. _OP_SPAN_CAP = 1 << 16 - def _attach_opcode_bytes(self, page: list[Head]) -> list[Head]: - """Fill ``raw`` (opcode bytes) for the code heads in ``page`` via one - bulk read over their extent (variable-length safe).""" - code = [h for h in page if h.kind == "code" and h.size > 0] - if not code: - return page - lo = code[0].ea - hi = code[-1].ea + code[-1].size - if hi - lo <= 0 or hi - lo > self._OP_SPAN_CAP: - return page - data = self._prog.read_bytes(lo, hi - lo) + def _build_page(self, rows: list) -> list[Head]: + """Turn the tool's raw rows into ``Head``s with their opcode bytes + already attached, via one bulk read over the code extent. + + The bytes are read BEFORE the Heads are built rather than patched in + afterwards: ``dataclasses.replace`` re-runs ``__init__`` with every + field, so filling ``raw`` after the fact meant constructing each code + head twice -- once per listing row, on the path a jump-to-address walks + hundreds of thousands of times. + """ + lo = hi = -1 + for r in rows: + if r.get("kind") == "code" and r.get("size"): + ea = _as_int(r["ea"]) + if lo < 0: + lo = ea + hi = ea + int(r["size"]) + data = None + if 0 <= lo < hi and hi - lo <= self._OP_SPAN_CAP: + try: + data = self._prog.read_bytes(lo, hi - lo) + except Exception: # noqa: BLE001 -- opcode bytes are decoration + data = None + page: list[Head] = [] biggest = self._max_raw - out = [] - for h in page: - if h.kind == "code" and h.size > 0: - off = h.ea - lo - b = bytes(data[off:off + h.size]) - biggest = max(biggest, len(b)) - out.append(replace(h, raw=b)) - else: - out.append(h) - with self._lock: - self._max_raw = biggest - return out + for r in rows: + raw = None + if data is not None and r.get("kind") == "code": + size = int(r.get("size") or 0) + if size > 0: + off = _as_int(r["ea"]) - lo + raw = bytes(data[off:off + size]) + if len(raw) > biggest: + biggest = len(raw) + try: + page.append(Head.from_raw(r, raw)) + except (KeyError, ValueError, TypeError): + continue + if biggest != self._max_raw: + with self._lock: + self._max_raw = biggest + return page def max_raw_len(self) -> int: with self._lock: @@ -647,18 +795,18 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(frm), count=self.PAGE, annotate=True) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} - page = [] - for r in rows: - try: - page.append(Head.from_raw(r)) - except (KeyError, ValueError, TypeError): - continue - page = self._attach_opcode_bytes(page) + page = self._build_page(rows) with self._lock: + gen = self._text_gen + self._page_head.append(len(self._heads)) + self._page_addr.append(frm) + self._page_digest.append(payload.get("digest") + if isinstance(payload, dict) else None) + self._page_rows.append(len(rows)) for h in page: # Banner/label rows (function headers, separators, code labels) # are display-only; don't index them so navigation lands on the @@ -667,6 +815,7 @@ class ListingModel: self._by_ea.setdefault(h.ea, self._rows) self._row_at.append(self._rows) self._head_eas.append(h.ea) + self._head_gen.append(gen) self._heads.append(h) self._rows += self._span(h) nxt = cur.get("next") @@ -684,7 +833,9 @@ class ListingModel: def _phys(self, row: int) -> tuple[int, int]: """(physical head index, byte offset into it) for logical ``row``.""" - import bisect + # bisect is imported at module scope; re-importing it here cost a + # sys.modules lookup on a function that runs once per rendered row and + # once per row a search reads. i = bisect.bisect_right(self._row_at, row) - 1 if i < 0: return (-1, 0) @@ -796,6 +947,156 @@ class ListingModel: def __len__(self) -> int: return self.loaded() + def truncate_from(self, ea: int) -> bool: + """Drop the walk from the page an edit at ``ea`` could have moved. + + An item edit changes structure, but only *locally*: every head before it + keeps its address and its row number. Throwing the whole model away made + the reload re-walk the segment -- 4.9 seconds on bash to make one byte + into data, for an edit the user made at the row they were looking at. + + Two pages are dropped rather than one, because undefining can coalesce + backwards into the run in front of it. Beyond that the caller marks the + kept prefix text-stale, so every kept page is digest-checked on the next + read and a page that really did move fails its sequence check and forces + a rebuild. Safe by construction, not by argument. + + Returns False if nothing worth keeping is left. + """ + with self._lock: + if not (self.seg_start <= ea < self.seg_end): + return True # another segment; nothing moved here + if len(self._page_head) < 3: + return False # barely walked; a rebuild is cheaper + p = bisect.bisect_right(self._page_addr, ea) - 1 + p = max(p - 1, 0) + if p <= 0: + return False # the edit is in the first pages + keep = self._page_head[p] + if keep <= 0: + return False + for h in self._heads[keep:]: + self._by_ea.pop(h.ea, None) + del self._heads[keep:] + del self._head_eas[keep:] + del self._head_gen[keep:] + del self._row_at[keep:] + del self._page_head[p:] + self._next = self._page_addr[p] + del self._page_addr[p:] + del self._page_digest[p:] + del self._page_rows[p:] + last = self._heads[-1] + self._rows = self._row_at[-1] + self._span(last) + self._done = False + self._ubytes.clear() # undefined-run bytes behind the drop point + return True + + def invalidate_text(self) -> None: + """A rename changed how rows READ, not which rows exist. + + Item boundaries are untouched by a rename, so every row keeps its + address and its row number — which the edit path already relies on, since + it restores the cursor by INDEX afterwards. Dropping the whole model + instead means the next jump re-walks the segment from its start: 6.4 + seconds on bash's .text, after every single rename. + + So keep the walk and mark the rendered text stale; :meth:`_ensure_text` + re-renders a block at a time, and refuses to splice anything back if the + head sequence has moved under it (which a rename cannot do, but a + mis-routed structural edit could). + """ + with self._lock: + self._text_gen += 1 + self._renamed = True + + def _ensure_text(self, j0: int, j1: int) -> None: + """Re-render physical heads [j0, j1) if a rename staled them. + + Works a PAGE at a time -- the same unit the loader fetched. A page is + exactly what ``heads(addr, count=PAGE)`` produced, so asking again with + the same arguments reproduces the same row sequence; nothing has to be + snapped out to whole address groups (a function start emits three banner + rows sharing one address, and an arbitrary boundary through those never + lines up again). It also means every head in a page can share one + generation marker, so "is this fresh?" is a single probe. + """ + with self._lock: + n = len(self._heads) + j1 = min(j1, n) + j0 = max(j0, 0) + if j1 <= j0: + return + p = max(bisect.bisect_right(self._page_head, j0) - 1, 0) + last = bisect.bisect_left(self._page_head, j1) + while p < last: + p = self._ensure_page(p) + + def _page_bounds(self, p: int) -> tuple[int, int]: + """[first, last) head index of page ``p`` (caller holds the lock).""" + lo = self._page_head[p] + hi = (self._page_head[p + 1] if p + 1 < len(self._page_head) + else len(self._heads)) + return lo, hi + + def _ensure_page(self, p: int) -> int: + """Freshen page ``p``; returns the next page to consider.""" + with self._lock: + if not (0 <= p < len(self._page_head)): + return p + 1 + gen = self._text_gen + lo, hi = self._page_bounds(p) + if hi <= lo or self._head_gen[lo] == gen: + return p + 1 + addr = self._page_addr[p] + want_digest = self._page_digest[p] + want_rows = self._page_rows[p] + want = [(h.ea, h.kind) for h in self._heads[lo:hi]] + # Tell the worker what we already hold. It builds the rows either way + # (there is no knowing a line is unchanged without rendering it), but if + # they still hash to the same value it keeps them: the pickling, the + # transfer, the unpickling and the Head rebuild are about 40% of what a + # page costs, and after a rename almost every page is unchanged. Sending + # the expectation rather than asking first means a page that HAS changed + # still costs one round trip. + try: + payload = self._prog.client.invoke( + "heads", addr=hex(addr), count=self.PAGE, annotate=True, + expect="" if want_digest is None else str(want_digest)) + except Exception: # noqa: BLE001 -- keep the old text rather than blank + return p + 1 + if (isinstance(payload, dict) and "heads" not in payload + and payload.get("count") == want_rows): + with self._lock: + if self._text_gen == gen and len(self._heads) >= hi: + for k in range(lo, hi): + self._head_gen[k] = gen + return p + 1 + rows = payload.get("heads", []) if isinstance(payload, dict) else [] + page = self._build_page(rows) + with self._lock: + if self._text_gen != gen or len(self._heads) < hi: + return p + 1 + if [(h.ea, h.kind) for h in page] != want: + # Something moved the walk, which a rename cannot do -- so this + # was not one. Say so and let Program.listing() rebuild, rather + # than sit here re-fetching a page that will never line up (and + # showing the old names while doing it). + self.stale_structure = True + for k in range(lo, hi): + self._head_gen[k] = gen + return p + 1 + self._heads[lo:hi] = page + # The stored digest has to describe what the client now HOLDS, not + # what it once loaded. Leaving it stale is how a literal cycling + # hex -> dec -> hex ends up declared "unchanged" while the row still + # shows the decimal it was refetched with in between. + self._page_digest[p] = (payload.get("digest") + if isinstance(payload, dict) else None) + for k in range(lo, hi): + self._head_gen[k] = gen + return p + 1 + def get(self, i: int) -> Head | None: with self._lock: if not (0 <= i < self._rows): @@ -803,8 +1104,23 @@ class ListingModel: j, off = self._phys(i) if j < 0: return None - span = self._span(self._heads[j]) - h = self._heads[j] + stale = self._renamed and self._head_gen[j] != self._text_gen + if not stale: + span = self._span(self._heads[j]) + h = self._heads[j] + if stale: + # A rename staled this row's text; re-render its block (one call for + # the block around it, so a viewport costs one round trip). Only + # this path re-takes the lock -- the ordinary read stays atomic. + self._ensure_text(j, j + 1) + with self._lock: + if not (0 <= i < self._rows): + return None + j, off = self._phys(i) + if j < 0: + return None + span = self._span(self._heads[j]) + h = self._heads[j] # Synthesis reads bytes, so do it OUTSIDE the lock: an RPC under the # model lock deadlocks the page loader that is filling it. return self._row_head(j, off) if span > 1 else h @@ -813,6 +1129,16 @@ class ListingModel: """``count`` logical rows from ``start`` (synthesising undefined ones).""" self.ensure(start + count) with self._lock: + # _renamed stays set once a rename has happened; _ensure_text then + # does the precise, range-limited staleness check. Before the first + # rename this is one boolean and the read is exactly as it was. + dirty = self._renamed + if dirty: + j0 = max(self._phys(max(start, 0))[0], 0) + j1 = self._phys(max(min(self._rows, start + count) - 1, 0))[0] + 1 + if dirty: + self._ensure_text(j0, j1) + with self._lock: rows = min(self._rows, start + count) spans = [self._phys(i) for i in range(max(start, 0), max(rows, 0))] heads = self._heads @@ -839,7 +1165,6 @@ class ListingModel: def _head_index_at(self, ea: int) -> int: """Index of the physical head containing ``ea`` (caller holds the lock).""" - import bisect eas = self._head_eas i = bisect.bisect_right(eas, ea) - 1 return i if 0 <= i < len(self._heads) else -1 @@ -949,7 +1274,7 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): + def __init__(self, client: "CodeModeClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" @@ -958,7 +1283,14 @@ class Program: self._disasm: dict[int, DisasmModel] = {} self._listings: dict[int, ListingModel] = {} # keyed by segment start self._decomp: dict[int, tuple[Decompilation, int]] = {} + #: {func ea: ({line: [(x0, x1, value)]}, name generation)} — literal + #: positions in the pseudocode, cached alongside the decompilation. + self._pc_nums: dict[int, tuple[dict, int]] = {} self._decomp_maps: dict[int, tuple[list[list[int]], int]] = {} # line->ea sets + #: {func ea: (Flowchart, name generation)} — the CFG plus its block rows. + #: Keyed off _name_gen, which BOTH bump_names and bump_items raise: the + #: rows carry live symbol names, so a rename must refetch them too. + self._flowcharts: dict[int, tuple["Flowchart", int]] = {} self._strings: list["StrLit"] | None = None # whole-binary string literals self._linkage: tuple[list["Linkage"], list["Linkage"]] | None = None self._name_gen = 0 # bumped on rename; invalidates stale name caches @@ -966,7 +1298,7 @@ class Program: self._sections: list[tuple[int, int, str]] | None = None self._fileregions: list[tuple[int, int, int]] | None = None self._hexmodel: "HexModel | None" = None - self._no_read_raw = False # set if the server lacks the read_raw tool + self._no_read_raw = False # compatibility fallback for alternate clients self._lock = threading.Lock() # -- prefetch plumbing ------------------------------------------------- # @@ -993,17 +1325,14 @@ class Program: """Sorted raw segment map [(start, end, file_off, name)] — the single source for sections()/file_regions()/image_range. Cached. - Uses the injected ``file_regions`` tool (a plain segment walk, ~ms). - This deliberately AVOIDS ``survey_binary``, which also computes function - counts / strings / stats and takes *seconds* on a large IDB (it was the - cause of the multi-second hex-pane open). Falls back to survey_binary - only if the injected tool is missing. + Uses the Code Mode adapter's ``file_regions`` operation (a plain segment + walk, ~ms), avoiding broad binary surveys on the hex-pane open path. """ if self._segments_cache is not None: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.call("file_regions") + r = self.client.invoke("file_regions") for d in (r.get("regions", []) if isinstance(r, dict) else []): if isinstance(d, dict) and "start" in d: segs.append((_as_int(d["start"]), _as_int(d["end"]), @@ -1012,7 +1341,7 @@ class Program: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.call("survey_binary") + sb = self.client.invoke("survey_binary") for s in (sb.get("segments", []) if isinstance(sb, dict) else []): try: segs.append((_as_int(s["start"]), _as_int(s["end"]), -1, @@ -1054,8 +1383,7 @@ class Program: def file_regions(self) -> list[tuple[int, int, int]]: """Sorted [(start, end, file_off)] mapping loaded segments to raw file - offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs - the injected ``file_regions`` server tool.""" + offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached.""" if self._fileregions is not None: return self._fileregions regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] @@ -1073,15 +1401,14 @@ class Program: def read_bytes(self, ea: int, n: int) -> bytes: """Raw bytes [ea, ea+n) from IDA (gaps read as zero). - Fast path: the injected ``read_raw`` tool returns one contiguous hex - string (C-speed both ends). Falls back to the stock ``get_bytes`` (a - per-byte '0x..'-with-spaces string) on an older server without it. + The Code Mode adapter returns one contiguous hex string (C-speed in IDA). + A legacy ``get_bytes`` decoding fallback remains for alternate clients. """ if n <= 0: return b"" if not self._no_read_raw: try: - r = self.client.call("read_raw", addr=hex(ea), size=int(n)) + r = self.client.invoke("read_raw", addr=hex(ea), size=int(n)) h = r.get("hex") if isinstance(r, dict) else None if isinstance(h, str): out = bytes.fromhex(h) @@ -1095,7 +1422,7 @@ class Program: except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.call("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1135,6 +1462,8 @@ class Program: start, end, name = seg with self._lock: m = self._listings.get(start) + if m is not None and m.stale_structure: + m = None # a refresh found the walk had moved; start over if m is None: m = ListingModel(self, start, end, name) self._listings[start] = m @@ -1144,7 +1473,7 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.call("search_structs", filter=filter) + payload = self.client.invoke("search_structs", filter=filter) res = payload.get("result", []) if isinstance(payload, dict) else [] out = [Struct.from_raw(d) for d in res if isinstance(d, dict) and d.get("name") @@ -1154,9 +1483,9 @@ class Program: def struct_source(self, name: str) -> str: """A C definition for ``name`` reconstructed from its member layout - (the server exposes members, not printable source). Faithful to IDA's + (the remote operation exposes members, not printable source). Faithful to IDA's field names/types; array dims are moved after the field name.""" - payload = self.client.call( + payload = self.client.invoke( "type_inspect", queries=[{"name": name, "include_members": True}]) res = payload.get("result", []) if isinstance(payload, dict) else [] info = res[0] if res and isinstance(res[0], dict) else {} @@ -1179,7 +1508,7 @@ class Program: def declare_type(self, decl: str) -> str | None: """Create or update a C type. Returns None on success, else the parse error. (Re-declaring a name updates it in place.)""" - payload = self.client.call("declare_type", decls=decl) + payload = self.client.invoke("declare_type", decls=decl) res = payload.get("result", []) if isinstance(payload, dict) else [] if res and isinstance(res[0], dict): return res[0].get("error") @@ -1188,10 +1517,9 @@ class Program: # -- function / variable types ---------------------------------------- # def func_types(self, ea: int) -> FuncTypes | None: """Structured decompiler types for the function at ``ea`` (prototype + - local variables). None if ``ea`` isn't a decompilable function. Requires - the injected ``func_types`` server tool.""" + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.call("func_types", addr=hex(ea)) + r = self.client.invoke("func_types", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1204,7 +1532,7 @@ class Program: def set_function_type(self, ea: int, signature: str) -> str | None: """Set a function's prototype. None on success, else an error string.""" - r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} if row.get("ok"): @@ -1213,9 +1541,9 @@ class Program: def data_type(self, ea: int) -> dict | None: """Current type info for a data item/global: {addr,name,type,size,is_func}. - None if the tool is unavailable or the address isn't mapped.""" + None if the operation fails or the address isn't mapped.""" try: - r = self.client.call("data_type", addr=hex(ea)) + r = self.client.invoke("data_type", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1224,7 +1552,7 @@ class Program: def set_data_type(self, ea: int, decl: str) -> str | None: """Set a global/data item's type. None on success, else an error string.""" - r = self.client.call( + r = self.client.invoke( "set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} @@ -1233,9 +1561,9 @@ class Program: return row.get("error") or "failed to set the type" def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: - """Set a decompiler local variable's type (via the injected server tool). + """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) if isinstance(r, dict) and r.get("error"): return r["error"] if isinstance(r, dict) and not r.get("ok"): @@ -1244,15 +1572,14 @@ class Program: def delete_type(self, name: str) -> str | None: """Delete a named type. Returns None on success, else an error string. - Requires a server-side ``del_type`` tool; if absent, a clear message is - returned instead of raising.""" + Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.call("del_type", name=name) + self.client.invoke("del_type", name=name) return None except IDAToolError as e: msg = e.message if "not found" in msg.lower() and "del_type" in msg: - return "delete needs a 'del_type' tool on the ida-pro-mcp server" + return "the connected Code Mode runtime cannot delete local types" return msg # -- disassembly ------------------------------------------------------- # @@ -1266,13 +1593,7 @@ class Program: # -- decompilation ----------------------------------------------------- # def decompile(self, ea: int, refresh: bool = False) -> Decompilation: - """Full pseudocode for a function. - - The server truncates responses over 50KB (strings clipped to 1000 - chars) but caches the full output and exposes it at - ``_meta.ida_mcp.download_url``. We transparently fetch that so the view - always gets the complete body, not a 1KB stub. - """ + """Full pseudocode for a function, returned directly by Code Mode.""" if not refresh: with self._lock: hit = self._decomp.get(ea) @@ -1281,10 +1602,10 @@ class Program: dec, hit_gen = hit if hit_gen == gen: return dec - # Cached before a rename: names may be stale. Drop the server's - # Hex-Rays cache so the refetch reflects the new names. + # Cached before a rename: names may be stale. Drop Hex-Rays' + # cache so the refetch reflects the new names. try: - self.client.call("force_recompile", items=[{"addr": hex(ea)}]) + self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 pass # Bound the decompile: a function Hex-Rays can't handle tends to stall @@ -1294,7 +1615,10 @@ class Program: # rpcclient socket timeout, and cache the failure below so a re-request # returns instantly instead of re-grinding. try: - envelope = self.client.call_envelope( + # Code Mode returns the complete JSON result directly; unlike the + # old MCP tool transport there is no structured-content envelope or + # out-of-band download URL to unwrap. + payload = self.client.invoke( "decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT ) except Exception as e: # noqa: BLE001 -- surface as a failed decompile @@ -1302,15 +1626,6 @@ class Program: with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec - result = envelope.get("result", {}) - payload = result.get("structuredContent") - if payload is None: # fall back to text content - payload = self.client._extract_payload("decompile", result) - meta = (result.get("_meta") or {}).get("ida_mcp") - if isinstance(meta, dict) and meta.get("download_url"): - full = self._fetch_output(meta["download_url"]) - if isinstance(full, dict) and full.get("code"): - payload = full dec = _parse_decompilation(ea, payload) with self._lock: self._decomp[ea] = (dec, self._name_gen) @@ -1318,30 +1633,60 @@ class Program: def bump_names(self) -> None: """Signal that symbol names changed (a rename). Disasm/listing names are - live in the IDB, so clearing the cached rows is enough for those; - decompilation is generation-checked and force-recompiled lazily.""" + live in the IDB, so the cached rows have to be re-rendered; decompilation + is generation-checked and force-recompiled lazily. + + The listing keeps its WALK. A rename cannot move an item boundary, so + every row keeps its address and its row number -- the edit path already + assumes exactly that, since it restores the cursor by index afterwards. + Dropping the segment model instead made the reload re-walk it from the + start, which is 6.4 seconds on bash after every rename. + """ with self._lock: self._name_gen += 1 models = list(self._disasm.values()) - self._listings.clear() # listing head rows cache names -> refetch + listings = list(self._listings.values()) + self._pc_nums.clear() # a reformat moves every literal on its line for m in models: m.invalidate() + for lm in listings: + lm.invalidate_text() - def bump_items(self) -> None: + def bump_items(self, ea: int | None = None) -> None: """Signal that item/function STRUCTURE changed (define code/data/func, undefine). Unlike a rename this can move instruction boundaries and - change function membership anywhere, so drop the disasm block caches, - the decompilation cache and the cached function indices outright, and - bump the name generation too (labels/names may appear or vanish).""" + change function membership, so drop the disasm block caches, the + decompilation cache and the cached function indices outright, and bump + the name generation too (labels/names may appear or vanish). + + Given the address that was edited, the segment listing keeps the walk in + front of it instead of being thrown away: the rows before an edit keep + their addresses and their row numbers. Without ``ea`` this falls back to + discarding the listings, as it always did. + """ with self._lock: self._name_gen += 1 self._indices.clear() self._decomp.clear() - self._listings.clear() + self._pc_nums.clear() models = list(self._disasm.values()) self._disasm.clear() + listings = list(self._listings.items()) + if ea is None: + self._listings.clear() for m in models: m.invalidate() + if ea is None: + return + for start, lm in listings: + if lm.truncate_from(ea): + # Names can move too; the kept prefix is re-rendered on demand, + # and that is also what catches a page the edit really did move. + lm.invalidate_text() + else: + with self._lock: + if self._listings.get(start) is lm: + del self._listings[start] # -- item / function structure edits (IDA c/d/u/p) --------------------- # @staticmethod @@ -1358,22 +1703,59 @@ class Program: Undefine first so it works even when the bytes are currently part of a data/align item — ``create_insn`` refuses to carve into a live item.""" try: - self.client.call("undefine", items=[{"addr": hex(ea)}]) + self.client.invoke("undefine", items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.call("define_code", items=[{"addr": hex(ea)}])) + self.client.invoke("define_code", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_code", f"@ {ea:#x}: {res['error']}") + def decomp_error(self, ea: int) -> str: + """Hex-Rays' own reason for refusing ``ea``, or "" if it won't say.""" + try: + r = self.client.invoke("decomp_error", addr=hex(ea)) + except IDAToolError: + return "" + if not isinstance(r, dict): + return "" + reason = str(r.get("reason") or "") + if reason and r.get("bitness") == 64 and "64-bit" in reason: + # Say the FIX, not the diagnosis. Hex-Rays' own sentence ("only + # 64-bit functions can be decompiled in the current database") is + # accurate and useless: it describes the database, not what to do, + # and it's long enough that a status bar cuts off the end — which is + # exactly where an appended hint would live. This is unfixable in + # place (bitness is decided at load), so the whole message is the + # instruction. + return "this database is 64-bit \u2014 Ctrl+L, pick arm:ARMv7-A" + return reason + + def thumb_scan(self, start: int, end: int, apply: bool = True) -> dict: + """Find Thumb entry points from odd pointers in ``[start, end)``.""" + r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end), + apply=bool(apply)) + if not isinstance(r, dict) or r.get("error"): + raise IDAToolError("thumb_scan", + f"@ {start:#x}: {(r or {}).get('error', 'failed')}") + return r + + def set_thumb(self, ea: int, mode: str = "toggle") -> dict: + """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" + r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode) + if not isinstance(r, dict) or r.get("error"): + raise IDAToolError("set_thumb", + f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + return r + def define_code_run(self, ea: int, limit: int = 20000) -> dict: """Disassemble consecutively from ``ea`` until something stops it. - Falls back to a single instruction when the worker predates the tool, so - an old worker degrades to the previous behaviour instead of failing. + Falls back to a single instruction for alternate clients that do not + provide the run operation. """ try: - r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} @@ -1382,19 +1764,31 @@ class Program: f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") return r - def define_func(self, ea: int) -> None: - """Create a function starting at ``ea`` (IDA's 'p').""" - res = self._first_result( - self.client.call("define_func", items=[{"addr": hex(ea)}])) - if res.get("error"): - raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") + def define_func(self, ea: int) -> dict: + """Create a function starting at ``ea`` (IDA's 'p'). + + Prefers the Code Mode operation, which works out the end when IDA can't; + falls back to a plain create for alternate clients. + """ + try: + r = self.client.invoke("define_func_run", addr=hex(ea)) + except IDAToolError: + res = self._first_result( + self.client.invoke("define_func", items=[{"addr": hex(ea)}])) + if res.get("error"): + raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") + return {"ok": True, "how": "legacy"} + if not isinstance(r, dict) or not r.get("ok"): + raise IDAToolError("define_func", + f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + return r def undefine(self, ea: int, size: int | None = None) -> None: """Undefine the item at ``ea`` back to raw bytes (IDA's 'u').""" item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.call("undefine", items=[item])) + res = self._first_result(self.client.invoke("undefine", items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1404,7 +1798,7 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.call("make_data", items=[item])) + res = self._first_result(self.client.invoke("make_data", items=[item])) if res.get("ok") is False or res.get("error"): raise IDAToolError( "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}") @@ -1412,13 +1806,79 @@ class Program: def make_string(self, ea: int, length: int = 0, kind: str = "c") -> str: """Create a string literal at ``ea`` (IDA's 'A'); auto-length when 0. Returns the decoded contents.""" - r = self.client.call("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( "make_string", f"@ {ea:#x}: {res.get('error') or 'rejected'}") return res.get("text", "") + # -- literal display formats (IDA's 'o': hex / dec / char / offset) ---- # + def op_format(self, ea: int, mode: str = "cycle", col: int = -1, + n: int = -1) -> dict: + """Change how the literal at ``ea`` is DISPLAYED in the listing. + + ``col`` is a column inside the rendered line, which is how the cursor + says *which* operand it means; ``n`` names one outright. ``mode`` is + ``cycle``/``back`` (step the stops that make sense for this value) or a + format by name. ``show`` reports without changing anything. + """ + r = self.client.invoke("op_format", addr=hex(ea), mode=str(mode), + col=int(col), n=int(n)) + res = r if isinstance(r, dict) else {} + if res.get("error"): + raise IDAToolError("op_format", f"@ {ea:#x}: {res['error']}") + if not res: + raise IDAToolError("op_format", f"@ {ea:#x}: no answer") + return res + + def pc_nums(self, fn_ea: int) -> dict[int, list[tuple[int, int, str, int, int]]]: + """{pseudocode line: [(x0, x1, value, ea, opnum), ...]} — every number + literal in a function's decompilation, so the view can show which one + the cursor is on. One worker call per decompilation (cached with it); + the alternative is a round trip per cursor move. + + ``ea``/``opnum`` identify a literal across a reformat: the text reflows + (``48`` becomes ``0x30``) and a column no longer means the same thing. + """ + with self._lock: + hit = self._pc_nums.get(fn_ea) + gen = self._name_gen + if hit is not None and hit[1] == gen: + return hit[0] + try: + r = self.client.invoke("pc_nums", addr=hex(fn_ea)) + except Exception: # noqa: BLE001 -- an older worker hasn't got the tool + r = {} + out: dict[int, list[tuple[int, int, str, int, int]]] = {} + for rec in (r or {}).get("nums", []): + try: + out.setdefault(int(rec["line"]), []).append( + (int(rec["x0"]), int(rec["x1"]), str(rec.get("value", "")), + _as_int(rec["ea"]), int(rec.get("opnum", 0)))) + except Exception: # noqa: BLE001 -- skip a malformed row + continue + with self._lock: + self._pc_nums[fn_ea] = (out, gen) + return out + + def pc_num_format(self, fn_ea: int, mode: str = "cycle", line: int = -1, + col: int = -1) -> dict: + """The same, for a number in the DECOMPILATION of ``fn_ea``. + + Hex-Rays keeps number formats of its own, per (address, operand) — the + listing's format doesn't reach the pseudocode and vice versa, so this is + a separate call rather than a flag on ``op_format``. + """ + r = self.client.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode), + line=int(line), col=int(col)) + res = r if isinstance(r, dict) else {} + if res.get("error"): + raise IDAToolError("pc_num_format", f"@ {fn_ea:#x}: {res['error']}") + if not res: + raise IDAToolError("pc_num_format", f"@ {fn_ea:#x}: no answer") + return res + def region_label(self, ea: int) -> str: """Display name for a non-function address (segment-qualified).""" try: @@ -1427,15 +1887,6 @@ class Program: sec = None return f"{sec} @ {ea:#x}" if sec else f"<no function> @ {ea:#x}" - @staticmethod - def _fetch_output(url: str, timeout: float = 15.0): - """GET the server's cached full-output blob (plain HTTP, not MCP).""" - try: - with urllib.request.urlopen(url, timeout=timeout) as r: - return json.loads(r.read().decode("utf-8", "replace")) - except Exception: # noqa: BLE001 -- fall back to the truncated preview - return None - def strings(self, min_len: int = 4, refresh: bool = False) -> list[StrLit]: """Every string literal in the binary (IDA's Shift+F12 list), paged in full and cached. ``[]`` if the tool is unavailable.""" @@ -1448,7 +1899,7 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.call( + payload = self.client.invoke( "list_strings", offset=offset, count=page, min_len=min_len, refresh=(refresh and offset == 0)) except IDAToolError: @@ -1473,13 +1924,13 @@ class Program: def linkage(self) -> tuple[list[Linkage], list[Linkage]]: """``(imports, exports)`` for this binary, cached. ``([], [])`` if the - tool is unavailable — an old worker must not break the caller.""" + operation is unavailable — an alternate client must not break the caller.""" with self._lock: hit = self._linkage if hit is not None: return hit try: - payload = self.client.call("list_linkage", kind="both") + payload = self.client.invoke("list_linkage", kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): @@ -1499,6 +1950,84 @@ class Program: self._linkage = out return out + def annotations(self, limit: int = 4000) -> tuple[list["Comment"], list["NamedItem"]]: + """``(comments, names)`` -- everything a person added to this database. + + Not cached: it is the *current* state of your work, and the one caller + (the findings export) asks for it once. ``([], [])`` if the backend has + no such operation, so an alternate client degrades instead of breaking. + """ + try: + payload = self.client.invoke("list_annotations", limit=int(limit)) + except IDAToolError: + return ([], []) + if not isinstance(payload, dict): + return ([], []) + comments = [ + Comment(addr=_as_int(r.get("addr", 0)), text=str(r.get("text", "")), + repeatable=bool(r.get("repeatable")), + whole_func=bool(r.get("whole_func")), + line=str(r.get("line", "") or ""), + seg=str(r.get("seg", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None)) + for r in payload.get("comments", []) if isinstance(r, dict) and r.get("text")] + names = [ + NamedItem(addr=_as_int(r.get("addr", 0)), name=str(r.get("name", "")), + is_func=bool(r.get("func")), size=int(r.get("size", 0) or 0), + proto=(r.get("proto") or None), seg=str(r.get("seg", "") or "")) + for r in payload.get("names", []) if isinstance(r, dict) and r.get("name")] + return (comments, names) + + + def search(self, query: str, mode: str = "text", *, limit: int = 500, + regex: bool = False, case: bool = False, + ) -> tuple[list["SearchHit"], str | None, bool]: + """Search the whole database. Returns ``(hits, error, truncated)``. + + A failed search is DATA (a message to show), not an exception: a bad + regex or an unparsable byte pattern is something the user typed, and + the palette wants to say so without unwinding. + """ + op = "search_bytes" if mode == "bytes" else "search_text" + args: dict = {"limit": int(limit), "case": bool(case)} + if mode == "bytes": + # Validate HERE, not just in the UI: IDA's find_bytes answers a + # malformed pattern with zero hits and no error, which reads as + # "not present" -- the most misleading answer a search can give. + from .search import normalise_pattern, pattern_problem + problem = pattern_problem(query) + if problem: + return ([], problem, False) + args["pattern"] = normalise_pattern(query) + else: + args["query"] = query + args["regex"] = bool(regex) + try: + payload = self.client.invoke(op, **args) + except IDAToolError as e: + return ([], str(e), False) + if not isinstance(payload, dict): + return ([], "the backend returned nothing searchable", False) + hits = [ + SearchHit(addr=_as_int(r.get("addr", 0)), + head=_as_int(r.get("head", r.get("addr", 0))), + line=str(r.get("line", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") + else None), + seg=str(r.get("seg", "") or "")) + for r in payload.get("hits", []) if isinstance(r, dict)] + return (hits, payload.get("error") or None, bool(payload.get("truncated"))) + + def journal_get(self) -> str: + """The findings journal blob stored in this database ('' if none).""" + payload = self.client.invoke("journal_get") + return str(payload.get("data", "")) if isinstance(payload, dict) else "" + + def journal_put(self, data: str) -> None: + self.client.invoke("journal_put", data=str(data)) + def decomp_map(self, ea: int) -> list[list[int]]: """Per-pseudocode-line instruction coverage for the split-view region highlight: a list aligned to the decompiled lines, each the EAs the @@ -1510,7 +2039,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("decomp_map", addr=hex(ea)) + payload = self.client.invoke("decomp_map", addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] @@ -1520,16 +2049,135 @@ class Program: self._decomp_maps[ea] = (out, gen) return out + # -- control-flow graph ------------------------------------------------ # + def flowchart(self, ea: int) -> "Flowchart | None": + """The basic-block CFG of the function containing ``ea``, with each + block's listing rows attached. + + Two calls, not one per block: ``flowchart`` for the shape, then a single + ``heads`` walk over the function's extent which is sliced up by address. + A hundred blocks would otherwise be a hundred round trips. + + Cached per function + item generation, so it survives cursor movement + but not an edit that changes the code. + """ + fn = self.function_of(ea) + key = fn.addr if fn else ea + with self._lock: + hit = self._flowcharts.get(key) + gen = self._name_gen + if hit is not None and hit[1] == gen: + return hit[0] + try: + payload = self.client.invoke("flowchart", addr=hex(ea)) + except IDAToolError: + return None + if not isinstance(payload, dict) or payload.get("error"): + return None + raw = payload.get("blocks") or [] + if not raw: + return None + blocks = [] + for b in raw: + try: + blocks.append(BasicBlock( + id=int(b["id"]), start=_as_int(b["start"]), + end=_as_int(b["end"]), + succs=[(int(d), str(k)) for d, k in (b.get("succs") or [])])) + except (KeyError, ValueError, TypeError): + continue + if not blocks: + return None + f = payload.get("func") or {} + lo = min(b.start for b in blocks) + rows = self._block_rows(blocks) + eas = [h.ea for h in rows] + for b in blocks: + # bisect, not a scan per block: a 400-block function against a few + # thousand rows is a million comparisons done for nothing. + b.rows = rows[bisect.bisect_left(eas, b.start): + bisect.bisect_left(eas, b.end)] + fcv = Flowchart( + func_ea=_as_int(f.get("addr", lo)), + name=str(f.get("name") or f"sub_{lo:X}"), + entry=int(payload.get("entry", 0) or 0), + blocks=blocks, + ) + with self._lock: + self._flowcharts[key] = (fcv, gen) + return fcv + + #: Bytes of padding between two blocks that are still worth fetching in one + #: call. Alignment gaps are a few bytes; a function chunk is far away. + _BLOCK_GAP = 256 + + def _block_rows(self, blocks: list[BasicBlock]) -> list[Head]: + """Listing rows covering ``blocks``, address-ordered. + + Fetches the blocks' merged extents, NOT their convex hull. IDA puts a + function's cold/tail chunks a long way from its entry, so the hull of a + 1.4 KB function can be 680 KB wide: walking it fetched 128 000 listing + rows and took three seconds to draw a graph, all but 300 of them thrown + away immediately. Adjacent blocks coalesce, so an ordinary contiguous + function is still exactly one call. + """ + spans: list[list[int]] = [] + for start, end in sorted((b.start, b.end) for b in blocks): + if spans and start <= spans[-1][1] + self._BLOCK_GAP: + if end > spans[-1][1]: + spans[-1][1] = end + else: + spans.append([start, end]) + out: list[Head] = [] + for start, end in spans: + out.extend(self._heads_between(start, end)) + return out + + def _heads_between(self, lo: int, hi: int) -> list[Head]: + """Listing rows for [lo, hi), paged. Same tool and same ``Head`` shape + the listing view renders, so the graph inherits IDA's colour tags and + operand marks for free.""" + out: list[Head] = [] + addr = lo + for _ in range(64): # bounded: ~128k heads + if addr >= hi: + break + payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), + count=2000) + rows = payload.get("heads", []) if isinstance(payload, dict) else [] + if not rows: + break + for r in rows: + try: + h = Head.from_raw(r) + except (KeyError, ValueError, TypeError): + continue + # Banners and separators are listing furniture; a box already + # has a border and a label of its own. + if h.kind in ("sep", "funchdr"): + continue + if lo <= h.ea < hi: + out.append(h) + cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} + nxt = cur.get("next") + if nxt is None or cur.get("done"): + break + n = _as_int(nxt) + if n <= addr: + break + addr = n + return out + # -- cross-references & containing function --------------------------- # def function_of(self, ea: int) -> Func | None: """Return the function containing ``ea`` (resolves mid-function addrs).""" - payload = self.client.call("lookup_funcs", queries=[hex(ea)]) + payload = self.client.invoke("lookup_funcs", queries=[hex(ea)]) res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None return Func.from_raw(fn) if fn else None def xrefs_from(self, ea: int) -> list[Xref]: - payload = self.client.call( + payload = self.client.invoke( "xref_query", queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) @@ -1541,9 +2189,9 @@ class Program: try: # xref_types adds a fine-grained `kind` (call/read/write/...) for the # xref dialog; fall back to xref_query (code/data only) if absent. - payload = self.client.call("xref_types", queries=q) + payload = self.client.invoke("xref_types", queries=q) except IDAToolError: - payload = self.client.call("xref_query", queries=q) + payload = self.client.invoke("xref_query", queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -1561,17 +2209,17 @@ class Program: # (loc_/locret_): lookup_funcs would map a label to its *containing* # function's entry, so double-clicking a label jumped to the wrong place. try: - payload = self.client.call("resolve_names", queries=[s]) + payload = self.client.invoke("resolve_names", queries=[s]) res = payload.get("result", []) if isinstance(payload, dict) else [] ea = res[0].get("ea") if res and isinstance(res[0], dict) else None if ea: return _as_int(ea) except IDAToolError: - pass # older server without resolve_names -> fall back below + pass # alternate client without resolve_names -> fall back below # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.call("lookup_funcs", queries=[s]) + payload = self.client.invoke("lookup_funcs", queries=[s]) except IDAToolError as e: raise KeyError(f"cannot resolve {target!r}: {e}") from e res = payload.get("result", []) if isinstance(payload, dict) else [] @@ -1609,7 +2257,7 @@ class Program: """Set (empty text clears) the comment at ``ea``; affects both the disasm and decompiler views. Returns the raw payload so the caller can surface a soft per-item error. The caller must invalidate/recompile to see it.""" - return self.client.call("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: diff --git a/idatui/drive.py b/idatui/drive.py index b6fa641..6e5c21a 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -64,7 +64,7 @@ def _fmt_where(st: dict) -> str: ea = fn.get("ea") loc = f"{name} @ {ea:#x}" if isinstance(ea, int) else "(none)" extra = "" - if cur.get("kind") in ("decomp", "disasm"): + if cur.get("kind") in ("decomp", "listing"): extra = f" L{cur.get('line')} C{cur.get('col')} word={cur.get('word')!r}" elif cur.get("kind") == "hex": extra = f" va={cur.get('va'):#x}" if isinstance(cur.get("va"), int) else "" @@ -87,10 +87,10 @@ def cmd_go(c, args): def _show_view(c, want): """Make the requested code pane the visibly-active view (best effort). - Tab toggles disasm<->decomp, and leaves hex back to the preferred code + Tab toggles listing<->decomp, and leaves hex back to the preferred code view; so at most two toggles reach either code view from any state. If - the decompiler fails for the current function the view falls back to - disasm and we simply stop — the caller still returns its text as before. + the decompiler fails for the current function the view falls back to the + listing and we simply stop — the caller still returns its text as before. """ for _ in range(2): if c.call("state").get("active") == want: @@ -121,7 +121,8 @@ def cmd_pc(c, args): lines = d["code"].splitlines() if needle: nlow = needle.lower() - lines = [f"{i:4} {l}" for i, l in enumerate(lines) if nlow in l.lower()] + lines = [f"{i:4} {line}" for i, line in enumerate(lines) + if nlow in line.lower()] return "\n".join(lines) or f"(no line matches {needle!r})" return d["code"] @@ -132,7 +133,7 @@ def cmd_dis(c, args): # Drive the real UI so viewers see the disassembly, not just the driver. if target is not None: c.call("goto", target=target, delay_ms=0) - _show_view(c, "disasm") + _show_view(c, "listing") d = c.call("disassembly", target=target, max=n) return "\n".join(f"{ln['ea']:#010x} {ln['text']}" for ln in d.get("lines", [])) @@ -219,10 +220,15 @@ def cmd_mv(c, args): def cmd_note(c, args): if len(args) < 2: raise SystemExit("usage: note <fn> <text...>") - c.call("goto", target=args[0], delay_ms=0) - c.call("cursor", line=0, col=0) + st = c.call("goto", target=args[0], delay_ms=0) + # goto already lands on the function's first line. The old `cursor line=0` + # meant "the top of the function" only in the decompiler; in the listing + # line 0 is the top of the whole SEGMENT, so the note landed at address 0 -- + # and on a 42k-line firmware listing the scroll to get there timed the + # caller out, which read as "comments are broken". c.call("comment", text=" ".join(args[1:]), delay_ms=0) - return f" noted {args[0]}" + cur = (st.get("function") or {}).get("name") or args[0] + return f" noted {cur} @ {(st.get('cursor') or {}).get('ea', 0):#x}" def cmd_retype(c, args): @@ -233,11 +239,86 @@ def cmd_retype(c, args): return _fmt_where(st) +def cmd_define(c, args): + """define <kind> [target ...] — the raw-image workflow (thumb/code/func). + + Several targets are common on a firmware image (a list of entry points from + a symbol file), so take them all and report per-target. + """ + if not args: + raise SystemExit("usage: define <code|func|undef|thumb|thumbscan|data|" + "string> [target ...]") + kind, targets = args[0], (args[1:] or [None]) + out = [] + for t in targets: + try: + st = c.call("define", kind=kind, **({"target": t} if t else {})) + out.append(f" {t or '.'}: {st.get('status', '')}") + except RpcError as e: + out.append(f" {t or '.'}: FAILED: {e}") + return "\n".join(out) + + +def cmd_fmt(c, args): + """fmt [mode] [word] — how the literal under the cursor is DISPLAYED. + + ``fmt`` alone cycles (IDA's 'o'); a mode name sets it outright. A trailing + word puts the cursor on that token first, so you can name the literal + instead of steering the column there. + + fmt # cycle the literal under the cursor + fmt dec # show it in decimal + fmt hex 18h # find '18h' on screen, then make it hex + """ + mode = args[0] if args else "cycle" + params = {"mode": mode} + if len(args) > 1: + params["word"] = args[1] + st = c.call("opfmt", **params) + return " " + (st.get("opfmt", {}).get("status") or st.get("status", "")) + + +def cmd_syms(c, args): + """syms <file.json> — bulk-apply a symbol file ([{addr|start|ea, name}]).""" + if len(args) != 1: + raise SystemExit("usage: syms <symbols.json>") + r = c.call("rename_many", file=os.path.abspath(os.path.expanduser(args[0]))) + m = r.get("rename_many", {}) + out = [f" {m.get('ok', 0)}/{m.get('requested', 0)} renamed" + f" (skipped {m.get('skipped', 0)}, failed {m.get('failed', 0)})"] + for e in m.get("errors", []): + out.append(f" {e.get('addr')}: {e.get('error')}") + return "\n".join(out) + + def cmd_save(c, args): c.call("save") return " saved" +def cmd_find(c, args): + """find <query...> -- search the database; bytes if it looks like bytes.""" + if not args: + raise SystemExit("usage: find <text | 48 8b ?? c3 | hex:...>") + r = c.call("find", query=" ".join(args)) + hits = r.get("hits", []) + out = [f" [{r.get('mode')}] {len(hits)}{'+' if r.get('truncated') else ''} hits"] + for h in hits[:40]: + out.append(f" {h['addr']} {(h.get('func') or h.get('seg') or ''):<20.20} " + f"{h.get('line', '')}") + if len(hits) > 40: + out.append(f" … {len(hits) - 40} more") + return "\n".join(out) + + +def cmd_export(c, args): + """export [path] -- write the session's findings as markdown.""" + r = c.call("export", **({"path": args[0]} if args else {})) + return (f" {r.get('path')} ({r.get('bytes', 0)} bytes: " + f"{r.get('comments', 0)} comments, {r.get('names', 0)} names, " + f"{r.get('types', 0)} types)") + + def cmd_screen(c, args): return c.call("screen").get("text", "") @@ -256,7 +337,9 @@ COMMANDS = { "where": cmd_where, "go": cmd_go, "pc": cmd_pc, "dis": cmd_dis, "callees": cmd_callees, "callers": cmd_callers, "names": cmd_names, "rename": cmd_rename, "mv": cmd_mv, "note": cmd_note, "retype": cmd_retype, - "save": cmd_save, "screen": cmd_screen, "raw": cmd_raw, + "save": cmd_save, "screen": cmd_screen, "raw": cmd_raw, "define": cmd_define, + "syms": cmd_syms, "fmt": cmd_fmt, "export": cmd_export, + "find": cmd_find, "binaries": cmd_binaries, "switch": cmd_switch, } diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py new file mode 100644 index 0000000..52566ca --- /dev/null +++ b/idatui/edit_ctl.py @@ -0,0 +1,738 @@ +"""Everything that writes to the database: rename, comment, retype, define. + +Six edits with the same shape -- work out what the cursor is on, ask for a +value, apply it in a worker, then invalidate the right caches and put the view +back where it was -- and that last step is the one that is easy to get subtly +wrong. The rules are collected here rather than rediscovered per edit: + +* An edit that only changes *names* (rename, comment, retype, literal format) + goes through :meth:`reload_active_code`, which reloads in place from indices. +* An edit that changes *item structure* (make data, define code/func/string, + undefine) goes through a :class:`~idatui.app.ViewAnchor`, because row indices + do not survive it -- defining code collapses four undefined byte rows into + one instruction row. +* Every one of them bumps a cache generation, sets ``_dirty``, and hands its + message over as a flash rather than writing it directly, because the reload + it just triggered will write its own status afterwards. + +The message handlers and the ``@work`` entry points stay on ``IdaTui``: Textual +dispatches ``on_<message>`` by name on the DOMNode, and its worker machinery +wants a DOMNode host. They are one-line delegates into here. +""" +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +from textual.widgets import DataTable + +from . import diag +from .errors import IDAToolError + +if TYPE_CHECKING: # pragma: no cover + from .app import IdaTui + +_app_mod = None + + +def _M(): + """The widget classes, imported lazily to avoid a cycle with app.py.""" + global _app_mod + if _app_mod is None: + from . import app as _m + _app_mod = _m + return _app_mod + + +#: A C type wide enough for N bytes, for prefilling a retype/define prompt. +_BY_SIZE = {1: "unsigned __int8", 2: "unsigned __int16", + 4: "unsigned __int32", 8: "unsigned __int64"} + + +class EditController: + """The database edits, and the bookkeeping each one owes afterwards.""" + + def __init__(self, app: "IdaTui") -> None: + self.app = app + + # -- shared aftermath --------------------------------------------------- # + def reload_active_code(self) -> None: + """Refresh whichever code view is showing after an edit (comment/rename/ + retype), in place: re-decompile if in the decompiler, else reload the + listing.""" + app = self.app + M = _M() + cur = app._cur + if cur is None: + return + if app.is_decomp: + # Snapshot the LIVE pseudocode position before forcing a recompile. + # dec_scroll_y isn't tracked on every move, so without this the reload + # falls into show()'s derive path (a bare scroll_to) and leaves a + # stale frame until the next cursor move; capturing the real scroll + # makes show() take the robust _apply_scroll path and repaint now. + dec = app.query_one(M.DecompView) + if cur.ea == dec.loaded_ea: + cur.dec_cursor = dec.cursor + cur.dec_cursor_x = dec.cursor_x + cur.dec_scroll_y = round(dec.scroll_offset.y) + cur.dec_scroll_x = round(dec.scroll_offset.x) + dec.loaded_ea = None # force re-decompile + app._show_active() + else: + # Capture the LIVE position from the widget (the source of truth) + # rather than trusting nav-entry tracking, which goes stale. Capture + # it as ADDRESSES via the anchor: bump_names() discards the segment + # model so the reload rebuilds it, and an edit that changes how many + # rows an item takes makes the old indices point somewhere else. + # Index capture is CORRECT here and an anchor is not: a rename or + # comment doesn't change how many rows anything takes, and the model + # this rebuilds is constructed empty — index_of_ea on it returns -1 + # until pages load, so an anchor would resolve to nothing while + # costing an extra model build on the UI thread. Address anchoring is + # for the edit paths that DO change row structure (see do_edit_item). + lst = app.query_one(M.ListingView) + cur.view = "listing" + if lst.model is not None: + cur.cursor = lst.cursor + cur.cursor_x = lst.cursor_x + cur.scroll_y = round(lst.scroll_offset.y) + app._open_entry(cur, push=False) + + def edit_done(self, anchor) -> None: # type: ignore[no-untyped-def] + """One place where an edit's aftermath is settled. + + The reload this edit triggered will write its own status when it lands — + after this — so the message is handed over as a flash rather than + written and lost. + """ + app = self.app + app._dirty = True + if anchor.flash: + app._status(anchor.flash, priority=True) + if anchor.refresh_functions: + # Creating (or destroying) a function changes the index that the + # names pane, Ctrl+N and the "no functions" hint all read. Without + # this, `p` gave you a function the rest of the app couldn't see. + app._reindex_functions() + + def _rename_index(self, addr: int, new: str) -> None: + """Point the function index, the nav history and the names table at a + function's new name. All three are caches of it, and a rename that + updates only some of them is how `functions`/resolve/the palette end up + reporting that the rename never happened.""" + app = self.app + if app._func_index is not None: + app._func_index.update_name(addr, new) + for e in app._nav: + if e.ea == addr: + e.name = new + # Update the one cell in place (a full rebuild would race the initial + # streaming load and duplicate row keys). + try: + table = app.query_one("#func-table", DataTable) + name_col = list(table.columns.keys())[1] + table.update_cell(str(addr), name_col, new) + except Exception: # noqa: BLE001 -- row filtered out / not yet streamed + pass + + # -- rename (IDA 'n') --------------------------------------------------- # + @staticmethod + def is_pseudocode_label(view, name: str) -> bool: + """True if ``name`` is a Hex-Rays goto label in ``view``. The rename tool + has no label category (only func/global/local/stack), so renaming one + fails with a misleading 'local variable not found'; detect it up front + and explain instead. A label is the default ``LABEL_n`` or any token used + as a ``goto`` target.""" + if not isinstance(view, _M().DecompView): + return False + if re.fullmatch(r"LABEL_\d+", name): + return True + body = "\n".join(getattr(view, "_texts", []) or []) + return re.search(rf"\bgoto\s+{re.escape(name)}\b", body) is not None + + def request_rename(self, msg) -> None: # type: ignore[no-untyped-def] + app = self.app + M = _M() + # In the flat listing, 'n' names the ADDRESS under the cursor (create a + # label), not a symbol-by-name. This is what lets you name a bare/ + # undefined byte — e.g. the free byte at addr+1 after shrinking a u16 to + # a u8 — which the word-under-cursor path can't do (no symbol to rename). + if isinstance(msg.view, (M.ListingView, M.GraphView)): + ea = msg.view._cursor_ea() + if ea is None: + app._status("no address on this line to name") + return + head = msg.view.cur_head() + word = msg.view.word_under_cursor() + mnem = head.text.split(" ", 1)[0] if (head and head.text) else "" + # If the cursor is on a symbol token (a call/branch target, a data + # reference, or this head's own label) rename THAT symbol; otherwise + # create/rename a label at the head's address (bare/undefined bytes). + if (word and app._looks_like_symbol(word) and word != mnem + and word.lower() not in M._ASM_KEYWORDS): + app.prompts.rename.show( + f"rename '{word}' — Enter=apply Esc=cancel", + word, ctx=(msg.view, word, None)) + else: + cur = head.name if (head is not None and head.name) else "" + app.prompts.rename.show( + f"name @ {ea:#x} — Enter=apply Esc=cancel", + cur, ctx=(msg.view, cur, ea)) + return + if not msg.name: + app._status("nothing to rename under the cursor") + return + if self.is_pseudocode_label(msg.view, msg.name): + app._status( + f"can't rename pseudocode label '{msg.name}' " + "(Hex-Rays goto labels aren't renamable via the API)") + return + app.prompts.rename.show( + f"rename '{msg.name}' — Enter=apply Esc=cancel", + msg.name, ctx=(msg.view, msg.name, None)) + + def submit_rename(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] + view, old, addr = ctx + if addr is not None: # listing: name this address (create a label) + if value and value != old: + self.app._do_name_addr(addr, value) + return + if view is not None and value and value != old: + self.app._do_rename(view, old, value) + + def do_rename(self, view, old: str, new: str) -> None: # worker context + app = self.app + assert app.program is not None + prog, cur = app.program, app._cur + kind = "data" + addr: int | None = None + batch: dict = {"data": {"old": old, "new": new}} + resolved: int | None = None + try: + resolved = prog.resolve(old) + except Exception as e: # noqa: BLE001 + # Not cosmetic: an unresolved name is renamed as DATA instead of as + # a function, so a lookup that failed for a transport reason quietly + # applies the wrong kind of edit. + diag.note(f"rename: resolve({old!r})", e) + resolved = None + if resolved is not None: + fn = prog.function_of(resolved) + if fn is not None and fn.addr == resolved: + kind, addr = "func", resolved + batch = {"func": {"addr": hex(resolved), "name": new}} + else: + kind, batch = "data", {"data": {"old": old, "new": new}} + elif isinstance(view, _M().DecompView) and cur is not None: + dec = prog.decompile(cur.ea) + ref = next((r for r in dec.refs if r.name == old), None) + if ref is not None: + fn = prog.function_of(ref.addr) + if fn is not None and fn.addr == ref.addr: + kind, addr = "func", ref.addr + batch = {"func": {"addr": hex(ref.addr), "name": new}} + else: + kind, batch = "data", {"data": {"old": old, "new": new}} + else: + kind = "local" + batch = {"local": {"func_addr": hex(cur.ea), "old": old, "new": new}} + elif cur is not None: # disasm view + if old.startswith(("var_", "arg_")): + kind = "stack" + batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} + try: + res = prog.client.invoke("rename", batch=batch) + except IDAToolError as e: + app.call_from_thread(app._status, f"rename failed: {e.message}") + return + summary = res.get("summary", {}) if isinstance(res, dict) else {} + if not (summary.get("ok", 0) > 0 and summary.get("failed", 0) == 0): + msg = "rename failed" + for catk in ("func", "data", "local", "stack"): + items = res.get(catk) if isinstance(res, dict) else None + if isinstance(items, list) and items and items[0].get("error"): + msg = f"rename failed: {items[0]['error']}" + app.call_from_thread(app._status, msg) + return + app.call_from_thread(self.after_rename, kind, addr, old, new) + + def after_rename(self, kind: str, addr: int | None, old: str, new: str) -> None: + app = self.app + # A renamed symbol can appear in many functions, so invalidate globally; + # each function refreshes its names the next time it's viewed. + app.program.bump_names() + self.reload_active_code() + if kind == "func" and addr is not None: + self._rename_index(addr, new) + app._dirty = True + app.journal.record("rename", addr, f"{old} → {new}", {"kind": kind}) + app._status(f"renamed {old} → {new} (Ctrl+S to save)") + + def do_name_addr(self, addr: int, name: str) -> None: # worker context + """Set a label at ``addr`` (listing 'n'). Works on a bare/undefined byte + — unlike the symbol-by-name path, this names the address directly.""" + app = self.app + assert app.program is not None + try: + res = app.program.client.invoke( + "rename", batch={"data": {"addr": hex(addr), "new": name}}) + except IDAToolError as e: + app.call_from_thread(app._status, f"name failed: {e.message}") + return + summary = res.get("summary", {}) if isinstance(res, dict) else {} + if not (summary.get("ok", 0) > 0 and summary.get("failed", 0) == 0): + err = "name failed" + items = res.get("data") if isinstance(res, dict) else None + if isinstance(items, list) and items and items[0].get("error"): + err = f"name failed: {items[0]['error']}" + app.call_from_thread(app._status, err) + return + # The label shows in the listing's head rows -> invalidate + reopen. + app.program.bump_items(addr) + # Naming the address *of a function start* is a function rename by any + # other name. Without this the cached index kept the old name, so + # `functions`/`names`/resolve/the palette all reported the rename had + # not happened -- and a driver that trusts those readbacks redoes work + # it already did. + try: + fn = app.program.function_of(addr) + except Exception as e: # noqa: BLE001 + # If this throws we don't learn the address IS a function start, so + # the index keeps the old name and every readback says the rename + # never happened. + diag.note(f"name: function_of({addr:#x})", e) + fn = None + is_func_start = fn is not None and fn.addr == addr + lm = app.program.listing(addr) + label = name if is_func_start else app.program.region_label(addr) + idx = max(lm.ensure_ea(addr), 0) if lm is not None else 0 + app.call_from_thread(self.open_at_named, label, addr, idx, name, + is_func_start) + + def open_at_named(self, label: str, addr: int, idx: int, name: str, + is_func_start: bool = False) -> None: + app = self.app + if is_func_start: + app.program.bump_names() + self._rename_index(addr, name) + app._open_at(addr, label, idx, False, -1, 0, True) + app._dirty = True + app.journal.record("rename", addr, f"→ {name}") + app._status(f"named {addr:#x} → {name} (Ctrl+S to save)") + + # -- comments (IDA ';') ------------------------------------------------- # + @staticmethod + def existing_comment(view) -> str: + """Current line comment (for prefill), parsed from the rendered text. In + pseudocode a comment is `// text` before the trailing /*0xEA*/ markers; + C has no `//` operator, so the last `//` is unambiguously the comment.""" + if isinstance(view, _M().DecompView) and 0 <= view.cursor < len(view._texts): + s = re.sub(r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", + view._texts[view.cursor]) + i = s.rfind("//") + return s[i + 2:].strip() if i >= 0 else "" + return "" + + def request_comment(self, msg) -> None: # type: ignore[no-untyped-def] + app = self.app + ea = app._line_ea_for(msg.view) + # Signature / local-declaration lines carry no address; fall back to the + # function's entry ea so commenting the header annotates the function. + func_level = ea is None + if func_level: + ea = app._cur.ea if app._cur else None + if ea is None: + app._status("no address on this line to comment") + return + existing = "" if func_level else self.existing_comment(msg.view) + what = "function comment" if func_level else "comment" + app.prompts.comment.show( + f"{what} @ {ea:#x} — Enter=apply (empty=clear) Esc=cancel", + existing, ctx=(msg.view, ea, existing)) + + def submit_comment(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] + view, ea, existing = ctx + if view is not None and value != existing: # empty value clears it + self.app._do_comment(ea, value) + + def do_comment(self, ea: int, text: str) -> None: # worker context + app = self.app + assert app.program is not None + # The prompt is single-line, so a literal '\n' (backslash-n) means a real + # newline — Hex-Rays renders each as its own '//' line. Lets long notes + # wrap instead of running off the right edge and clipping. + text = text.replace("\\n", "\n") + try: + res = app.program.set_comment(ea, text) + except IDAToolError as e: + app.call_from_thread(app._status, f"comment failed: {e.message}") + return + data = res.get("result") if isinstance(res, dict) else None + if (isinstance(data, list) and data and isinstance(data[0], dict) + and data[0].get("error")): + app.call_from_thread(app._status, + f"comment failed: {data[0]['error']}") + return + app.call_from_thread(self.after_comment, ea, text) + + def after_comment(self, ea: int, text: str) -> None: + app = self.app + # A comment shows in both views but only after Hex-Rays recompiles, so + # reuse the name-generation invalidation (bumps gen -> decompile is + # force_recompiled lazily; disasm/listing caches are cleared). + app.program.bump_names() + self.reload_active_code() + app._dirty = True + app.journal.record("comment" if text else "uncomment", ea, text) + verb = "cleared comment" if not text else "commented" + app._status(f"{verb} @ {ea:#x} (Ctrl+S to save)") + + # -- retype (set type, IDA 'y') ---------------------------------------- # + @staticmethod + def guess_data_type(size: int) -> str: + """A sensible prefill when a global carries no type yet.""" + return _BY_SIZE.get(size, f"char[{size}]" if size > 0 else "void *") + + def prepare_retype(self, view, word: str | None) -> None: # worker context + """Work out whether the cursor is on a local variable or a function, and + fetch the current type/prototype to prefill the prompt.""" + app = self.app + assert app.program is not None and app._cur is not None + ft = app.program.func_types(app._cur.ea) + kind: str | None = None + subject: int = app._cur.ea + prefill = "" + # 1) a local variable (or arg) of the current function + if word and ft is not None: + lv = next((v for v in ft.lvars if v.name == word), None) + if lv is not None: + kind, prefill = "lvar", lv.type + # 2) a symbol under the cursor: a function (retype its prototype) or a + # global/data item (retype the variable). Without the data case a + # global fell through to (3) and silently retyped the ENCLOSING + # function's prototype instead. + if kind is None and app._looks_like_symbol(word): + try: + tgt = app.program.resolve(word) + except Exception as e: # noqa: BLE001 + # Falls through to case (3), which retypes the enclosing + # function -- a different edit from the one asked for. + diag.note(f"retype: resolve({word!r})", e) + tgt = None + if tgt is not None: + tft = app.program.func_types(tgt) + if tft is not None: + kind, subject, prefill = "func", tgt, tft.prototype + else: + dt = app.program.data_type(tgt) + if dt is not None and not dt.get("is_func"): + kind, subject = "data", tgt + prefill = dt.get("type") or self.guess_data_type( + dt.get("size") or 0) + # 3) fall back to the current function itself + if kind is None and ft is not None: + kind, subject, prefill = "func", app._cur.ea, ft.prototype + if kind is None: + app.call_from_thread(app._status, + "nothing to retype under the cursor") + return + app.call_from_thread(self.open_retype, view, kind, subject, + word or "", prefill) + + def open_retype(self, view, kind: str, subject: int, word: str, + prefill: str) -> None: # type: ignore[no-untyped-def] + label = "prototype" if kind == "func" else f"type for '{word}'" + self.app.prompts.retype.show(f"{label} — Enter=apply Esc=cancel", + prefill, ctx=(view, kind, subject, word)) + + def submit_retype(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] + view, kind, subject, word = ctx + if view is not None and value: + self.app._do_retype(kind, subject, word, value) + + def do_retype(self, kind: str, subject: int, word: str, + new: str) -> None: # worker context + app = self.app + assert app.program is not None + if kind == "func": + err = app.program.set_function_type(subject, new) + elif kind == "data": # a global / data item referenced in the body + err = app.program.set_data_type(subject, new) + else: # lvar of the current function + err = app.program.set_lvar_type(app._cur.ea, word, new) + if err: + app.call_from_thread(app._status, f"retype failed: {err}") + return + app.call_from_thread(self.after_retype, kind, word) + + def after_retype(self, kind: str, word: str) -> None: + app = self.app + # A type change alters the pseudocode (and disasm operand types), so + # recompile via the name-generation invalidation and reopen in place. + app.program.bump_names() + self.reload_active_code() + app._dirty = True + app.journal.record("retype", getattr(app._cur, "ea", None), word, + {"kind": kind}) + what = "prototype" if kind == "func" else f"'{word}'" + app._status(f"retyped {what} (Ctrl+S to save)") + + # -- typed data definition (make_data, IDA 'd') ------------------------ # + @staticmethod + def default_data_type(head) -> str: # type: ignore[no-untyped-def] + """A sensible prefill C type for defining data over ``head``.""" + sz = getattr(head, "size", 0) or 0 + return _BY_SIZE.get(sz, f"char[{sz}]" if sz > 0 else "unsigned __int8") + + def request_make_data(self, msg) -> None: # type: ignore[no-untyped-def] + app = self.app + view = msg.view + is_listing = isinstance(view, _M().ListingView) + ea = view._cursor_ea() if is_listing else None + if ea is None: + app._status("no address on this line to define data") + return + head = view.cur_head() if is_listing else None + app.prompts.makedata.show( + f"data type @ {ea:#x} (e.g. int, char[16], my_struct)" + " — Enter=apply Esc=cancel", + self.default_data_type(head) if head is not None else "int", + ctx=(view, ea)) + + def submit_make_data(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] + view, ea = ctx + if view is not None and value: + self.app._do_make_data(ea, value, self.app._anchor()) + + def do_make_data(self, ea: int, type_decl: str, + anchor=None) -> None: # worker context + app = self.app + assert app.program is not None + try: + app.program.make_data(ea, type_decl) + except Exception as e: # noqa: BLE001 + diag.note(f"make_data({ea:#x}, {type_decl!r})", e) + app.call_from_thread(app._status, f"make data: {e}") + return + app.program.bump_items(ea) + anchor = anchor or _M().ViewAnchor() + anchor.flash = f"data ({type_decl}) @ {ea:#x} (Ctrl+S to save)" + name = app.program.region_label(ea) + lm = app.program.listing(ea) + idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 + _cur, top = app._anchor_rows(anchor, lm, ea) + app.call_from_thread( + app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app.call_from_thread(self.edit_done, anchor) + + # -- literal display formats (IDA 'o') --------------------------------- # + def request_op_format(self, msg) -> None: # type: ignore[no-untyped-def] + app = self.app + M = _M() + if app.program is None or app._cur is None: + return + view = msg.view + if isinstance(view, M.ListingView): + ea = view._cursor_ea() + if ea is None: + app._status("no address on this line to reformat") + return + head = view.cur_head() + if head is not None and head.kind in ("sep", "funchdr", "label"): + # A banner/label row carries the NEXT item's address so that + # navigation lands somewhere real — but it has no operands of + # its own, and a column measured against it would point into + # that item at random. + app._status("no literal on this line to reformat", priority=True) + return + app._do_op_format(msg.mode, "listing", ea, view.op_col()) + return + if isinstance(view, M.DecompView): + # The pseudocode's formats are keyed on the FUNCTION Hex-Rays + # decompiled, not on the line's own address. + fn = view.loaded_ea if view.loaded_ea is not None else app._cur.ea + app._do_op_format(msg.mode, "decomp", fn, view.cursor_x, view.cursor) + + def do_op_format(self, mode: str, where: str, ea: int, col: int, + line: int = -1) -> None: # worker context + app = self.app + assert app.program is not None + try: + if where == "listing": + r = app.program.op_format(ea, mode=mode, col=col) + what = f"op{r.get('n', 0)} " + else: + r = app.program.pc_num_format(ea, mode=mode, line=line, col=col) + what = "" + # These are the RESULT of a keypress, so they go on the bar with + # priority. Without it a refusal is swallowed by the previous edit's + # flash and both the screen and the RPC snapshot still show the last + # success -- a call that did nothing reads as one that worked. + except IDAToolError as e: + app.call_from_thread(app._status, f"format: {e.message}", True) + return + except Exception as e: # noqa: BLE001 -- surface transport failures too + diag.note(f"op_format({where}, {ea:#x})", e) + app.call_from_thread(app._status, f"format: {e}", True) + return + text = " ".join((r.get("text") or "").split()) + prev, fmt = r.get("prev", ""), r.get("format", "?") + if mode == "show": + # A question, not an edit: say what this literal is and what it + # could be, and leave the database (and the view) alone. + app.call_from_thread( + app._status, + f"{what}{fmt} {r.get('value') or ''}" + f" [{', '.join(r.get('choices', []))}]", True) + return + step = f"{prev} \u2192 {fmt}" if prev and prev != fmt else fmt + desc = f"{what}{step}: {text[:96]}" + if r.get("warn"): + desc += f" \u26a0 {r['warn']}" + # Which literal this was, so the cursor can be put back on it after the + # reload: the line reflows and the old column stops meaning the same + # thing (48 -> 0x30 shifts everything to its right). + keep: tuple | None = None + if where == "listing": + if r.get("n") is not None: + keep = ("listing", int(r["n"])) + elif r.get("ea"): + keep = ("decomp", int(str(r["ea"]), 0), int(r.get("opnum", 0))) + app.call_from_thread(self.after_op_format, desc, keep) + + def after_op_format(self, desc: str, keep: tuple | None = None) -> None: + app = self.app + M = _M() + # Only the rendering changed, but it changed in the database: drop the + # cached rows (and bump the generation, so the decompiler re-runs and + # picks up its own new number format) and reopen where we are. + app.program.bump_names() + if keep is not None: + # Set before the reload: both views consume this one-shot when their + # new content lands, which is always after this handler returns. + if keep[0] == "listing": + app.query_one(M.ListingView)._pending_op = keep[1] + else: + app.query_one(M.DecompView)._keep_lit = (keep[1], keep[2]) + self.reload_active_code() + app._dirty = True + app._status(f"{desc} (Ctrl+S to save)", priority=True) + + # -- item structure edits (IDA c/p/u) ---------------------------------- # + def request_edit_item(self, msg) -> None: # type: ignore[no-untyped-def] + app = self.app + view = msg.view + ea = view._cursor_ea() if isinstance(view, _M().ListingView) else None + if ea is None: + app._status("no address on this line to (re)define") + return + app._do_edit_item(msg.kind, ea, app._anchor()) + + def do_edit_item(self, kind: str, ea: int, + anchor=None) -> None: # worker context + app = self.app + assert app.program is not None + verb = {"code": "defined code", "func": "created function", + "undef": "undefined", "string": "made string", + "thumb": "switched decoding", "thumbscan": "scanned"}[kind] + try: + if kind == "code": + # Keep going until something stops it: one instruction is rarely + # what you want, and on a raw image it means pressing `c` once + # per opcode for the length of a function. + r = app.program.define_code_run(ea) + n, why = int(r.get("count", 0)), r.get("stopped", "") + if n == 0 and why == "defined": + # Already code/data here — a no-op, not a failure. Saying + # "failed to create instruction" for it would be a lie. + app.call_from_thread( + app._status, f"already defined @ {ea:#x}") + return + if n == 0: + raise IDAToolError("define_code", + f"@ {ea:#x}: Failed to create instruction") + end = int(str(r.get("end", hex(ea))), 0) + reason = {"undecodable": "hit bytes that don't decode", + "flow": "control flow ends here", + "defined": "ran into existing code/data", + "segment": "end of segment", + "limit": "instruction limit"}.get(why, why) + verb = (f"defined {n} instruction{'s' if n != 1 else ''} " + f"({ea:#x}\u2013{end:#x}) \u2014 {reason}") + elif kind == "thumbscan": + # A vector table is a list of Thumb entry points that IDA won't + # follow on a headerless image, because nothing tells it those + # words are pointers. Scan from the cursor. + anchor.refresh_functions = True + r = app.program.thumb_scan(ea, ea + 0x400) + n, applied = int(r.get("n", 0)), int(r.get("applied", 0)) + if not n: + verb = (f"no Thumb entry pointers in {ea:#x}\u2013{ea+0x400:#x}" + " (odd words pointing into the image)") + else: + verb = (f"{n} Thumb entr{'y' if n == 1 else 'ies'} found, " + f"{applied} disassembled") + elif kind == "thumb": + # Switch the mode, then disassemble in it: flipping T and + # leaving the bytes undefined shows nothing, and the reason you + # flipped it was to read the code. + r = app.program.set_thumb(ea) + run = app.program.define_code_run(ea) + n = int(run.get("count", 0)) + mode = "Thumb" if r.get("thumb") else "ARM" + verb = f"{mode} @ {ea:#x}" + if r.get("forced_32bit"): + verb += " (segment set to 32-bit; Thumb needs ARM32)" + if r.get("db_64bit"): + # Disassembly will look right and F5 will never work. + verb += (" \u26a0 this database is 64-bit, so Hex-Rays " + "won't decompile it \u2014 Ctrl+L and pick " + "arm:ARMv7-A") + verb += (f" \u2014 {n} instruction{'s' if n != 1 else ''}" + if n else " \u2014 still doesn't decode") + # falls through to the shared reload: same cache bump, same + # anchor restore, same flash. That is the whole point of having + # one path. + elif kind == "func": + anchor.refresh_functions = True + r = app.program.define_func(ea) + if r.get("start") and r.get("end"): + verb = (f"created function {r['start']}\u2013{r['end']}" + + (" (end worked out from the code)" + if r.get("how") == "explicit-end" else "")) + elif kind == "string": + s = app.program.make_string(ea) + verb = f"made string ({s[:24]!r})" if s else verb + else: + # Undefining can destroy a function as easily as `p` creates one. + anchor.refresh_functions = True + app.program.undefine(ea) + except Exception as e: # noqa: BLE001 -- surface soft/hard tool errors + diag.note(f"edit_item({kind}, {ea:#x})", e) + app.call_from_thread(app._status, f"{kind}: {e}") + return + # Structure changed: drop all item/function/decomp caches. The segment + # listing keeps its walk in front of `ea` -- rows before an edit keep + # their addresses and their row numbers. + app.program.bump_items(ea) + # Re-resolve: a define_func upgrades the region to a real function view; + # anything else re-reads the (still function-less) listing in place. + anchor = anchor or _M().ViewAnchor() + anchor.flash = f"{verb} @ {ea:#x} (Ctrl+S to save)" + fn = app.program.function_of(ea) + if fn is not None: + model = app.program.disasm(fn.addr, fn.name) + idx = 0 if ea == fn.addr else model.index_of_ea(ea) + _cur, top = app._anchor_rows(anchor, model, ea) + app.call_from_thread( + app._open_at, fn.addr, fn.name, idx, False, -1, 0, False, + None, top) + else: + name = app.program.region_label(ea) + lm = app.program.listing(ea) + idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 + _cur, top = app._anchor_rows(anchor, lm, ea) + app.call_from_thread( + app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app.call_from_thread(self.edit_done, anchor) diff --git a/idatui/errors.py b/idatui/errors.py index 29b09ae..aaf2dc5 100644 --- a/idatui/errors.py +++ b/idatui/errors.py @@ -1,10 +1,8 @@ -"""Transport-agnostic error hierarchy and the Session model. +"""TUI-facing error hierarchy and lightweight database session model. -These were originally defined in client.py (the ida-pro-mcp HTTP client), but the -idalib worker path (worker_client / domain / app) needs the same exception types -and Session dataclass without dragging in the HTTP transport. They live here so -both backends share one definition; client.py re-exports them for backwards -compatibility with the (deprecated) mcp tooling and the stress tests. +The Code Mode adapter normalizes ``ida_codemode.client`` transport and execution +errors into these types so the domain and Textual layers do not depend on HTTP or +registry implementation details. """ from __future__ import annotations diff --git a/idatui/findings.py b/idatui/findings.py new file mode 100644 index 0000000..58a69f6 --- /dev/null +++ b/idatui/findings.py @@ -0,0 +1,384 @@ +"""Export a reverse-engineering session as a markdown report. + +The output of an RE session is not the database, it is what you *learned* -- +and that lives scattered across comments, names and types inside a `.i64` that +only IDA can read. This turns it into one document you can paste into an +advisory, a writeup or a ticket. + +Two halves, deliberately separated: + +* :func:`gather` talks to a :class:`~idatui.domain.Program` (the only part that + needs IDA) and returns a plain :class:`Findings`. +* :func:`render` turns a :class:`Findings` into markdown and knows nothing about + IDA, so the formatting -- grouping, sorting, escaping, the empty cases -- is + tested offline in ``tests/test_findings.py``. + +**On authorship.** IDA records "this address has a real name" but not *who* +named it, so a stripped binary's report is exactly your renames while a binary +with symbols also lists the ones it shipped with. The report says which case it +is rather than claiming credit; comments and types have no such ambiguity. +""" + +from __future__ import annotations + +import os +import re +import time +from dataclasses import dataclass, field + + +@dataclass +class Findings: + """Everything the report can show, already fetched. Plain data on purpose.""" + + binary: str = "" + path: str = "" + #: (start, end, name) segments, for the overview + sections: list[tuple[int, int, str]] = field(default_factory=list) + n_functions: int = 0 + #: idatui.domain.Comment + comments: list = field(default_factory=list) + #: idatui.domain.NamedItem + names: list = field(default_factory=list) + #: (idatui.domain.Struct, source or "") + types: list[tuple[object, str]] = field(default_factory=list) + #: names IDA supplied from imports/exports -- excluded from "named", since + #: they are the linker's work, not anyone's finding + linked: set[str] = field(default_factory=set) + stripped: bool = True + truncated: bool = False + #: annotations dropped as the loader's own work, reported as a count + skipped_loader: int = 0 + #: Addresses idatui recorded itself editing (idatui/journal.py). When this + #: is non-empty the report is EXACT -- it is what you did, not what the + #: database happens to contain. Empty means nobody journalled this database + #: (worked on in the IDA GUI, or before this feature), and the report falls + #: back to filtering by shape, which it says out loud. + recorded: set[int] = field(default_factory=set) + #: type names the journal saw declared, for the same reason + recorded_types: set[str] = field(default_factory=set) + n_recorded: int = 0 + generated_at: float = field(default_factory=time.time) + + +#: Segments the *loader* owns rather than the program: the ELF/PE header and +#: friends. IDA annotates those itself -- "File format: \x7FELF", "File class: +#: 64-bit", `elf_gnu_hash_nbuckets` -- through the very same set_cmt/set_name +#: calls a person uses, and the database does not record who called them. So a +#: report that trusted `has_user_name` alone opened with forty lines of ELF +#: header trivia. Anything here is the file describing itself; it is reported as +#: a count, never as a finding. +_LOADER_SEGS = frozenset({"LOAD", "HEADER", "MEMORY", "UNDEF", "abs", "extern"}) + +#: Same idea for names the loader derives from format structures. +_LOADER_NAME = re.compile(r"^(?:elf|pe|macho|coff|dos)_", re.I) + + +def from_loader(seg: str, name: str = "") -> bool: + """True if this annotation is the file format describing itself.""" + return (seg or "") in _LOADER_SEGS or bool(_LOADER_NAME.match(name or "")) + + +#: IDA's *analyzer* also writes comments, with `set_cmt`, and the database keeps +#: no record that they are its own -- verified: `get_cmt` at a switch returns +#: "switch jump" with exactly the flags a hand-written comment has. These are +#: its stereotyped shapes, which no one types by accident. +_ANALYZER = re.compile( + r"^(?:switch \d+ cases?|switch jump|jumptable [0-9A-Fa-f]+\b.*|" + r"indirect table for switch.*|jump table for switch.*)$", re.I) + +#: The other family is argument hints (`s1`, `locale`, `domainname`), which IDA +#: copies from the callee's prototype onto each argument-setup instruction. They +#: have no distinguishing shape -- but they REPEAT, once per call site, while a +#: note you wrote is yours alone. Three occurrences of one whitespace-free text +#: is the threshold; anything filtered is counted in the report, never dropped +#: in silence. +_HINT_REPEATS = 3 + + +def analyzer_texts(comments) -> set[str]: + """The comment texts in ``comments`` that look like IDA's own work.""" + counts: dict[str, int] = {} + for c in comments: + text = (c.text or "").strip() + if text and not text.split()[1:]: # a single whitespace-free token + counts[text] = counts.get(text, 0) + 1 + out = {t for t, n in counts.items() if n >= _HINT_REPEATS} + out |= {(c.text or "").strip() for c in comments + if _ANALYZER.match((c.text or "").strip())} + return out + + +#: Names IDA invents when nobody has said otherwise. An address carrying one of +#: these has not been understood by anybody, so it is not a finding. +_DUMMY = re.compile( + r"^(?:(?:sub|loc|locret|off|seg|asc|byte|word|dword|qword|xmmword|ymmword|" + r"flt|dbl|tbyte|stru|algn|unk|nullsub|def|jpt|jsub)_[0-9A-Fa-f]+" + # j_strlen: a thunk name IDA derives from its target, not from a person. + r"|j_\w+)$") + + +def is_dummy(name: str) -> bool: + """True for an IDA-generated placeholder name (``sub_1234``, ``loc_A0``…).""" + return bool(_DUMMY.match(name or "")) + + +def gather(program, path: str = "", *, limit: int = 4000, + types: bool = True, journal=None) -> Findings: + """Collect a :class:`Findings` from a live :class:`Program`. + + ``path`` is the binary the app opened -- ``Program`` speaks to a database + and does not know the file name the user would recognise. + + Every step is individually guarded: a report that is missing its types + section is worth far more than an exception at the end of a long session. + """ + out = Findings() + out.path = path or "" + out.binary = os.path.basename(out.path) if out.path else "" + if journal is not None: + try: + out.recorded = journal.addresses() + out.recorded_types = {e.get("d", "") for e in journal.entries + if e.get("k") == "type" and e.get("d")} + out.n_recorded = len(journal) + except Exception: # noqa: BLE001 + out.recorded, out.recorded_types, out.n_recorded = set(), set(), 0 + try: + out.sections = list(program.sections()) + except Exception: # noqa: BLE001 + out.sections = [] + try: + comments, names = program.annotations(limit=limit) + except Exception: # noqa: BLE001 + comments, names = [], [] + out.comments, out.names = list(comments), list(names) + try: + imports, exports = program.linkage() + out.linked = {i.name for i in imports} | {e.name for e in exports} + except Exception: # noqa: BLE001 + out.linked = set() + try: + idx = program.functions() + idx.load_all() + out.n_functions = len(idx) + # "Stripped" is a judgement about the report, not about the ELF: if + # almost every function is still sub_XXXX, a real name IS a finding. + named = sum(1 for f in idx.all_loaded() if not is_dummy(f.name)) + out.stripped = named <= max(4, out.n_functions // 20) + except Exception: # noqa: BLE001 + pass + if types: + try: + for st in program.list_structs(): + try: + src = program.struct_source(st.name) + except Exception: # noqa: BLE001 + src = "" + out.types.append((st, src)) + except Exception: # noqa: BLE001 + out.types = [] + return out + + +def _esc(text: str) -> str: + """Make one line safe inside a markdown TABLE cell.""" + return (text or "").replace("|", "\\|").replace("\n", " ").strip() + + +def _fence(text: str) -> str: + """Fence body text so a comment containing backticks cannot break out.""" + ticks = "`" * max(3, max((len(m) for m in re.findall(r"`+", text or "")), + default=0) + 1) + return f"{ticks}\n{(text or '').rstrip()}\n{ticks}" + + +def _user_names(f: Findings) -> list: + """The names worth reporting. + + With a journal, that is exactly the addresses we recorded renaming. Without + one, it is a judgement: a real name, not the linker's, not the loader's. + """ + names = [n for n in f.names + if not is_dummy(n.name) and n.name not in f.linked + and not from_loader(n.seg, n.name)] + if f.recorded: + return [n for n in names if n.addr in f.recorded] + return names + + +def _user_types(f: Findings) -> list: + """The types worth reporting. A database is seeded with the type libraries + IDA loaded, so with a journal we show only the ones declared here; without + one, all of them, newest ordinal first (yours are the newest).""" + if f.recorded or f.recorded_types: + return [t for t in f.types + if getattr(t[0], "name", "") in f.recorded_types] + return list(f.types) + + +def _user_comments(f: Findings) -> list: + """The comments a person wrote: not the loader's, not the analyzer's, and + not the same comment reported twice.""" + auto = analyzer_texts(f.comments) + out, seen = [], set() + for c in f.comments: + text = (c.text or "").strip() + if not text or from_loader(c.seg) or text in auto: + continue + if f.recorded and c.addr not in f.recorded: + continue + # A comment on a function's first instruction comes back BOTH as an + # instruction comment and as the function comment; report it once. + key = (c.addr, text) + if key in seen: + continue + seen.add(key) + out.append(c) + return out + + +def render(f: Findings) -> str: + """Render a :class:`Findings` as a markdown document.""" + when = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(f.generated_at)) + names = sorted(_user_names(f), key=lambda n: n.addr) + funcs = [n for n in names if n.is_func] + data = [n for n in names if not n.is_func] + comments = sorted(_user_comments(f), key=lambda c: (c.func_addr or c.addr, + c.addr)) + dropped = (len(f.comments) - len(comments)) + (len(f.names) - len(names)) + types = _user_types(f) + + L: list[str] = [] + title = f.binary or "database" + L.append(f"# Findings — {title}") + L.append("") + L.append(f"*{len(funcs)} named functions · {len(data)} named data · " + f"{len(comments)} comments · {len(types)} local types — " + f"exported {when} by idatui*") + L.append("") + if f.path: + L.append(f"- **binary**: `{f.path}`") + if f.n_functions: + L.append(f"- **functions**: {f.n_functions}") + if f.sections: + segs = ", ".join(f"`{nm}` {s:#x}–{e:#x}" for s, e, nm in f.sections[:8]) + more = f" (+{len(f.sections) - 8} more)" if len(f.sections) > 8 else "" + L.append(f"- **segments**: {segs}{more}") + if f.recorded or f.recorded_types: + n_at = len(f.recorded) + L.append(f"- **source**: idatui's edit journal — {f.n_recorded} recorded " + f"edits across {n_at} address{'' if n_at == 1 else 'es'}. " + "Everything below is work done here, not the analyzer's.") + else: + L.append("- **source**: a scan of the database. Nothing in a `.i64` " + "records *who* wrote a comment or a name — IDA's own analyzer " + "uses the same calls — so this is filtered by shape and may " + "include its work as well as yours.") + if not f.stripped: + L.append("- **note**: this binary has its own symbols, so the names " + "below include ones it shipped with.") + if dropped and (f.recorded or f.recorded_types): + L.append(f"- **note**: {dropped} other annotations in this database " + "were not made here (the analyzer's, the loader's, the " + "linker's) and are left out.") + elif dropped: + L.append(f"- **note**: {dropped} annotations left out as the loader's " + "own (file headers, dummy names, imports).") + if f.truncated: + L.append("- **note**: the scan hit its limit; this report is partial.") + L.append("") + + # -- comments: the actual reasoning, so they lead ----------------------- # + L.append("## Comments") + L.append("") + if not comments: + L.append("*None. (Comments are the part of a database nobody else can " + "reconstruct — they are worth writing.)*") + L.append("") + else: + by_func: dict[str, list] = {} + for c in comments: + by_func.setdefault(c.func or "", []).append(c) + for fn in sorted(by_func, key=lambda k: (k == "", k)): + rows = by_func[fn] + head = f"### `{fn}`" if fn else "### outside any function" + if fn and rows[0].func_addr is not None: + head += f" ({rows[0].func_addr:#x})" + L.append(head) + L.append("") + for c in rows: + if c.whole_func: + L.append(f"- **{c.addr:#x}** — *whole function*: " + f"{_esc(c.text)}") + elif c.line: + L.append(f"- **{c.addr:#x}** `{_esc(c.line)}` \n" + f" {_esc(c.text)}") + else: + L.append(f"- **{c.addr:#x}** — {_esc(c.text)}") + L.append("") + + # -- names -------------------------------------------------------------- # + L.append("## Named functions") + L.append("") + if not funcs: + L.append("*None.*") + L.append("") + else: + L.append("| address | name | size | prototype |") + L.append("|---|---|---|---|") + for n in funcs: + proto = f"`{_esc(n.proto)}`" if n.proto else "" + L.append(f"| `{n.addr:#x}` | `{_esc(n.name)}` | " + f"{n.size:#x} | {proto} |") + L.append("") + if data: + L.append("## Named data") + L.append("") + L.append("| address | name | segment |") + L.append("|---|---|---|") + for n in data: + L.append(f"| `{n.addr:#x}` | `{_esc(n.name)}` | {_esc(n.seg)} |") + L.append("") + + # -- types -------------------------------------------------------------- # + if types: + L.append("## Local types") + L.append("") + if not (f.recorded or f.recorded_types): + L.append("*Newest first. A database is seeded with types from the " + "libraries IDA loaded, so the ones you defined are the " + "ones with the highest ordinals — at the top of this " + "list.*") + L.append("") + ordered = sorted(types, key=lambda t: -getattr(t[0], "ordinal", 0)) + for st, src in ordered: + kw = "union" if getattr(st, "is_union", False) else "struct" + L.append(f"### `{kw} {st.name}` " + f"({getattr(st, 'size', 0):#x} bytes, " + f"{getattr(st, 'members', 0)} fields)") + L.append("") + if src: + L.append("```c") + L.append(src.rstrip()) + L.append("```") + L.append("") + return "\n".join(L).rstrip() + "\n" + + +def default_path(program_path: str) -> str: + """Where a report lands if nobody says otherwise: beside the binary.""" + base = program_path or "findings" + return f"{base}.findings.md" + + +def export(program, binary_path: str = "", out_path: str | None = None, *, + limit: int = 4000, types: bool = True, + journal=None) -> tuple[str, Findings]: + """Gather, render and WRITE the report. Returns ``(path, findings)``.""" + f = gather(program, binary_path, limit=limit, types=types, journal=journal) + out = out_path or default_path(f.path) + out = os.path.abspath(os.path.expanduser(out)) + with open(out, "w", encoding="utf-8") as fh: + fh.write(render(f)) + return out, f diff --git a/idatui/formats.py b/idatui/formats.py index b2f784d..f09bb99 100644 --- a/idatui/formats.py +++ b/idatui/formats.py @@ -92,9 +92,18 @@ def needs_load_options(path: str) -> bool: #: reach for first — arm64, aarch64, mips, m68k — are all invalid. Re-run the #: script before adding to this list. PROCESSORS: tuple[tuple[str, str], ...] = ( - # 'arm' covers AArch64 too; arm64/aarch64 are NOT valid -p names, so they - # live in the label where the filter can still find them. - ("arm", "ARM / AArch64 / arm64 — little-endian"), + # Bare 'arm' gives a 64-BIT database in IDA 9 (AArch64). That matters far + # more than it looks: Hex-Rays refuses a 32-bit function in a 64-bit database + # ("only 64-bit functions can be decompiled in the current database"), and + # Thumb doesn't exist in AArch64 at all — so a 32-bit ARM image loaded as + # plain 'arm' disassembles wrongly and can never be decompiled. The database + # bitness is fixed at load; it cannot be corrected afterwards (setting it + # post-hoc makes the decompiler INTERR). Pick the right one here. + ("arm", "ARM64 / AArch64 / arm64 — 64-bit, little-endian"), + ("arm:ARMv7-A", "ARM 32-bit (ARMv7-A) — Thumb capable, most firmware"), + ("arm:ARMv7-M", "ARM 32-bit (ARMv7-M) — Cortex-M, Thumb only"), + ("arm:ARMv6-M", "ARM 32-bit (ARMv6-M) — Cortex-M0/M0+"), + ("arm:ARMv5TE", "ARM 32-bit (ARMv5TE) — older SoCs"), ("armb", "ARM — big-endian"), ("metapc", "x86 / x86-64"), ("mipsl", "MIPS — little-endian"), diff --git a/idatui/graph.py b/idatui/graph.py new file mode 100644 index 0000000..baee597 --- /dev/null +++ b/idatui/graph.py @@ -0,0 +1,750 @@ +"""Layered control-flow-graph layout, in character cells. + +Pure python: no IDA, no Textual, no I/O. That is deliberate — it means the whole +layout can be unit-tested offline in milliseconds (``tests/test_graph.py``) and +iterated without an idalib worker, and it keeps the hard algorithmic part away +from the UI. + +The pipeline is textbook Sugiyama, the same shape IDA's own graph uses: + + 1. break cycles DFS gray-set; back edges are reversed for layout only + 2. layer longest-path ranking on the resulting DAG + 3. dummies an edge spanning k layers becomes a chain of k-1 dummy + nodes, so every segment is between ADJACENT layers and long + edges reserve real horizontal space (this is what makes it + impossible for an edge to need to cross a box) + 4. order median sweeps + adjacent transposition, to cut crossings + 5. x-coords priority/median sweeps, variable node widths + 6. route ports on node borders, one lane-packed channel per layer gap + +Sizing is injected (``sizer``) rather than computed here, so the caller decides +how wide a block is at the current zoom level without this module knowing +anything about text. + +The result is NOT a painted canvas. A big function lays out to millions of +cells, so ``Painting`` is an *index* — per-row horizontal runs, a bucketed +interval index of vertical runs, and point marks — and the view asks it for one +row at a time (``cells_at_row``), exactly like the listing's ``render_line``. +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field + +# Terminal cells are about twice as tall as they are wide, so horizontal gaps +# need roughly 2x the cell count of vertical gaps to look square. +HGAP = 3 # min columns between two boxes in a layer +VGAP = 1 # min rows between a layer band and the channel below it + +# Edge classes, used as style keys by the renderer. +E_UNCOND = "uncond" +E_TRUE = "jump" +E_FALSE = "fall" +E_SWITCH = "switch" +E_BACK = "back" + + +@dataclass +class Block: + """One basic block, as the backend reports it.""" + + id: int + start: int + end: int + succs: list[tuple[int, str]] = field(default_factory=list) + selfloop: bool = False + + +@dataclass +class Node: + """A laid-out box (``block`` set) or a routing dummy (``block`` None).""" + + id: int + block: Block | None = None + label: str = "" + rank: int = 0 + order: int = 0 + x: int = 0 # left column + y: int = 0 # top row + w: int = 1 + h: int = 1 + + @property + def dummy(self) -> bool: + return self.block is None + + @property + def cx(self) -> float: + return self.x + self.w / 2 + + @property + def bottom(self) -> int: + return self.y + self.h - 1 + + @property + def right(self) -> int: + return self.x + self.w - 1 + + def contains(self, row: int, col: int) -> bool: + return self.y <= row <= self.bottom and self.x <= col <= self.right + + def inside(self, row: int, col: int) -> bool: + """Strictly inside the border (where text lives).""" + return (self.y < row < self.bottom) and (self.x < col < self.right) + + +@dataclass +class Edge: + src: int + dst: int + kind: str = E_UNCOND + back: bool = False + chain: list[int] = field(default_factory=list) + + @property + def style(self) -> str: + return E_BACK if self.back else self.kind + + +class _Graph: + def __init__(self) -> None: + self.nodes: dict[int, Node] = {} + self.edges: list[Edge] = [] + self._next = 0 + + def add(self, n: Node) -> Node: + self.nodes[n.id] = n + self._next = max(self._next, n.id + 1) + return n + + def new_dummy(self) -> Node: + n = Node(id=self._next, w=1, h=1) + return self.add(n) + + +# ------------------------------------------------------------ 1. cycles + +def _break_cycles(g: _Graph, root: int) -> None: + """Reverse back edges (DFS gray-set) so layering sees a DAG.""" + color: dict[int, int] = {} + adj: dict[int, list[Edge]] = {i: [] for i in g.nodes} + for e in g.edges: + adj[e.src].append(e) + for start in [root] + [i for i in g.nodes if i != root]: + if color.get(start): + continue + color[start] = 1 + stack = [(start, iter(adj[start]))] + while stack: + node, it = stack[-1] + for e in it: + c = color.get(e.dst, 0) + if c == 1: + e.back = True + elif c == 0: + color[e.dst] = 1 + stack.append((e.dst, iter(adj[e.dst]))) + break + else: + color[node] = 2 + stack.pop() + for e in g.edges: + if e.back: + e.src, e.dst = e.dst, e.src + + +# --------------------------------------------------------- 2. layering + +def _assign_ranks(g: _Graph, root: int) -> None: + """Longest-path layering: rank(v) = 1 + max(rank(preds)). + + Kahn, but it never trusts that ``_break_cycles`` left a perfect DAG: if the + ready queue drains with nodes left over (a residual cycle, or a block only + reachable through a reversed edge) it force-releases the most-constrained + survivor instead of stranding it at rank 0. Getting this wrong collapses the + whole graph into three layers and looks like a layout bug, not a ranking one. + """ + indeg = {i: 0 for i in g.nodes} + adj: dict[int, list[int]] = {i: [] for i in g.nodes} + for e in g.edges: + indeg[e.dst] += 1 + adj[e.src].append(e.dst) + + rank = {i: 0 for i in g.nodes} + done: set[int] = set() + ready = [i for i in g.nodes if indeg[i] == 0] or [root] + pending = dict(indeg) + while len(done) < len(g.nodes): + if not ready: + left = [i for i in g.nodes if i not in done] + ready = [min(left, key=lambda i: (pending[i], rank[i], i))] + i = ready.pop(0) + if i in done: + continue + done.add(i) + for j in adj[i]: + if rank[j] < rank[i] + 1: + rank[j] = rank[i] + 1 + pending[j] -= 1 + if pending[j] <= 0 and j not in done: + ready.append(j) + for i, n in g.nodes.items(): + n.rank = rank[i] + + +# ---------------------------------------------------------- 3. dummies + +def _add_dummies(g: _Graph) -> None: + for e in list(g.edges): + span = g.nodes[e.dst].rank - g.nodes[e.src].rank + if span <= 0: + e.back = True # residual cycle: colour it, route it flat + chain = [e.src] + if span > 1: + for r in range(g.nodes[e.src].rank + 1, g.nodes[e.dst].rank): + d = g.new_dummy() + d.rank = r + chain.append(d.id) + chain.append(e.dst) + e.chain = chain + + +def _layers_of(g: _Graph) -> list[list[int]]: + top = max((n.rank for n in g.nodes.values()), default=0) + layers: list[list[int]] = [[] for _ in range(top + 1)] + for i, n in g.nodes.items(): + layers[n.rank].append(i) + return layers + + +def _segments(g: _Graph) -> list[tuple[int, int, Edge]]: + out = [] + for e in g.edges: + for a, b in zip(e.chain, e.chain[1:]): + out.append((a, b, e)) + return out + + +# ---------------------------------------------------------- 4. ordering + +def _neighbors(g: _Graph) -> tuple[dict[int, list[int]], dict[int, list[int]]]: + down: dict[int, list[int]] = {i: [] for i in g.nodes} + up: dict[int, list[int]] = {i: [] for i in g.nodes} + for a, b, _ in _segments(g): + down[a].append(b) + up[b].append(a) + return down, up + + +def _cross_below(layer: list[int], down: dict[int, list[int]], + pos: dict[int, int]) -> int: + """Crossings between this layer and the one below, counted as inversions + with a Fenwick tree: O(E log E). The naive O(E^2) version is the entire + runtime on a 400-block function (20s vs 150ms), so it is not an option.""" + pairs = [] + for u in layer: + for v in down[u]: + pairs.append((pos[u], pos[v])) + if not pairs: + return 0 + pairs.sort() + size = max(p[1] for p in pairs) + 2 + tree = [0] * (size + 1) + total = seen = 0 + for _, v in pairs: + i = v + 1 + acc, j = 0, i + while j > 0: + acc += tree[j] + j -= j & -j + total += seen - acc + seen += 1 + j = i + while j <= size: + tree[j] += 1 + j += j & -j + return total + + +def _pair_cross(a: int, b: int, side: dict[int, list[int]], + pos: dict[int, int]) -> int: + """Crossings from a's and b's edges to one neighbouring layer given a sits + immediately LEFT of b. Local — O(deg(a)*deg(b)) — so the transposition pass + never has to recount the whole graph per candidate swap.""" + n = 0 + for u in side[a]: + pu = pos[u] + for v in side[b]: + if pu > pos[v]: + n += 1 + return n + + +def _swap_delta(a: int, b: int, down: dict[int, list[int]], + up: dict[int, list[int]], pos: dict[int, int]) -> tuple[int, int]: + """``(keep, swap)`` for the adjacent pair (a, b), both sides, in one pass. + + The same as calling :func:`_pair_cross` four times, which is what the + transposition loop used to do: every neighbour pair was visited twice (once + per direction) and each visit was a python call. Counting both outcomes + while the pair is in hand halves the comparisons and removes three calls per + candidate swap — and this runs a third of a million times over a corpus. + """ + keep = swap = 0 + for side in (down, up): + va = side[a] + vb = side[b] + if not va or not vb: + continue + pbs = [pos[v] for v in vb] + for u in va: + pu = pos[u] + for pv in pbs: + if pu > pv: + keep += 1 + elif pu < pv: + swap += 1 + return keep, swap + + +def crossings(layers: list[list[int]], down: dict[int, list[int]], + pos: dict[int, int]) -> int: + return sum(_cross_below(l, down, pos) for l in layers) + + +def _order_layers(g: _Graph, root: int, sweeps: int = 6) -> list[list[int]]: + layers = _layers_of(g) + down, up = _neighbors(g) + + # Seed with a DFS preorder so the picture already resembles control flow + # (fallthrough-first); barycenter alone does not recover that. + seed: dict[int, int] = {} + stack, seen, tick = [root], {root}, 0 + while stack: + i = stack.pop() + seed[i] = tick + tick += 1 + for j in reversed(down.get(i, [])): + if j not in seen: + seen.add(j) + stack.append(j) + for layer in layers: + layer.sort(key=lambda i: seed.get(i, 10 ** 9)) + pos = {i: k for layer in layers for k, i in enumerate(layer)} + + def median(i: int, side: dict[int, list[int]]) -> float: + # Almost every node in a control-flow graph has one or two neighbours + # on a given side, so answer those without building and sorting a list: + # this runs tens of thousands of times per corpus layout. + js = side[i] + n = len(js) + if n == 1: + return float(pos[js[0]]) + if n == 2: + return (pos[js[0]] + pos[js[1]]) / 2 + if not n: + return -1.0 + ps = sorted(pos[j] for j in js) + m = n // 2 + return float(ps[m]) if n % 2 else (ps[m - 1] + ps[m]) / 2 + + best, best_x = [list(l) for l in layers], crossings(layers, down, pos) + for s in range(sweeps): + rng = range(1, len(layers)) if s % 2 == 0 else range(len(layers) - 2, -1, -1) + side = up if s % 2 == 0 else down + for r in rng: + layer = layers[r] + keys = {i: median(i, side) for i in layer} + layer.sort(key=lambda i: (keys[i] if keys[i] >= 0 else pos[i], pos[i])) + for k, i in enumerate(layer): + pos[i] = k + for _ in range(4): + improved = False + for layer in layers: + for k in range(len(layer) - 1): + a, b = layer[k], layer[k + 1] + keep, swap = _swap_delta(a, b, down, up, pos) + if swap < keep: + layer[k], layer[k + 1] = b, a + pos[a], pos[b] = k + 1, k + improved = True + if not improved: + break + x = crossings(layers, down, pos) + if x < best_x: + best, best_x = [list(l) for l in layers], x + + layers = best + for layer in layers: + for k, i in enumerate(layer): + g.nodes[i].order = k + return layers + + +# --------------------------------------------------------- 5. x coords + +def _assign_x(g: _Graph, layers: list[list[int]], sweeps: int = 8) -> None: + down, up = _neighbors(g) + for layer in layers: + x = 0 + for i in layer: + g.nodes[i].x = x + x += g.nodes[i].w + HGAP + + def pack(layer: list[int]) -> None: + for k in range(1, len(layer)): + a, b = g.nodes[layer[k - 1]], g.nodes[layer[k]] + if b.x < a.x + a.w + HGAP: + b.x = a.x + a.w + HGAP + for k in range(len(layer) - 2, -1, -1): + a, b = g.nodes[layer[k]], g.nodes[layer[k + 1]] + if a.x + a.w + HGAP > b.x: + a.x = b.x - HGAP - a.w + + for s in range(sweeps): + rng = range(1, len(layers)) if s % 2 == 0 else range(len(layers) - 2, -1, -1) + side = up if s % 2 == 0 else down + for r in rng: + layer = layers[r] + # dummies first: keeping long edges straight matters most + order = sorted(layer, key=lambda i: (not g.nodes[i].dummy, + g.nodes[i].order)) + for i in order: + nb = side[i] + if not nb: + continue + cs = sorted(g.nodes[j].cx for j in nb) + m = len(cs) // 2 + target = cs[m] if len(cs) % 2 else (cs[m - 1] + cs[m]) / 2 + g.nodes[i].x = int(round(target - g.nodes[i].w / 2)) + pack(layer) + + lo = min((g.nodes[i].x for layer in layers for i in layer), default=0) + for n in g.nodes.values(): + n.x -= lo + + +# ------------------------------------------------------------ 6. route + +def _ports(g: _Graph) -> tuple[dict, dict]: + """Spread a node's out-edges along its bottom border and its in-edges along + its top, each ordered by the other end's x so they don't cross at the node.""" + out_port: dict[tuple, int] = {} + in_port: dict[tuple, int] = {} + by_src: dict[int, list] = {} + by_dst: dict[int, list] = {} + for a, b, e in _segments(g): + by_src.setdefault(a, []).append((a, b, e)) + by_dst.setdefault(b, []).append((a, b, e)) + + def spread(n: Node, k: int, count: int) -> int: + if n.dummy or count <= 1: + return int(n.cx) + usable = max(n.w - 4, 1) + step = usable / (count + 1) + return int(n.x + 2 + step * (k + 1)) + + for i, lst in by_src.items(): + lst.sort(key=lambda t: g.nodes[t[1]].cx) + for k, (a, b, e) in enumerate(lst): + out_port[(a, b, id(e))] = spread(g.nodes[i], k, len(lst)) + for i, lst in by_dst.items(): + lst.sort(key=lambda t: g.nodes[t[0]].cx) + for k, (a, b, e) in enumerate(lst): + in_port[(a, b, id(e))] = spread(g.nodes[i], k, len(lst)) + return out_port, in_port + + +@dataclass +class Route: + edge: Edge + pts: list[tuple[int, int]] + head: bool = True # arrowhead (target is a real block) + tail: bool = True # port tee (source is a real block) + + +def _route(g: _Graph, layers: list[list[int]]) -> list[Route]: + out_port, in_port = _ports(g) + segs = _segments(g) + by_rank: dict[int, list] = {} + for a, b, e in segs: + by_rank.setdefault(g.nodes[a].rank, []).append((a, b, e)) + + lanes: dict[tuple, int] = {} + channels = [1] * len(layers) + for r, lst in by_rank.items(): + runs = [] + for a, b, e in lst: + x0, x1 = out_port[(a, b, id(e))], in_port[(a, b, id(e))] + if x0 != x1: # a straight drop needs no lane + runs.append((min(x0, x1), max(x0, x1), (a, b, id(e)))) + runs.sort(key=lambda t: (t[1] - t[0], t[0])) + occupied: list[list[tuple[int, int]]] = [] + for lo, hi, key in runs: + for li, used in enumerate(occupied): + if all(hi < u_lo or lo > u_hi for u_lo, u_hi in used): + used.append((lo, hi)) + lanes[key] = li + break + else: + occupied.append([(lo, hi)]) + lanes[key] = len(occupied) - 1 + channels[r] = max(len(occupied), 1) + + # y: a band per layer, then the routing channel underneath it. Horizontal + # runs live in the channel BELOW A WHOLE LAYER, never at a per-node offset + # -- that is what stops an edge sawing through a taller neighbour. + chan_y = [] + y = 0 + for r, layer in enumerate(layers): + h = max((g.nodes[i].h for i in layer if not g.nodes[i].dummy), default=1) + for i in layer: + n = g.nodes[i] + n.y = y + if n.dummy: + n.h = h # the band is its pass-through + chan_y.append(y + h - 1 + VGAP) + y += h - 1 + VGAP + channels[r] + VGAP + 1 + + def exit_y(n: Node) -> int: + # a dummy leaves from the TOP of its band: its own outgoing segment + # draws the vertical that passes through the band. + return n.y if n.dummy else n.bottom + + routes = [] + for a, b, e in segs: + na, nb = g.nodes[a], g.nodes[b] + x0, x1 = out_port[(a, b, id(e))], in_port[(a, b, id(e))] + y0, y1 = exit_y(na), nb.y + if x0 == x1: + pts = [(y0, x0), (y1, x1)] + else: + ych = chan_y[na.rank] + lanes.get((a, b, id(e)), 0) + pts = [(y0, x0), (ych, x0), (ych, x1), (y1, x1)] + routes.append(Route(edge=e, pts=pts, + head=not nb.dummy, tail=not na.dummy)) + return routes + + +# ------------------------------------------------------------ painting + +BOX = {"tl": "\u250c", "tr": "\u2510", "bl": "\u2514", "br": "\u2518", + "h": "\u2500", "v": "\u2502"} +LINE_CHARS = set("\u2502\u2500\u250c\u2510\u2514\u2518\u251c\u2524\u252c\u2534" + "\u253c\u256d\u256e\u2570\u256f") +MERGE = { + frozenset("\u2502\u2500"): "\u253c", + frozenset("\u2502\u250c"): "\u251c", frozenset("\u2502\u2510"): "\u2524", + frozenset("\u2502\u2514"): "\u251c", frozenset("\u2502\u2518"): "\u2524", + frozenset("\u2500\u250c"): "\u252c", frozenset("\u2500\u2510"): "\u252c", + frozenset("\u2500\u2514"): "\u2534", frozenset("\u2500\u2518"): "\u2534", + frozenset("\u2502\u256d"): "\u251c", frozenset("\u2502\u256e"): "\u2524", + frozenset("\u2502\u2570"): "\u251c", frozenset("\u2502\u256f"): "\u2524", + frozenset("\u2500\u256d"): "\u252c", frozenset("\u2500\u256e"): "\u252c", + frozenset("\u2500\u2570"): "\u2534", frozenset("\u2500\u256f"): "\u2534", +} +CORNER = { + ("D", "R"): "\u2570", ("D", "L"): "\u256f", ("R", "D"): "\u256e", + ("L", "D"): "\u256d", ("R", "U"): "\u256f", ("L", "U"): "\u2570", + ("U", "R"): "\u256d", ("U", "L"): "\u256e", +} +BUCKET = 32 # rows per vertical-run index bucket + + +def _dir(p: tuple[int, int], q: tuple[int, int]) -> str: + if p[0] == q[0]: + return "R" if q[1] > p[1] else "L" + return "D" if q[0] > p[0] else "U" + + +class Painting: + """A queryable drawing of the edges. Never a full canvas: a 400-block + function is ~13M cells, so runs are stored as intervals and asked for one + row at a time.""" + + def __init__(self) -> None: + self.hruns: dict[int, list[tuple[int, int, str, int]]] = {} + self.vruns: list[tuple[int, int, int, str, int]] = [] + self.vindex: dict[int, list[int]] = {} + self.marks: dict[int, list[tuple[int, str, str, int]]] = {} + + def add_h(self, row: int, c0: int, c1: int, style: str, eid: int) -> None: + self.hruns.setdefault(row, []).append((min(c0, c1), max(c0, c1), style, eid)) + + def add_v(self, r0: int, r1: int, col: int, style: str, eid: int) -> None: + lo, hi = (r0, r1) if r0 <= r1 else (r1, r0) + idx = len(self.vruns) + self.vruns.append((lo, hi, col, style, eid)) + for b in range(lo // BUCKET, hi // BUCKET + 1): + self.vindex.setdefault(b, []).append(idx) + + def add_mark(self, row: int, col: int, ch: str, style: str, eid: int) -> None: + self.marks.setdefault(row, []).append((col, ch, style, eid)) + + def cells_at_row(self, row: int, c0: int, c1: int + ) -> dict[int, tuple[str, str, int]]: + """{col: (char, style, edge_id)} for ``row`` within [c0, c1).""" + out: dict[int, tuple[str, str, int]] = {} + + def put(col: int, ch: str, style: str, eid: int, force: bool = False) -> None: + if col < c0 or col >= c1: + return + old = out.get(col) + if old and not force and old[0] != ch \ + and old[0] in LINE_CHARS and ch in LINE_CHARS: + ch = MERGE.get(frozenset((old[0], ch)), ch) + out[col] = (ch, style, eid) + + for lo, hi, style, eid in self.hruns.get(row, ()): + for c in range(max(lo, c0), min(hi + 1, c1)): + put(c, BOX["h"], style, eid) + for i in self.vindex.get(row // BUCKET, ()): + lo, hi, col, style, eid = self.vruns[i] + if lo <= row <= hi: + put(col, BOX["v"], style, eid) + for col, ch, style, eid in self.marks.get(row, ()): + put(col, ch, style, eid, force=True) + return out + + +@dataclass +class Layout: + """The finished drawing: boxes, an edge index, and enough structure for the + view to hit-test, navigate and highlight.""" + + nodes: list[Node] # real blocks only, layout order + by_id: dict[int, Node] + edges: list[Edge] + painting: Painting + width: int + height: int + entry: int + rows: dict[int, list[int]] # row -> real node ids covering it + incident: dict[int, set[int]] # node id -> edge ids touching it + succ: dict[int, list[tuple[int, str]]] # node id -> [(node id, style)] + pred: dict[int, list[tuple[int, str]]] + stats: dict + + def node_at(self, row: int, col: int) -> Node | None: + for nid in self.rows.get(row, ()): + n = self.by_id[nid] + if n.x <= col <= n.right: + return n + return None + + def nodes_at_row(self, row: int) -> list[Node]: + return [self.by_id[i] for i in self.rows.get(row, ())] + + def edge_at(self, row: int, col: int) -> Edge | None: + cells = self.painting.cells_at_row(row, col, col + 1) + hit = cells.get(col) + if hit is None: + return None + for e in self.edges: + if id(e) == hit[2]: + return e + return None + + +def layout(blocks: list[Block], sizer, entry: int | None = None) -> Layout: + """Lay out ``blocks``. ``sizer(block) -> (width, height)`` in cells.""" + t0 = time.perf_counter() + g = _Graph() + for b in blocks: + w, h = sizer(b) + g.add(Node(id=b.id, block=b, label=f"loc_{b.start:X}", + w=max(int(w), 4), h=max(int(h), 3))) + for b in blocks: + outs = [(d, k) for d, k in b.succs if d in g.nodes] + for dst, kind in outs: + if dst == b.id: + # A self-loop constrains nothing and would deadlock the Kahn + # ranking (its own in-degree never drains). Drawn as a marker. + b.selfloop = True + continue + if len(outs) == 1: + kind = E_UNCOND + g.edges.append(Edge(src=b.id, dst=dst, kind=kind)) + + root = entry if entry in g.nodes else (min(g.nodes) if g.nodes else 0) + if g.nodes: + _break_cycles(g, root) + _assign_ranks(g, root) + _add_dummies(g) + layers = _order_layers(g, root) + _assign_x(g, layers) + routes = _route(g, layers) + else: + layers, routes = [], [] + + # ---- paint into the index ----------------------------------------- + p = Painting() + real = [n for n in g.nodes.values() if not n.dummy] + rows: dict[int, list[int]] = {} + for n in real: + for r in range(n.y, n.y + n.h): + rows.setdefault(r, []).append(n.id) + for lst in rows.values(): + lst.sort(key=lambda i: g.nodes[i].x) + + def blocked(row: int, col: int) -> bool: + for nid in rows.get(row, ()): + if g.nodes[nid].inside(row, col): + return True + return False + + incident: dict[int, set[int]] = {n.id: set() for n in real} + for rt in routes: + e, style, eid = rt.edge, rt.edge.style, id(rt.edge) + incident.setdefault(e.src, set()).add(eid) + incident.setdefault(e.dst, set()).add(eid) + for (r0, c0), (r1, c1) in zip(rt.pts, rt.pts[1:]): + if r0 == r1: + p.add_h(r0, c0, c1, style, eid) + else: + p.add_v(r0, r1, c0, style, eid) + for k in range(1, len(rt.pts) - 1): + a, b, c = rt.pts[k - 1], rt.pts[k], rt.pts[k + 1] + ch = CORNER.get((_dir(a, b), _dir(b, c))) + if ch and not blocked(*b): + p.add_mark(b[0], b[1], ch, style, eid) + # A back edge was reversed for layering, so its polyline runs from the + # loop HEAD down to the tail: the arrow belongs at the start, pointing + # up into the block control returns to. + first, last = rt.pts[0], rt.pts[-1] + if e.back: + if rt.tail: + p.add_mark(first[0], first[1], "\u25b2", style, eid) + if rt.head: + p.add_mark(last[0], last[1], "\u2534", style, eid) + else: + if rt.tail: + p.add_mark(first[0], first[1], "\u252c", style, eid) + if rt.head: + p.add_mark(last[0], last[1], "\u25bc", style, eid) + + succ: dict[int, list[tuple[int, str]]] = {n.id: [] for n in real} + pred: dict[int, list[tuple[int, str]]] = {n.id: [] for n in real} + for e in g.edges: + a, b = (e.dst, e.src) if e.back else (e.src, e.dst) # undo reversal + if a in succ: + succ[a].append((b, e.style)) + if b in pred: + pred[b].append((a, e.style)) + + width = max((n.right + 1 for n in real), default=1) + height = max((n.y + n.h for n in real), default=1) + order = sorted(real, key=lambda n: (n.rank, n.order)) + stats = { + "blocks": len(blocks), + "nodes": len(g.nodes), + "dummies": len(g.nodes) - len(real), + "layers": len(layers), + "edges": len(g.edges), + "back": sum(1 for e in g.edges if e.back), + "ms": (time.perf_counter() - t0) * 1000, + } + return Layout(nodes=order, by_id={n.id: n for n in g.nodes.values()}, + edges=g.edges, painting=p, width=width, height=height, + entry=root, rows=rows, incident=incident, + succ=succ, pred=pred, stats=stats) diff --git a/idatui/highlight.py b/idatui/highlight.py index 5f504bf..0bd76af 100644 --- a/idatui/highlight.py +++ b/idatui/highlight.py @@ -1,10 +1,16 @@ -"""Pygments-based C highlighting for Hex-Rays pseudocode. +"""Pygments-based C highlighting for Hex-Rays pseudocode and the struct editor. Textual's TextArea has no C/C++ tree-sitter grammar (its bundled languages are python/rust/go/... only), so ``language="cpp"`` silently does nothing. Pygments (already a Rich/Textual dependency) has a solid C lexer, so we tokenize once and map tokens to Rich styles, producing per-line Segment lists the virtualized view can cache and paint instantly. + +The same tokenizer feeds two consumers, so one palette covers both: + +* ``highlight_c`` -> Segment lists for the read-only pseudocode view. +* ``CTextArea`` -> an *editable* TextArea (the struct editor) highlighted by + filling TextArea's own ``_highlights`` map, the hook tree-sitter would use. """ from __future__ import annotations @@ -13,8 +19,12 @@ from rich.segment import Segment from rich.style import Style from pygments.lexers import CLexer from pygments.token import Token +from textual.widgets import TextArea +from textual.widgets.text_area import TextAreaTheme -# Token -> style, checked in priority order (first hierarchical match wins). +# Token -> (highlight name, style), checked in priority order (first +# hierarchical match wins). The names are what TextArea's theme maps to styles; +# the styles are what the pseudocode view paints directly. # # Same measured palette as the listing (see the theme notes): one hue = one # meaning across BOTH panes, so a string is the same green and a symbol the same @@ -25,27 +35,59 @@ from pygments.token import Token # Control keywords take the brightest NEUTRAL rather than a hue, mirroring the # mnemonic column: they're the skeleton you scan for, and a hue there would # claim a meaning the rest of the palette already assigns. -_STYLES: list[tuple[object, Style]] = [ - (Token.Comment, Style(color="#7c8b9e", italic=True)), # 5.2:1 commentary - (Token.Keyword.Type, Style(color="#93aee0")), # 8.1:1 type info - (Token.Keyword, Style(color="#e8ecf2", bold=True)), # 15.3:1 control flow - (Token.Name.Builtin, Style(color="#93aee0")), # 8.1:1 type info - (Token.Literal.String, Style(color="#9ece6a")), # 9.9:1 strings - (Token.Literal.Number, Style(color="#d8a657")), # 8.2:1 data/number - (Token.Operator, Style(color="#c3cad3")), # 11.0:1 body - (Token.Punctuation, Style(color="#626c7a")), # 3.4:1 structure - (Token.Name, Style(color="#7aa2f7")), # 7.2:1 symbol names +_PALETTE: list[tuple[str, object, Style]] = [ + ("comment", Token.Comment, Style(color="#7c8b9e", italic=True)), # 5.2:1 commentary + ("type", Token.Keyword.Type, Style(color="#93aee0")), # 8.1:1 type info + ("keyword", Token.Keyword, Style(color="#e8ecf2", bold=True)), # 15.3:1 control flow + ("builtin", Token.Name.Builtin, Style(color="#93aee0")), # 8.1:1 type info + ("string", Token.Literal.String, Style(color="#9ece6a")), # 9.9:1 strings + ("number", Token.Literal.Number, Style(color="#d8a657")), # 8.2:1 data/number + ("operator", Token.Operator, Style(color="#c3cad3")), # 11.0:1 body + ("punctuation", Token.Punctuation, Style(color="#626c7a")), # 3.4:1 structure + ("name", Token.Name, Style(color="#7aa2f7")), # 7.2:1 symbol names ] +_STYLES: list[tuple[object, Style]] = [(t, s) for _, t, s in _PALETTE] _DEFAULT = Style(color="#c3cad3") # 11.0:1 body +_DEFAULT_NAME = "text" + +#: highlight name -> style, for TextArea themes (see ``CTextArea``). +SYNTAX_STYLES: dict[str, Style] = {name: style for name, _, style in _PALETTE} +SYNTAX_STYLES[_DEFAULT_NAME] = _DEFAULT _lexer = CLexer(stripnl=False, ensurenl=False) +#: Resolved styles by token type. Pygments token types are interned singletons +#: and a whole decompilation only ever uses about eighteen of them, but +#: ``token in ttype`` is a hierarchy walk and _STYLES is scanned in order -- so +#: without this every token in the body pays up to nine of those walks. It was a +#: quarter of the time spent highlighting a function. +_STYLE_CACHE: dict[object, Style] = {} +_NAME_CACHE: dict[object, str] = {} + + def _style_for(token) -> Style: - for ttype, style in _STYLES: - if token in ttype: - return style - return _DEFAULT + style = _STYLE_CACHE.get(token) + if style is None: + style = _DEFAULT + for ttype, candidate in _STYLES: + if token in ttype: + style = candidate + break + _STYLE_CACHE[token] = style + return style + + +def _name_for(token) -> str: + name = _NAME_CACHE.get(token) + if name is None: + name = _DEFAULT_NAME + for candidate, ttype, _ in _PALETTE: + if token in ttype: + name = candidate + break + _NAME_CACHE[token] = name + return name def highlight_c(code: str) -> list[list[Segment]]: @@ -55,6 +97,9 @@ def highlight_c(code: str) -> list[list[Segment]]: if not value: continue style = _style_for(token) + if "\n" not in value: # the common case: a token inside one line + lines[-1].append(Segment(value, style)) + continue parts = value.split("\n") for i, part in enumerate(parts): if i > 0: @@ -65,3 +110,82 @@ def highlight_c(code: str) -> list[list[Segment]]: if len(lines) > 1 and not lines[-1]: lines.pop() return lines + + +def highlight_c_spans(code: str) -> dict[int, list[tuple[int, int, str]]]: + """Return ``{row: [(start_byte, end_byte, highlight_name), ...]}`` for ``code``. + + The shape TextArea's ``_highlights`` map wants. Columns are **byte** offsets + into the line, not character offsets -- that's the tree-sitter convention + TextArea's renderer decodes with ``build_byte_to_codepoint_dict``, so a + non-ASCII identifier or string would smear its styling one cell per extra + byte if we handed it character offsets. + """ + spans: dict[int, list[tuple[int, int, str]]] = {} + row = 0 + col = 0 + for token, value in _lexer.get_tokens(code): + if not value: + continue + name = _name_for(token) + parts = value.split("\n") + for i, part in enumerate(parts): + if i: + row += 1 + col = 0 + if not part: + continue + width = len(part) if part.isascii() else len(part.encode("utf-8")) + if part.strip(): # whitespace carries no visible style + spans.setdefault(row, []).append((col, col + width, name)) + col += width + return spans + + +#: TextArea theme carrying our palette. Everything else (background, cursor, +#: selection) is deliberately left unset so it keeps falling back to the app's +#: CSS -- this theme only says how C tokens are coloured. +C_TEXTAREA_THEME = TextAreaTheme(name="idatui-c", syntax_styles=SYNTAX_STYLES) + + +class CTextArea(TextArea): + """An editable TextArea that syntax-highlights C. + + ``language="cpp"`` is not available (no bundled grammar), so instead of a + tree-sitter query we fill the very same ``_highlights`` map the tree-sitter + path fills, from the Pygments lexer above. Everything downstream -- + per-line style application, the line cache, selection, the cursor -- is + stock TextArea, and the colours match the pseudocode pane token for token. + """ + + #: Above this, re-lexing on every keystroke would cost more than the colour + #: is worth. Struct definitions are a few hundred bytes; this is a guard, + #: not a limit anyone should hit. + MAX_HIGHLIGHT_CHARS = 200_000 + + def __init__(self, text: str = "", **kwargs) -> None: + super().__init__(text, **kwargs) + self.register_theme(C_TEXTAREA_THEME) + self.theme = C_TEXTAREA_THEME.name + # __init__ built the document (and so the highlight map) before the + # theme existed; redo it now that tokens can resolve to styles. + self._build_highlight_map() + + def _build_highlight_map(self) -> None: + """Lex the buffer and publish per-line highlight spans. + + Called by TextArea on every document change, so it must be cheap and it + must never raise: a lexer hiccup should cost colour, not the editor. + """ + self._line_cache.clear() + highlights = self._highlights + highlights.clear() + text = self.document.text + if not text or len(text) > self.MAX_HIGHLIGHT_CHARS: + return + try: + spans = highlight_c_spans(text) + except Exception: # noqa: BLE001 - highlighting is never load-bearing + return + for row, row_spans in spans.items(): + highlights[row].extend(row_spans) diff --git a/idatui/journal.py b/idatui/journal.py new file mode 100644 index 0000000..a5a07de --- /dev/null +++ b/idatui/journal.py @@ -0,0 +1,107 @@ +"""A record of the edits idatui makes, kept inside the database. + +**Why this has to exist.** A findings report wants to say "here is what *you* +worked out", and the database cannot answer that. IDA's own analyzer writes +comments with the same `set_cmt` a person uses (`; s1` on an argument setup, +`; switch 73 cases` on a jump table), and the loader writes both comments and +names for the file header. Four separate probes agree that nothing tells them +apart: the `FF_COMM` flag is identical, `get_cmt` returns them all, the colour +tag in `generate_disasm_line` is `COLOR_REGCMT` for every one of them, and they +survive with auto-comments switched off. So authorship is not recoverable after +the fact -- it has to be recorded as it happens, which is what this does. + +It lives in an IDA **netnode**, so it is saved into the `.i64` with everything +else and is still there next session. The entries are small and additive; the +journal is metadata *about* edits that themselves live in the database, so +losing it degrades the report to a heuristic rather than losing work. +""" + +from __future__ import annotations + +import json +import threading +import time + +#: Where the blob lives inside the database. +NODE = "$ idatui.journal" + +#: Cap: a long session is hundreds of edits, not hundreds of thousands, and the +#: blob is rewritten whole. Oldest entries fall off first. +MAX_ENTRIES = 20000 + + +class Journal: + """Append-only log of what was edited, with lazy load and explicit flush. + + Writing through to the database on every keystroke-sized edit would put a + round trip in the way of the user; the in-memory list is authoritative + during a session and :meth:`flush` persists it at the points that already + mean "keep this": saving, exporting, and quitting. + """ + + def __init__(self) -> None: + self.entries: list[dict] = [] + self._dirty = False + self._loaded = False + self._lock = threading.Lock() + + # -- recording ---------------------------------------------------------- # + def record(self, kind: str, ea: int | None = None, detail: str = "", + extra: dict | None = None) -> None: + """Note one edit: ``kind`` is 'rename' / 'comment' / 'retype' / …""" + entry = {"k": str(kind), "t": int(time.time())} + if ea is not None: + entry["ea"] = int(ea) + if detail: + entry["d"] = str(detail)[:400] + if extra: + entry.update(extra) + with self._lock: + self.entries.append(entry) + if len(self.entries) > MAX_ENTRIES: + del self.entries[:len(self.entries) - MAX_ENTRIES] + self._dirty = True + + def addresses(self, kinds: tuple[str, ...] | None = None) -> set[int]: + """Every address touched (optionally only by certain kinds of edit).""" + with self._lock: + return {e["ea"] for e in self.entries + if "ea" in e and (kinds is None or e.get("k") in kinds)} + + def __len__(self) -> int: + return len(self.entries) + + # -- persistence -------------------------------------------------------- # + def load(self, program) -> None: + """Read the journal out of the database, once. Never raises.""" + if self._loaded: + return + self._loaded = True + try: + raw = program.journal_get() + except Exception: # noqa: BLE001 -- an old database simply has none + return + if not raw: + return + try: + data = json.loads(raw) + except Exception: # noqa: BLE001 + return + if isinstance(data, list): + with self._lock: + # Prepend: what is already in memory happened later. + self.entries = [e for e in data if isinstance(e, dict)] + self.entries + + def flush(self, program) -> bool: + """Write the journal back if it changed. Returns whether it wrote.""" + with self._lock: + if not self._dirty: + return False + payload = json.dumps(self.entries, separators=(",", ":")) + try: + program.journal_put(payload) + except Exception: # noqa: BLE001 -- never let bookkeeping break an edit + return False + with self._lock: + self._dirty = False + return True diff --git a/idatui/kittygfx.py b/idatui/kittygfx.py new file mode 100644 index 0000000..e53fba3 --- /dev/null +++ b/idatui/kittygfx.py @@ -0,0 +1,271 @@ +"""Kitty graphics protocol: detect it, upload an image, place it on screen. + +Used for the startup splash, which otherwise falls back to the block-art +``logo.ans``. Two things about this were expensive to find out, so they are +written down here rather than rediscovered. + +**Support cannot be sniffed from the environment.** A multiplexer that passes +the protocol through (recent zellij, tmux with allow-passthrough) leaves TERM as +``xterm-256color`` with ``KITTY_WINDOW_ID``, ``TERM_PROGRAM`` and ``COLORTERM`` +all empty, while the protocol answers perfectly. Detection by terminal name +would disable graphics on exactly the terminals that support them. So we ask: +send a 1x1 graphics query together with a Primary Device Attributes request. +Every terminal answers DA1, so that reply is the sync point -- a ``_G...OK`` +before it means yes, DA1 alone means no. No timeouts to tune, no allowlist. + +**Unicode placeholders are not usable.** The tidy way to put an image in a TUI +is a virtual placement (``U=1``) plus U+10EEEE placeholder cells, which the +compositor then moves and clips like ordinary text. It is also what every +Textual image library is built on -- and this terminal answers +``ENOTSUPPORTED:unicode placeholders are not supported`` while supporting +everything else. So we use ordinary placement: the image is anchored at a screen +cell and stays there until deleted, which means the caller owns its lifetime +(place on mount and resize, delete on unmount) and must reserve blank cells +underneath. That is fine for a splash and deliberately not built up into a +general image widget. + +**Uploading and drawing happen on opposite sides of the alternate screen.** The +detection query must run BEFORE the app starts, because it needs a reply and +Textual reads stdin on its own thread. The IMAGE, though, must be uploaded AFTER +Textual has switched to the alternate screen: an image uploaded to the primary +screen cannot be placed from the alternate one -- placement reports no error, it +simply draws nothing. That combination is why the splash calls ``supported()`` +from the launcher and ``upload()`` from its own ``on_mount``. +""" +from __future__ import annotations + +import base64 +import os +import re +import select +import struct +import sys +import time + +#: One id for the splash. Ids are a terminal-wide namespace shared with whatever +#: else the user is running, so this is deliberately not 1. +LOGO_ID = 0x1DA7 + +_supported: bool | None = None +_uploaded: dict[int, tuple[int, int]] = {} # image id -> (pixel w, pixel h) +#: Terminal cell size in pixels, asked for in the same round trip as the +#: graphics query. Cells are nothing like a fixed 1:2 -- this box reports 9x22, +#: i.e. 1:2.44 -- and getting it wrong stretches the image. +_cell: tuple[int, int] | None = None + + +def log(msg: str) -> None: + """Trace to ``$IDATUI_KITTY_LOG``. The splash lives inside a full-screen TUI + on a tty we can't print to, so this is the only way to see what it decided.""" + path = os.environ.get("IDATUI_KITTY_LOG") + if not path: + return + try: + with open(path, "a") as f: + f.write(f"{time.time():.3f} {msg}\n") + except OSError: + pass + + +# --------------------------------------------------------------------------- # +# Detection +# --------------------------------------------------------------------------- # +def _query_tty(timeout: float = 2.0) -> bool: + import termios + import tty as ttymod + + try: + fd = os.open("/dev/tty", os.O_RDWR | os.O_NOCTTY) + except OSError: + return False + try: + old = termios.tcgetattr(fd) + except termios.error: + os.close(fd) + return False + try: + ttymod.setraw(fd) + # graphics query + cell-size query + DA1. DA1 is answered by everything, + # so it marks the end of the replies and nothing has to be timed. + os.write(fd, b"\033_Gi=31,s=1,v=1,a=q,t=d,f=24;AAAA\033\\\033[16t\033[c") + buf = b"" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + r, _, _ = select.select([fd], [], [], 0.15) + if not r: + continue + chunk = os.read(fd, 4096) + if not chunk: + break + buf += chunk + if re.search(rb"\033\[\?[0-9;]*c", buf): # DA1: the answers are in + break + global _cell + m = re.search(rb"\033\[6;(\d+);(\d+)t", buf) # CSI 6 ; height ; width t + if m: + ch, cw = int(m.group(1)), int(m.group(2)) + if 0 < cw < 100 and 0 < ch < 200: + _cell = (cw, ch) + log(f"cell size {cw}x{ch}px") + return bool(re.search(rb"\033_G[^\033]*;OK\033\\", buf)) + except OSError: + return False + finally: + try: + termios.tcsetattr(fd, termios.TCSANOW, old) + finally: + os.close(fd) + + +def supported() -> bool: + """True if the terminal speaks the kitty graphics protocol. + + ``$IDATUI_KITTY=0/1`` forces the answer, for a terminal that swallows the + query and for tests. Cached: the query costs a round trip and must not run + once Textual owns stdin. + """ + global _supported + if _supported is not None: + return _supported + env = os.environ.get("IDATUI_KITTY", "").strip().lower() + if env in ("1", "yes", "true", "on"): + _supported = True + elif env in ("0", "no", "false", "off"): + _supported = False + elif not (sys.__stdout__ and sys.__stdout__.isatty()): + _supported = False # pilot tests, pipes, redirected output + log("supported: stdout is not a tty") + else: + _supported = _query_tty() + log(f"supported() -> {_supported}") + return _supported + + +# --------------------------------------------------------------------------- # +# Upload / place / delete +# --------------------------------------------------------------------------- # +def png_size(path: str) -> tuple[int, int] | None: + """(width, height) from a PNG's IHDR, without decoding it.""" + try: + with open(path, "rb") as f: + head = f.read(26) + except OSError: + return None + if len(head) < 24 or head[:8] != b"\x89PNG\r\n\x1a\n": + return None + w, h = struct.unpack(">II", head[16:24]) + return (w, h) + + +def _write(data: str) -> bool: + """Write escapes to the same stream Textual writes frames to, so the two + can't be reordered. Called only from the app's own loop.""" + out = sys.__stdout__ + if out is None: + return False + try: + out.write(data) + out.flush() + return True + except (OSError, ValueError): + return False + + +def upload(path: str, image_id: int = LOGO_ID) -> bool: + """Send the PNG to the terminal WITHOUT placing it (``a=t``). + + Must be called once the app is already on the ALTERNATE screen -- an image + uploaded to the primary screen can't be placed from the alternate one, and + the placement fails silently. Idempotent, so callers can just ask. + """ + if image_id in _uploaded: + return True + size = png_size(path) + if size is None: + return False + try: + with open(path, "rb") as f: + payload = base64.standard_b64encode(f.read()) + except OSError: + return False + parts = [payload[i:i + 4096] for i in range(0, len(payload), 4096)] + if not parts: + return False + buf = [] + for i, part in enumerate(parts): + more = 1 if i < len(parts) - 1 else 0 + ctrl = (f"a=t,f=100,t=d,i={image_id},q=2,m={more}" if i == 0 + else f"m={more}") + buf.append("\033_G" + ctrl + ";" + part.decode("ascii") + "\033\\") + if not _write("".join(buf)): + log("upload: write failed") + return False + _uploaded[image_id] = size + log(f"upload -> ok id={image_id} px={size} chunks={len(parts)}") + return True + + +def is_uploaded(image_id: int = LOGO_ID) -> bool: + return image_id in _uploaded + + +def place(row: int, col: int, cols: int, rows: int, + image_id: int = LOGO_ID) -> bool: + """Draw the uploaded image at (``row``, ``col``), 0-based, sized in cells. + + Saves and restores the cursor, and asks the terminal not to move it + (``C=1``), so Textual's idea of where the cursor is stays true. + """ + size = _uploaded.get(image_id) + if size is None or cols <= 0 or rows <= 0: + log(f"place: refused size={size} cols={cols} rows={rows}") + return False + w, h = size + log(f"place row={row} col={col} c={cols} r={rows}") + return _write( + f"\033[s\033[{row + 1};{col + 1}H" + f"\033_Ga=p,i={image_id},s={w},v={h},c={cols},r={rows},C=1,q=2\033\\" + f"\033[u") + + +def clear(image_id: int = LOGO_ID) -> None: + """Remove the image's placements from the screen (it stays uploaded).""" + _write(f"\033_Ga=d,d=i,i={image_id},q=2\033\\") + + +def delete(image_id: int = LOGO_ID) -> None: + """Remove the placements AND free the image data in the terminal.""" + _write(f"\033_Ga=d,d=I,i={image_id},q=2\033\\") + _uploaded.pop(image_id, None) + + +def cell_size() -> tuple[int, int]: + """(width, height) of a terminal cell in pixels. + + Measured during the graphics query when the terminal answers CSI 16 t; + otherwise a 10x20 guess, which is only ever used to keep the aspect ratio + honest. + """ + return _cell or (10, 20) + + +def fit(px: tuple[int, int], max_cols: int, max_rows: int, + cell: tuple[int, int] | None = None) -> tuple[int, int]: + """Cell size that fits ``max_cols`` x ``max_rows`` keeping the aspect ratio. + + Cells are far from square -- this box reports 9x22 px -- so a naive + cols==rows box stretches the image; ``cell`` is that ratio in pixels and + defaults to what the terminal actually said. + """ + if cell is None: + cell = cell_size() + w, h = px + if w <= 0 or h <= 0: + return (max_cols, max_rows) + cw, ch = cell + cols = max_cols + rows = max(int(round((h / w) * cols * cw / ch)), 1) + if rows > max_rows: + rows = max_rows + cols = max(int(round((w / h) * rows * ch / cw)), 1) + return (max(cols, 1), max(rows, 1)) diff --git a/idatui/launch.py b/idatui/launch.py index f51e5af..0c53987 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -1,15 +1,13 @@ -"""One-shot launcher: ``ida-tui foo.elf`` and you're in the TUI. +"""One-shot launcher for the IDA Code Mode-backed TUI. -Spawns a private idalib worker (``idatui.worker``) that opens + auto-analyzes -THIS binary in its own process, talking to the TUI over a unix socket. No shared -supervisor, no HTTP: everything slow (open + analysis) happens behind the TUI's -loading overlay. +A path first resolves to a registered GUI database; when none matches, Code Mode +reuses or starts a managed idalib worker. With no path, a single registered +database is selected automatically. -Usage: +Usage:: - ida-tui /path/to/binary # open a binary and drive it - -Extras: --ttl, --no-keepalive, --rpc (all forwarded to the TUI). + ida-tui /path/to/binary + ida-tui # attach when exactly one database is registered """ from __future__ import annotations @@ -17,16 +15,17 @@ import argparse import os import sys -# The unpacked working-copy files IDA writes next to a `.i64` while a database is -# open. A hard-killed worker leaves them behind and the `.i64` then refuses to -# reopen ("Failed to open database"). Safe to delete when nothing holds the DB. -_LOCK_SUFFIXES = (".id0", ".id1", ".id2", ".nam", ".til") - - def _load_args(load: dict) -> str: - """``load`` as IDA switches, for the single-binary path (no project ref).""" + """``load`` as IDA switches, for the single-binary path (no project ref). + + The base goes through project._as_addr rather than bare int(): our own CLI + hands over an int, but a project file writes "0x8000000" as a string and + int() raises on that. One parser, so the two paths can't disagree about what + an address looks like. + """ from .formats import load_args - return load_args(load.get("processor", ""), int(load.get("base", 0) or 0), + from .project import _as_addr + return load_args(load.get("processor", ""), _as_addr(load.get("base", 0)), str(load.get("ida_args", "") or "")) @@ -34,24 +33,19 @@ def _log(msg: str) -> None: print(f"ida-tui: {msg}", file=sys.stderr) -def _sweep_locks(binary: str) -> int: - """Remove stale unpacked DB files next to ``binary``. Returns how many.""" - stem = os.path.splitext(binary)[0] - n = 0 - for base in (binary, stem): # IDA may key on the full name or the stem - for suf in _LOCK_SUFFIXES: - try: - os.remove(base + suf) - n += 1 - except OSError: - pass - return n +def _registered_databases() -> tuple[list[dict], list[dict]]: + """Ready and blocked Code Mode registrations, with normalized errors.""" + try: + from ida_codemode.registry import discover_instances + return discover_instances() + except Exception as exc: # discovery diagnostics belong at the CLI boundary + return [], [{"error": str(exc)}] def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", - description="Open a binary in the IDA TUI (private idalib worker).") + description="Open a registered GUI or managed idalib database in the IDA TUI.") p.add_argument("binary", nargs="*", help="binary to open and analyze (several with --project " "creates/extends that project)") @@ -59,11 +53,13 @@ def main(argv: list[str] | None = None) -> int: help="open a multi-binary project (created from the given " "binaries if FILE doesn't exist)") p.add_argument("--ttl", type=int, default=1800, - help="worker idle-TTL seconds (default 1800)") + help="deprecated compatibility option (Code Mode uses leases)") p.add_argument("--no-keepalive", action="store_true", - help="do not run the keepalive heartbeat") + help="deprecated compatibility option (the lease is the heartbeat)") p.add_argument("--rpc", metavar="PATH", help="listen for RPC on this unix socket (puppeteer the TUI)") + p.add_argument("--trace", metavar="FILE", + help="Tenet execution trace to explore alongside the binary") g = p.add_argument_group( "loading a headerless blob", "An ELF/PE/Mach-O says what it is. A raw firmware dump doesn't, and IDA " @@ -74,7 +70,7 @@ def main(argv: list[str] | None = None) -> int: g.add_argument("--base", metavar="ADDR", help="load address, e.g. 0x8000000 (any base; NOT paragraphs)") g.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="extra IDA command-line switches, passed through as-is") + help="legacy switches; only Code Mode-representable -p/-b/-T are accepted") args = p.parse_args(argv) load: dict = {} @@ -132,32 +128,61 @@ def main(argv: list[str] | None = None) -> int: _log(str(e)) return 2 else: - if len(args.binary) != 1: - _log("give exactly one binary, or use --project for several") + ready, blocked = _registered_databases() + if len(args.binary) > 1: + _log("give at most one binary, or use --project for several") return 2 - binary = os.path.abspath(os.path.expanduser(args.binary[0])) - if not os.path.isfile(binary): - _log(f"no such file: {binary}") + if args.binary: + binary = os.path.abspath(os.path.expanduser(args.binary[0])) + key = os.path.normcase(os.path.realpath(binary)) + registered = any( + key == os.path.normcase(os.path.realpath(str(item.get(field) or ""))) + for item in ready for field in ("exe_path", "idb_path") + if item.get(field) + ) + if not os.path.isfile(binary) and not registered: + _log(f"no such file or registered database: {binary}") + return 2 + elif len(ready) == 1: + item = ready[0] + binary = str(item.get("exe_path") or item.get("idb_path") or "") + _log(f"attaching to registered {item.get('backend')} database: {binary}") + elif not ready: + detail = f" ({blocked[0].get('error')})" if blocked else "" + _log(f"no registered Code Mode database; pass a binary path{detail}") return 2 - if not os.access(os.path.dirname(binary), os.W_OK): - _log(f"directory not writable (IDA writes a .i64 there): " - f"{os.path.dirname(binary)}") + else: + _log("several Code Mode databases are registered; pass one of these paths:") + for item in ready: + _log(f" {item.get('exe_path') or item.get('idb_path')} " + f"[{item.get('backend')}, {item.get('record_id')}]") return 2 - swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged - if swept: - _log(f"cleared {swept} stale lock file(s) from a crashed worker") - # Hand off to the TUI (imported late so --help works without textual). It - # spawns the worker behind its loading overlay while auto-analysis runs. + # Hand off to the TUI (imported late so --help works without Textual). Code + # Mode discovery/opening happens behind its loading overlay. try: from .app import IdaTui except ImportError as e: - _log(f"the TUI needs textual; run with ~/ida-venv/bin/python ({e})") + _log(f"TUI dependencies are missing; run `uv sync` ({e})") return 1 + # Ask the terminal about graphics support NOW: the query needs a reply from + # stdin, and once Textual starts it reads stdin on its own thread and would + # swallow it. Only the ANSWER is wanted here -- the image itself is uploaded + # later, by the splash, because an image uploaded to the primary screen + # cannot be placed once Textual has switched to the alternate one. Costs one + # round trip, and only when attached to a tty. + try: + from . import kittygfx + kittygfx.supported() + except Exception: # noqa: BLE001 -- graphics are decoration, never fatal + pass + rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None IdaTui(open_path=binary, keepalive=not args.no_keepalive, rpc_path=rpc_path, ttl=args.ttl, project=project, - load_args=_load_args(load)).run() + load_args=_load_args(load), + trace_path=(os.path.abspath(os.path.expanduser(args.trace)) + if args.trace else "")).run() return 0 diff --git a/idatui/pane.py b/idatui/pane.py index 53afef3..c31a93c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -1,8 +1,12 @@ -"""Spawn/stop/list idatui TUI panes in tmux, for an agent to drive over RPC. +"""Spawn/stop/list idatui TUI panes in tmux or zellij, for an agent to drive over RPC. -The agent (running inside a tmux pane) can open a fresh pane with the TUI -running against a binary, wait until it's ready, drive it over the RPC socket, -then close it — all without a human touching the keyboard. +The agent (running inside a tmux or zellij pane) can open a fresh pane with the +TUI running against a binary, wait until it's ready, drive it over the RPC +socket, then close it — all without a human touching the keyboard. + +The multiplexer is auto-detected ($ZELLIJ -> zellij, $TMUX -> tmux) and recorded +per pane in the registry, so stop/list/capture/keys keep working across both +(and across a mixed set of panes). $IDATUI_MUX forces a backend. # open a binary in a new pane, block until analysed + drivable, print JSON python -m idatui.pane spawn --open /abs/path/to/bin @@ -15,9 +19,13 @@ 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. 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. + # mux-agnostic screen scrape / key injection (for debugging the input layer) + python -m idatui.pane capture --pane <pane> + python -m idatui.pane keys --pane <pane> Escape + +Requires: running inside tmux or zellij. Each pane leases a registered GUI or +shared managed idalib database through Code Mode. Uses ~/ida-venv/bin/python for +the TUI (needs textual) unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -25,7 +33,6 @@ import argparse import json import os import secrets -import signal import subprocess import sys import time @@ -61,10 +68,47 @@ def _save_registry(rows: list[dict[str, Any]]) -> None: os.replace(tmp, _registry_path()) -def _pane_alive(pane: str) -> bool: - out = subprocess.run(["tmux", "list-panes", "-a", "-F", "#{pane_id}"], - capture_output=True, text=True) - return pane in out.stdout.split() +# --------------------------------------------------------------------------- # +# terminal multiplexer backends +# +# Everything that touches panes goes through here, so the rest of the module (and +# every caller) is mux-agnostic. tmux pane ids look like ``%7``; zellij ids look +# like ``terminal_3``, which is what ``zellij action new-pane`` prints, so a pane +# id alone is enough to route a later stop/capture even if the registry predates +# the ``mux`` field. +# --------------------------------------------------------------------------- # +MUXES = ("tmux", "zellij") + + +def _detect_mux() -> str: + """Which multiplexer we're running under: 'tmux', 'zellij', or '' if neither.""" + forced = os.environ.get("IDATUI_MUX", "").strip().lower() + if forced: + return forced if forced in MUXES else "?" + forced + # Check zellij first: a zellij session started from inside tmux inherits + # $TMUX, and the pane we can actually create there is the zellij one. + if os.environ.get("ZELLIJ"): + return "zellij" + if os.environ.get("TMUX"): + return "tmux" + return "" + + +def _mux_of_pane(pane: str) -> str: + """Infer the backend from a pane id ('%7' = tmux, 'terminal_3' = zellij).""" + if pane.startswith(("terminal_", "plugin_")): + return "zellij" + if pane.startswith("%"): + return "tmux" + return _detect_mux() or "tmux" + + +def _zellij_argv() -> list[str]: + """Base zellij argv, pinned to our session when we know it (so it still works + from a process that isn't itself attached).""" + session = (os.environ.get("IDATUI_ZELLIJ_SESSION") + or os.environ.get("ZELLIJ_SESSION_NAME")) + return ["zellij", "-s", session] if session else ["zellij"] def _tmux(*args: str) -> str: @@ -72,64 +116,168 @@ def _tmux(*args: str) -> str: check=True).stdout.strip() -# --------------------------------------------------------------------------- # -# idalib worker reaping -# -# ``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"idatui/worker\.py" +def _zellij(*args: str) -> str: + return subprocess.run([*_zellij_argv(), *args], capture_output=True, + text=True, check=True).stdout.strip() -def _worker_pids() -> list[int]: - """PIDs of our private per-pane idalib worker processes (idatui/worker.py), - never our own PID.""" +def _zellij_panes() -> list[dict[str, Any]]: try: - out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], + out = subprocess.run([*_zellij_argv(), "action", "list-panes", + "--state", "--json"], capture_output=True, text=True) - except OSError: + rows = json.loads(out.stdout or "[]") + except (OSError, ValueError): return [] - me = os.getpid() - pids: list[int] = [] - for tok in out.stdout.split(): - try: - pid = int(tok) - except ValueError: - continue - if pid != me: - pids.append(pid) - return pids + return rows if isinstance(rows, list) else [] + + +def _pane_alive(pane: str, mux: str | None = None) -> bool: + """True if the pane exists *and* its command is still running. + + zellij keeps an exited pane on screen (EXITED, holding its output) rather + than removing it like tmux does; that husk must count as dead or ``stop`` + would wait out its whole timeout and ``_wait_ready`` would never notice a + launcher that died on startup. + """ + if not pane: + return False + if (mux or _mux_of_pane(pane)) == "zellij": + want = pane.split("_", 1)[-1] + for row in _zellij_panes(): + if str(row.get("id")) == want and bool(row.get("is_plugin")) is False: + return not row.get("exited", False) + return False + out = subprocess.run(["tmux", "list-panes", "-a", "-F", "#{pane_id}"], + capture_output=True, text=True) + return pane in out.stdout.split() + + +def _pane_exists(pane: str, mux: str | None = None) -> bool: + """True if the pane is still on screen at all (including a zellij exit husk).""" + if not pane: + return False + if (mux or _mux_of_pane(pane)) == "zellij": + want = pane.split("_", 1)[-1] + return any(str(r.get("id")) == want and not r.get("is_plugin") + for r in _zellij_panes()) + return _pane_alive(pane, "tmux") + + +def _pane_kill(pane: str, mux: str | None = None) -> None: + """Remove the pane. Idempotent, and also clears a zellij exit husk.""" + if not pane: + return + if (mux or _mux_of_pane(pane)) == "zellij": + subprocess.run([*_zellij_argv(), "action", "close-pane", + "--pane-id", pane], capture_output=True) + else: + subprocess.run(["tmux", "kill-pane", "-t", pane], capture_output=True) + + +def _pane_split(inner: list[str], *, mux: str, vertical: bool, + size: str | None, detached: bool) -> str: + """Open a pane running ``inner`` (argv) in REPO, and return its pane id.""" + if mux == "zellij": + # zellij runs the argv directly (no shell) and takes the cwd as a flag, + # so there's nothing to quote. --name labels the pane in the UI. + argv = [*_zellij_argv(), "action", "new-pane", + "--direction", "down" if vertical else "right", + "--cwd", REPO, "--name", "idatui"] + argv += ["--", *inner] + pane = subprocess.run(argv, capture_output=True, text=True, + check=True).stdout.strip() + # zellij prints the new pane id ('terminal_3'); without it we could not + # target this pane later, so treat a missing id as a hard failure. + if not pane.startswith(("terminal_", "plugin_")): + raise RuntimeError(f"zellij new-pane did not return a pane id: {pane!r}") + if detached: + # zellij always focuses the pane it creates and has no -d; hop back + # to the pane we were called from. + origin = os.environ.get("ZELLIJ_PANE_ID") + if origin: + subprocess.run([*_zellij_argv(), "action", "focus-pane-id", + f"terminal_{origin}"], capture_output=True) + return pane + + cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) + split = ["split-window", "-v" if vertical else "-h", + "-P", "-F", "#{pane_id}"] + if size: + split += ["-l", str(size)] + if detached: + split += ["-d"] + anchor = os.environ.get("TMUX_PANE") + if anchor: + split += ["-t", anchor] + split.append(cmd) + return _tmux(*split) + + +def _pane_capture(pane: str, mux: str | None = None) -> str: + """The pane's visible screen as text.""" + mux = mux or _mux_of_pane(pane) + if mux == "zellij": + return _zellij("action", "dump-screen", "--pane-id", pane) + return _tmux("capture-pane", "-p", "-t", pane) + + +# tmux key names -> zellij key names (zellij rejects e.g. "Escape", wants "Esc"). +_ZELLIJ_KEYS = { + "escape": "Esc", "bspace": "Backspace", "space": "Space", + "pageup": "PageUp", "pagedown": "PageDown", "ppage": "PageUp", + "npage": "PageDown", "ic": "Insert", "dc": "Delete", +} + + +def _to_zellij_key(key: str) -> str: + """Accept tmux-flavoured key names so callers can stay mux-agnostic.""" + low = key.lower() + if low in _ZELLIJ_KEYS: + return _ZELLIJ_KEYS[low] + if len(key) > 2 and key[1] == "-" and key[0] in "CM": # C-a / M-x + return ("Ctrl " if key[0] == "C" else "Alt ") + key[2:] + return key + + +def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: + """Inject real terminal keystrokes into the pane (the input-layer cross-check).""" + mux = mux or _mux_of_pane(pane) + if mux == "zellij": + subprocess.run([*_zellij_argv(), "action", "send-keys", "--pane-id", pane, + *[_to_zellij_key(k) for k in keys]], check=True) + else: + subprocess.run(["tmux", "send-keys", "-t", pane, *keys], check=True) + + +# Code Mode owns database process lifetime: a closed pane drops its lease at the +# socket/kernel boundary and Code Mode decides whether a managed worker still +# has clients. There is nothing for the pane layer to reap. def _count_live_panes() -> int: - return sum(1 for r in _load_registry() if _pane_alive(r.get("pane", ""))) + return sum(1 for r in _load_registry() + if _pane_alive(r.get("pane", ""), r.get("mux"))) def _reap_orphan_workers(force: bool = False) -> int: - """Kill leaked idalib workers when it is safe (no live pane) or ``force``. - - Returns the number of workers signalled. Best-effort; never raises. - """ - if not force and _count_live_panes() > 0: - return 0 - reaped = 0 - for pid in _worker_pids(): - try: - os.kill(pid, signal.SIGKILL) - reaped += 1 - except OSError: - pass - return reaped + """Compatibility no-op: Code Mode workers are shared and lease-managed.""" + del force + return 0 # --------------------------------------------------------------------------- # # spawn # --------------------------------------------------------------------------- # def spawn(args) -> int: - if not os.environ.get("TMUX"): - print("error: not inside tmux (spawn creates a tmux pane)", file=sys.stderr) + mux = args.mux or _detect_mux() + if mux.startswith("?"): + print(f"error: unknown multiplexer {mux[1:]!r} (want tmux or zellij)", + file=sys.stderr) + return 2 + if not mux: + print("error: not inside tmux or zellij (spawn creates a pane there). " + "Set $IDATUI_MUX=tmux|zellij to force a backend.", file=sys.stderr) return 2 if not args.open and not getattr(args, "project", None): print("error: pass --open <binary> or --project <file>", file=sys.stderr) @@ -146,16 +294,8 @@ def spawn(args) -> int: print(f"error: no such project: {project}", file=sys.stderr) return 2 - # Reap workers leaked by previously-stopped/crashed panes so we don't spawn - # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever, - # never reaching ready). No-op while any pane is live. - reaped = _reap_orphan_workers() - if reaped: - print(f"reaped {reaped} orphaned idalib worker(s) before spawn", - file=sys.stderr) - - # 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. + # The pane owns only the TUI. Code Mode's lease cleanup handles crashes; + # kill-pane must never reap a shared GUI/idalib database. if project is not None: # launch takes: --project FILE [binaries...]; extra binaries are added to # the project (and a missing project file is created from them). @@ -165,27 +305,34 @@ def spawn(args) -> int: inner += ["--rpc", sock] else: inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] - cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) + # Loading a headerless blob: without these IDA reads a raw firmware image as + # x86 at 0 and analyses to nothing, and the pane comes up ready-but-empty. + # They are launch's options; spawn just forwards them (a project records + # them per binary, so they're only needed on the first open). + for opt in ("processor", "base", "ida_args"): + val = getattr(args, opt, None) + if val: + inner += ["--" + opt.replace("_", "-"), str(val)] + if getattr(args, "trace", None): + inner += ["--trace", os.path.abspath(os.path.expanduser(args.trace))] - split = ["split-window", "-v" if args.vertical else "-h", - "-P", "-F", "#{pane_id}"] - if args.size: - split += ["-l", str(args.size)] - if args.detached: - split += ["-d"] - anchor = os.environ.get("TMUX_PANE") - if anchor: - split += ["-t", anchor] - split.append(cmd) - pane = _tmux(*split) + if args.size and mux == "zellij": + print("note: --size is tmux-only; zellij tiles the new pane evenly", + file=sys.stderr) + try: + pane = _pane_split(inner, mux=mux, vertical=args.vertical, + size=args.size, detached=args.detached) + except (OSError, subprocess.CalledProcessError, RuntimeError) as e: + print(f"error: could not create a {mux} pane: {e}", file=sys.stderr) + return 2 - row = {"sock": sock, "pane": pane, "target": project or target, + row = {"sock": sock, "pane": pane, "mux": mux, "target": project or target, "kind": "project" if project else "open", "started": time.time()} reg = [r for r in _load_registry() if r.get("sock") != sock] reg.append(row) _save_registry(reg) - ready = _wait_ready(sock, args.timeout, pane) + ready = _wait_ready(sock, args.timeout, pane, mux=mux) row.update(ready) print(json.dumps(row)) return 0 if ready.get("ready") else 1 @@ -197,19 +344,18 @@ def _q(s: str) -> str: def _wait_ready(sock: str, timeout: float, pane: str, - stuck_after: float = 45.0) -> dict[str, Any]: + stuck_after: float = 45.0, mux: str | None = None) -> dict[str, Any]: """Poll the socket + ping until the TUI reports ready (or timeout). - Emits a one-time hint to stderr if it's still not ready after ``stuck_after`` - seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic - instead of an unexplained silent hang. + Emits a one-time hint if Code Mode discovery/opening is still not ready after + ``stuck_after`` seconds. """ start = time.time() deadline = start + timeout warned = False last: dict[str, Any] = {"ready": False} while time.time() < deadline: - if not _pane_alive(pane): + if not _pane_alive(pane, mux): return {"ready": False, "error": "pane exited during startup"} if os.path.exists(sock): try: @@ -223,9 +369,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, warned = True 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 — try " - f"`python -m idatui.pane reap`.", file=sys.stderr) + print(f"still waiting ({int(time.time() - start)}s): {why}. " + f"Check Code Mode registrations and worker logs.", file=sys.stderr) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -246,18 +391,32 @@ def stop(args) -> int: if not rows: print("error: no matching pane (need --sock or --pane)", file=sys.stderr) return 2 + killed: list[str] = [] for r in rows: - sock, pane = r.get("sock"), r.get("pane") + sock, pane, mux = r.get("sock"), r.get("pane"), r.get("mux") + quit_ok = False if sock and os.path.exists(sock): try: # ask it to quit gracefully first with RpcClient(sock) as c: c.call("quit") - time.sleep(0.4) + quit_ok = True except (OSError, RpcError, ConnectionError): pass - if pane and _pane_alive(pane): - subprocess.run(["tmux", "kill-pane", "-t", pane], - capture_output=True) + # Wait for the pane to actually go away. Quitting runs App.on_unmount, + # which writes every dirty database; a 90 MB .i64 takes tens of seconds. + # Killing the pane on a fixed short sleep truncated that save and + # silently destroyed the session's work, so block on the real signal. + if pane and quit_ok: + deadline = time.monotonic() + float(args.timeout) + while time.monotonic() < deadline and _pane_alive(pane, mux): + time.sleep(0.25) + if pane: + # Still running past the timeout = force kill (and warn). Otherwise + # it exited cleanly, but under zellij the pane lingers as an exit + # husk, so close it either way to leave the layout as we found it. + if _pane_alive(pane, mux): + killed.append(pane) + _pane_kill(pane, mux) if sock: try: os.unlink(sock) @@ -269,6 +428,12 @@ def stop(args) -> int: out = {"stopped": [r.get("sock") or r.get("pane") for r in rows]} if reaped: out["reaped_workers"] = reaped + if killed: + # Only ever reached on timeout: say so, because it means a save may have + # been cut short rather than "clean teardown". + out["force_killed"] = killed + out["warning"] = (f"pane(s) did not exit within {args.timeout}s and were " + "killed; unsaved database changes may be lost") print(json.dumps(out)) return 0 @@ -278,7 +443,8 @@ def list_panes(args) -> int: alive = [] for r in reg: r = dict(r) - r["pane_alive"] = _pane_alive(r.get("pane", "")) + r.setdefault("mux", _mux_of_pane(r.get("pane", ""))) + r["pane_alive"] = _pane_alive(r.get("pane", ""), r.get("mux")) r["sock_up"] = bool(r.get("sock") and os.path.exists(r["sock"])) if args.prune and not r["pane_alive"]: if r.get("sock") and os.path.exists(r["sock"]): @@ -286,6 +452,9 @@ def list_panes(args) -> int: os.unlink(r["sock"]) except OSError: pass + # a zellij pane whose command exited is still on screen; drop it + if r.get("pane") and _pane_exists(r["pane"], r.get("mux")): + _pane_kill(r["pane"], r.get("mux")) continue alive.append(r) if args.prune: @@ -298,32 +467,95 @@ def list_panes(args) -> int: def reap(args) -> int: - """Kill leaked idalib workers (safe when no pane is live; --force overrides).""" - live = _count_live_panes() - n = _reap_orphan_workers(force=args.force) - print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force})) - if n == 0 and not args.force and live > 0: - print(f"note: {live} live pane(s) — not reaping in-use workers; pass " - f"--force to reap anyway", file=sys.stderr) + """Deprecated no-op; shared Code Mode workers are managed by leases.""" + print(json.dumps({"reaped_workers": 0, "live_panes": _count_live_panes(), + "forced": args.force, "deprecated": True})) + return 0 + + +def capture(args) -> int: + """Print a pane's visible screen (tmux capture-pane / zellij dump-screen).""" + pane = args.pane or _resolve_pane(args.sock) + if not pane: + return 2 + try: + print(_pane_capture(pane, args.mux or None)) + except (OSError, subprocess.CalledProcessError) as e: + print(f"error: could not capture {pane}: {e}", file=sys.stderr) + return 1 + return 0 + + +def send_keys(args) -> int: + """Inject real terminal keystrokes (tmux send-keys / zellij send-keys). + + Key names are tmux-flavoured and translated per backend, so `keys --pane P + Escape` does the right thing under either mux. + """ + pane = args.pane or _resolve_pane(args.sock) + if not pane: + return 2 + try: + _pane_keys(pane, args.keys, args.mux or None) + except (OSError, subprocess.CalledProcessError) as e: + print(f"error: could not send keys to {pane}: {e}", file=sys.stderr) + return 1 return 0 +def _resolve_pane(sock: str | None) -> str | None: + """Pane id for a socket, or the single live pane if there's exactly one.""" + reg = _load_registry() + if sock: + for r in reg: + if r.get("sock") == sock: + return r.get("pane") + print(f"error: no tracked pane for {sock}", file=sys.stderr) + return None + live = [r for r in reg if _pane_alive(r.get("pane", ""), r.get("mux"))] + if len(live) == 1: + return live[0].get("pane") + if not live: + print("error: no live panes (pass --pane)", file=sys.stderr) + else: + print("error: several live panes, pass --pane or --sock:", file=sys.stderr) + for r in live: + print(f" {r.get('pane')} {r.get('sock')} {r.get('target')}", + file=sys.stderr) + return None + + def main(argv: list[str]) -> int: - p = argparse.ArgumentParser(prog="idatui.pane", - description="spawn/manage idatui TUI panes in tmux") + p = argparse.ArgumentParser( + prog="idatui.pane", + description="spawn/manage idatui TUI panes in tmux or zellij") 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 (its dir must be writable)") + sp.add_argument("--trace", metavar="FILE", + help="Tenet execution trace to load alongside the binary") sp.add_argument("--project", metavar="FILE", help="project file to open instead of a single binary; " "any --open paths are added to it (created if absent)") + sp.add_argument("--processor", metavar="NAME", + help="IDA processor for a headerless blob: arm, armb, " + "mipsb, metapc, … (passed to idatui.launch)") + sp.add_argument("--base", metavar="ADDR", + help="load address for a headerless blob, e.g. 0x8000000 " + "(16-byte aligned)") + sp.add_argument("--ida-args", metavar="STR", dest="ida_args", + help="extra IDA command-line switches, passed through") 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("--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("--size", help="new pane size (tmux -l value, e.g. 60%% or 120; " + "ignored under zellij)") sp.add_argument("--detached", action="store_true", help="don't focus the new pane") + sp.add_argument("--mux", choices=MUXES, default="", + help="multiplexer to spawn in (default: autodetect from " + "$ZELLIJ/$TMUX; $IDATUI_MUX overrides)") sp.add_argument("--timeout", type=float, default=300.0, help="seconds to wait for readiness (fresh --open analysis is slow)") sp.set_defaults(fn=spawn) @@ -331,17 +563,33 @@ def main(argv: list[str]) -> int: st = sub.add_parser("stop", help="graceful quit + kill the pane") st.add_argument("--sock") st.add_argument("--pane") + st.add_argument("--timeout", type=float, default=600.0, + help="seconds to wait for the pane to exit (it saves dirty " + "databases on the way out) before force-killing it") st.set_defaults(fn=stop) ls = sub.add_parser("list", help="list tracked panes") ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)") ls.set_defaults(fn=list_panes) - rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)") - rp.add_argument("--force", action="store_true", - help="reap even while panes are live (may kill an in-use analyser)") + rp = sub.add_parser("reap", help="deprecated no-op (Code Mode uses shared leases)") + rp.add_argument("--force", action="store_true", help=argparse.SUPPRESS) rp.set_defaults(fn=reap) + cp = sub.add_parser("capture", help="print a pane's visible screen") + cp.add_argument("--pane") + cp.add_argument("--sock", help="resolve the pane from this socket") + cp.add_argument("--mux", choices=MUXES, default="") + cp.set_defaults(fn=capture) + + kp = sub.add_parser("keys", help="inject real keystrokes into a pane " + "(tmux-style names, translated per mux)") + kp.add_argument("keys", nargs="+", help="e.g. Escape, Enter, C-a, g m a i n") + kp.add_argument("--pane") + kp.add_argument("--sock", help="resolve the pane from this socket") + kp.add_argument("--mux", choices=MUXES, default="") + kp.set_defaults(fn=send_keys) + args = p.parse_args(argv) return args.fn(args) diff --git a/idatui/pool.py b/idatui/pool.py index ae37c25..465dff2 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -1,23 +1,16 @@ -"""WorkerPool — keeps a live idalib worker per project binary, within a budget. +"""DatabasePool — LRU leases on Code Mode databases for a project. -One worker process holds exactly one database (idalib is single-DB and -main-thread-only), so a project with N binaries means up to N processes. They are -not cheap and they do not share: a worker on ``bash`` measures ~126 MB RSS / -117 MB PSS, and the database working set dominates for anything larger -(``libcrypto.so.3``'s ``.i64`` alone is 72 MB). +Code Mode may bind a lease to an existing IDA GUI or to a shared managed idalib +worker. The pool therefore owns *client interest*, never an IDA process. Releasing +an LRU entry persists managed IDBs but does not implicitly save a GUI, then closes +only this TUI's lease; other clients and GUI sessions remain alive. Managed workers exit themselves after their final lease. -Residency is therefore bounded by a **memory budget**, not a worker count — a -count is the wrong knob when one project holds both a 50 KB helper and a 6 MB -crypto library. Workers are spawned lazily on first use, kept resident while they -fit, and least-recently-used ones evicted when they don't. Eviction **saves the -database first**, so coming back is a load rather than a re-analysis. - -The pool never evicts the active binary, nor anything pinned. +The historical memory budget remains useful for managed idalib instances, while +GUI process memory is only advisory. The active and pinned databases are never +released to satisfy it. """ from __future__ import annotations -import os - from .project import BinaryRef, Project #: Fallback budget if /proc/meminfo can't be read (MB). @@ -36,11 +29,11 @@ def _total_ram_mb() -> int: def _pss_mb(pid: int | None) -> int: - """Proportional set size of a worker, in MB. + """Proportional set size of the leased instance process, in MB. - PSS (not RSS) is the honest per-worker cost: it splits shared pages between - the processes mapping them. In practice workers share very little, so the two - are close, but PSS is what makes summing across workers meaningful. + PSS is useful for managed idalib workers. For GUI/shared processes it is only + advisory because the TUI neither owns all that memory nor controls process + exit. """ if not pid: return 0 @@ -54,13 +47,19 @@ def _pss_mb(pid: int | None) -> int: return 0 -def _default_spawn(ref: BinaryRef, ttl: int): # pragma: no cover - needs idalib - from .worker_client import WorkerClient - return WorkerClient(ref.staged, ttl=ttl, load_args=ref.load_args) +def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA + from .codemode_client import CodeModeClient + return CodeModeClient( + ref.staged, + ttl=ttl, + load_args=ref.load_args, + output_database=ref.db, + new_database=new_database, + ) -class WorkerPool: - """Live workers for a project's binaries, keyed by label.""" +class DatabasePool: + """Live Code Mode database leases, keyed by project label.""" def __init__(self, project: Project, *, budget_mb: int | None = None, ttl: int = 1800, spawn=None, mem_fn=None) -> None: @@ -71,6 +70,7 @@ class WorkerPool: self._clients: dict[str, object] = {} self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() + self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() @@ -98,11 +98,11 @@ class WorkerPool: # -- acquire ----------------------------------------------------------- # def get(self, label: str, progress=None): - """A live client for ``label``, spawning it (and making room) if needed. + """A live client for ``label``, attaching or spawning as needed. - Staging and the scratch sweep happen here: a worker killed hard last time - leaves unpacked ``.id0/.id1/...`` behind, and the database then refuses to - reopen. Nothing else holds this DB (one worker per label), so it is safe. + Do not sweep IDA scratch files here: a registered GUI or another Code + Mode client may own the database. Code Mode's registry locks and health + probes are the authority for safe discovery and stale-record cleanup. """ client = self._clients.get(label) if client is not None: @@ -118,28 +118,29 @@ class WorkerPool: note(f"staging {ref.label}\u2026") self.project.stage(ref) - self.project.sweep_scratch(ref) note(f"opening {ref.label}\u2026") - client = self._spawn(ref, self._ttl) + fresh = label in self._recreate + client = (_default_spawn(ref, self._ttl, new_database=fresh) + if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() self._clients[label] = client + self._recreate.discard(label) self._lru.append(label) self._enforce_budget(protect=label) return client def prewarm(self, label: str, progress=None) -> bool: - """Spawn a worker for ``label`` only if it fits the budget AS IT STANDS. + """Attach a database for ``label`` only if it fits the current budget. Pre-warming must never cost residency: evicting a binary the user actually visited to speculatively load one they haven't is a straight downgrade, and the eviction would also throw away that binary's caches. So this refuses rather than making room, and returns False. - The cost of a worker that doesn't exist yet can only be estimated; the - largest resident one is the best evidence available (they are all the - same program with a different database). With nothing resident we have + The cost of a database not attached yet can only be estimated; the + largest resident instance is the best evidence available. With nothing resident we have no evidence at all, so we allow one — that is the case where the budget is certainly free. """ @@ -153,13 +154,19 @@ class WorkerPool: return False self.get(label, progress=progress) # get() enforces the budget protecting the NEW label; if that had to - # evict, our estimate was wrong and the speculative worker is the one + # evict, our estimate was wrong and the speculative lease is the one # that should go — never a binary the user chose. if self.memory_mb() > self.budget_mb and label != self.active: self.evict(label) return False return True + def recreate_on_next_open(self, label: str) -> None: + """Request a fresh IDB after the current lease has been released.""" + if self.project.by_label(label) is None: + raise KeyError(f"no such binary in the project: {label}") + self._recreate.add(label) + def _touch(self, label: str) -> None: if label in self._lru: self._lru.remove(label) @@ -171,16 +178,21 @@ class WorkerPool: self._touch(label) # -- release ----------------------------------------------------------- # - def evict(self, label: str, save: bool = True) -> bool: - """Drop a resident worker, persisting its database first.""" + def evict(self, label: str, save: bool = True, + save_gui: bool = False) -> bool: + """Release a resident lease, persisting a managed database first. + + A budget-driven eviction must not save somebody's GUI implicitly. GUI + saves are reserved for an explicit/defensive ``close_all(save=True)``. + """ client = self._clients.pop(label, None) if client is None: return False if label in self._lru: self._lru.remove(label) - if save: + if save and (save_gui or getattr(client, "backend", None) != "gui"): try: # persist analysis + edits so the next open is a load - client.call("idb_save") + client.save_database() except Exception: # noqa: BLE001 -- evict regardless pass try: @@ -198,7 +210,7 @@ class WorkerPool: return None def _enforce_budget(self, protect: str | None = None) -> int: - """Evict LRU workers until the pool fits its budget. Returns how many.""" + """Release LRU leases until the pool fits its budget. Returns how many.""" n = 0 while self.memory_mb() > self.budget_mb: victim = self._evictable(protect) @@ -210,7 +222,7 @@ class WorkerPool: def close_all(self, save: bool = True) -> None: for label in list(self._clients): - self.evict(label, save=save) + self.evict(label, save=save, save_gui=save) self.active = None # -- introspection ------------------------------------------------------ # @@ -231,5 +243,9 @@ class WorkerPool: return out def __repr__(self) -> str: # pragma: no cover - debug aid - return (f"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + return (f"<DatabasePool {len(self._clients)}/{len(self.project.refs)} resident " f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") + + +# Source compatibility for callers that imported the pre-Code-Mode name. +WorkerPool = DatabasePool diff --git a/idatui/project.py b/idatui/project.py index e2fc542..53f3b0a 100644 --- a/idatui/project.py +++ b/idatui/project.py @@ -23,7 +23,8 @@ firmware image, a cleaned build tree). A source whose size/mtime no longer matches the staged copy is re-staged, and its now-stale database is dropped (the DB describes the old bytes). -stdlib-only, like the domain/worker layers — the TUI is the only Textual consumer. +The model has no IDA imports. Staging consults ida_codemode's registry before +replacing files so it never mutates a database owned by a GUI/shared worker. """ from __future__ import annotations @@ -56,7 +57,7 @@ class BinaryRef: #: it is, a raw firmware image doesn't, and IDA defaults to metapc at 0. processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … base: int = 0 # load address (natural, e.g. 0x8000000) - ida_args: str = "" # escape hatch: extra IDA command-line switches + ida_args: str = "" # legacy -p/-b/-T switches accepted by Code Mode adapter @property def db(self) -> str: @@ -310,13 +311,25 @@ class Project: """Ensure ``ref`` is staged in the sidecar; returns the staged path. Re-staging a changed source drops its database: the DB describes the old - bytes, so keeping it would silently mismatch the disassembly (any renames - in it are lost, which is why callers should say so out loud). + bytes. Refuse while Code Mode reports a GUI/idalib owner; replacing a + staged executable or IDB underneath a shared live instance is corruption. """ if not os.path.isfile(ref.source): raise ProjectError(f"no such binary: {ref.source}") if not self.is_stale(ref): return ref.staged + try: + from .codemode_client import database_owner + owner = database_owner(ref.db, ref.staged) + except Exception as exc: + raise ProjectError( + f"cannot verify Code Mode ownership before staging {ref.label}: {exc}" + ) from exc + if owner is not None: + raise ProjectError( + f"cannot restage {ref.label}: Code Mode instance {owner.record_id} " + f"still owns {owner.idb_path}; close/release it first" + ) os.makedirs(self.bin_dir, exist_ok=True) tmp = ref.staged + ".staging" _unlink(tmp) @@ -338,10 +351,11 @@ class Project: return out def sweep_scratch(self, ref: BinaryRef) -> int: - """Delete IDA's unpacked working files (never the ``.i64``) for ``ref``. + """Delete unpacked working files (never the ``.i64``) for maintenance. - A hard-killed worker leaves them behind and the database then refuses to - reopen. Only safe when no worker holds it. + Runtime paths no longer call this: Code Mode instances are shared, so a + registry owner may still be using these files. Callers must independently + prove that no GUI/idalib instance owns the database. """ return sum(1 for suf in SCRATCH_SUFFIXES if _unlink(ref.staged + suf)) diff --git a/idatui/prompt.py b/idatui/prompt.py new file mode 100644 index 0000000..74956f6 --- /dev/null +++ b/idatui/prompt.py @@ -0,0 +1,117 @@ +"""The one-line prompts along the bottom of the app. + +There are six (`search` `rename` `comment` `retype` `makedata` `goto`) and every +one of them was opened, closed and escaped by its own copy of the same six +lines. The copies had drifted: some restored focus to the view, some didn't; +some cleared their context on close, some left it for the next caller to trip +over. + +Opening a prompt is: hide the status bar (they share a row), set the +placeholder and prefill, make it focusable, show it, focus it. Closing is the +same in reverse, plus handing focus back to whatever view asked for it. That is +all `Prompt` is -- but having it in one place is what makes "Esc closes whatever +is open" a loop instead of a six-branch ladder in `on_key`. + +Note the `can_focus` toggling: a hidden `Input` that stays focusable still takes +part in Tab focus-nav, so tabbing around a closed prompt used to land the cursor +in an invisible widget and swallow every subsequent keystroke. +""" +from __future__ import annotations + +from typing import TYPE_CHECKING + +from textual.widgets import Input, Static + +if TYPE_CHECKING: # pragma: no cover + from textual.app import App + + +class Prompt: + """One `Input` at the bottom of the screen, plus who to give focus back to. + + ``ctx`` is whatever the opener needs when the value is submitted (the view + that asked, the address being commented, the kind of retype...). It is held + here rather than in a parallel ``_rename_ctx`` attribute on the app so that + it cannot outlive the prompt that owns it. + """ + + def __init__(self, app: "App", ident: str) -> None: + self.app = app + self.id = ident + self.ctx: object = None + + @property + def input(self) -> Input: + return self.app.query_one(f"#{self.id}", Input) + + @property + def open(self) -> bool: + """Is this prompt on screen? Cheap enough to poll in `on_key`.""" + try: + return bool(self.input.display) + except Exception: # noqa: BLE001 -- a modal owns the screen + return False + + def show(self, placeholder: str, value: str = "", ctx: object = None) -> Input: + """Put the prompt up, prefilled and focused.""" + self.ctx = ctx + self.app.query_one("#status", Static).display = False + inp = self.input + inp.placeholder = placeholder + inp.can_focus = True + inp.display = True + inp.value = value + inp.focus() + return inp + + def close(self, refocus: bool = True) -> object: + """Take the prompt down and return the context it was holding. + + Returned rather than left readable, because every caller wants it + exactly once and the old parallel-attribute version kept handing the + next caller a stale one (the listing's `_rename_addr` had to be captured + by hand before `_end_rename` cleared it, or the name went to address 0). + """ + ctx, self.ctx = self.ctx, None + inp = self.input + inp.display = False + inp.can_focus = False + self.app.query_one("#status", Static).display = True + if refocus: + view = ctx[0] if isinstance(ctx, tuple) and ctx else None + if view is not None and hasattr(view, "focus"): + view.focus() + return ctx + + +class PromptBar: + """Every prompt the app owns, so "close whatever is open" is one call.""" + + def __init__(self, app: "App", *idents: str) -> None: + self.app = app + self._order = list(idents) + self._by_id = {i: Prompt(app, i) for i in idents} + + def __getitem__(self, ident: str) -> Prompt: + return self._by_id[ident] + + def __getattr__(self, name: str) -> Prompt: + try: + return self.__dict__["_by_id"][name] + except KeyError: + raise AttributeError(name) from None + + def active(self) -> Prompt | None: + """The prompt currently on screen, in declaration order. + + Only one is ever up -- they share the row above the footer -- but the + order is fixed anyway so this can't depend on dict iteration. + """ + for ident in self._order: + p = self._by_id[ident] + if p.open: + return p + return None + + def any_open(self) -> bool: + return self.active() is not None diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py new file mode 100644 index 0000000..6fb6436 --- /dev/null +++ b/idatui/remote_tools.py @@ -0,0 +1,1548 @@ +"""The IDAPython ida-tui runs inside the Code Mode sandbox. + +Two features have no ida-domain surface at all and are carried over VERBATIM +from the tools ida-tui was developed against (`server/patch_server.py`'s +injected BODY, which the Code Mode port deletes): + +* `heads` -- the continuous listing. ida-domain enumerates defined heads and + renders plain disassembly; the listing also needs coalesced undefined runs, + IDA colour-tag spans, PER-OPERAND EXTENTS, function banners, code labels and + expanded struct members, plus the digest/`expect` protocol the paging layer + uses to skip re-sending a page that has not changed. +* `op_format` / `pc_nums` / `pc_num_format` -- `o`/`O`. IDA's operand types and + Hex-Rays' per-(ea, opnum) numforms are separate sets, and neither is exposed. + +Keeping the originals rather than paraphrasing them is deliberate: this is the +most performance-tuned and most behaviour-sensitive code in the project (the +span walker is a single regex pass because a per-character loop was the most +expensive thing the listing did, and the cycle only offers stops that change +what you see). A re-implementation drifts from it silently. + +This file is SOURCE SHIPPED AS TEXT to the database process; it is never +imported here, because the ida_* modules do not exist in the TUI's interpreter. +`codemode_client` reads it and prepends it to the relevant snippets. Keep it +self-contained: no relative imports, nothing beyond what Code Mode provides. +""" +# ruff: noqa +import re as _re + +from typing import Annotated # the extracted tool signatures still carry these + + +class IDAError(Exception): + """The MCP host's error type; the tools raise/catch it by name.""" + + +def parse_address(addr): + """ida_pro_mcp.utils.parse_address: hex/decimal string, int, or a symbol.""" + if isinstance(addr, int): + return addr + try: + return int(addr, 0) + except ValueError: + import idaapi + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + if ea != idaapi.BADADDR: + return ea + raise IDAError(f"Not found: {addr!r}") + + +#: Byte-identical to ida_pro_mcp.utils._STRING_OR_SPACES_RE: the pseudocode +#: column coordinates the client holds depend on collapsing exactly the same way. +_IDATUI_STRING_OR_SPACES_RE = _re.compile( + r'"(?:[^"\\]|\\.)*"' # double-quoted string + r"|'(?:[^'\\]|\\.)*'" # single-quoted string / char + r"|[ \t]{2,}" # run of 2+ whitespace (outside strings) +) + + +def compact_whitespace(line: str) -> str: + """ida_pro_mcp.utils.compact_whitespace: collapse runs of 2+ spaces/tabs to + one, preserving string literals.""" + stripped = line.lstrip(" \t") + if not stripped: + return line + lead = line[: len(line) - len(stripped)] + + def _repl(m): + s = m.group() + if s[0] in ('"', "'"): + return s # preserve string content + return " " + + return lead + _IDATUI_STRING_OR_SPACES_RE.sub(_repl, stripped) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + if ida_bytes.is_code(f): + kind = "code" + elif ida_bytes.is_data(f): + kind = "data" + else: + kind = "unknown" + line = ida_lines.generate_disasm_line(ea, 0) + text, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + row["spans"] = spans + # Where each operand sits in `text`. Comes out of the same tag walk + # (free), and is what lets the client show WHICH literal a keypress + # would reformat before you press it. + if ops: + row["ops"] = ops + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + if ida_bytes.is_code(f): + kind = "code" + elif ida_bytes.is_data(f): + kind = "data" + else: + kind = "unknown" + line = ida_lines.generate_disasm_line(ea, 0) + text, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + row["spans"] = spans + # Where each operand sits in `text`. Comes out of the same tag walk + # (free), and is what lets the client show WHICH literal a keypress + # would reformat before you press it. + if ops: + row["ops"] = ops + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +@_idatui_functools.lru_cache(maxsize=_IDATUI_LINE_CACHE) +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +_IDATUI_SPAN_KINDS = { + "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), + "reg": ("SCOLOR_REG",), + "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), + "str": ("SCOLOR_STRING",), + # NB the real constant names: DATNAME/CODNAME, not "DNAME". Guessing here + # fails silently — an unmapped tag renders as plain body text, so symbols + # just quietly aren't blue and nothing tells you why. + "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", + "SCOLOR_IMPNAME", "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", + "SCOLOR_CNAME", "SCOLOR_DNAME", + "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), + "seg": ("SCOLOR_SEGNAME",), + "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), + "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), + "err": ("SCOLOR_ERROR",), +} + + +def _idatui_tag_map(): + """{tag character: kind}, built once from whatever this IDA actually has.""" + import ida_lines + out = {} + for kind, names in _IDATUI_SPAN_KINDS.items(): + for n in names: + v = getattr(ida_lines, n, None) + if isinstance(v, str) and v: + out[v[0]] = kind + elif isinstance(v, int): + out[chr(v)] = kind + return out + + +_IDATUI_TAGS = None + + +_IDATUI_OPND_TAGS = None + + +_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument + + +_IDATUI_TAGINFO = None + + +def _idatui_opnd_tag_map(): + """{tag character: operand index}. IDA wraps each operand of a disassembly + line in COLOR_OPND1..8, so the line already says where operand N starts and + ends -- no need to re-render operands with print_operand to find out (and + the two agree exactly; checked over thousands of instructions).""" + import ida_lines + out = {} + for i in range(1, 9): + v = getattr(ida_lines, "COLOR_OPND%d" % i, None) + if isinstance(v, int): + out[chr(v)] = i - 1 + elif isinstance(v, str) and v: + out[v[0]] = i - 1 + return out + + +def _idatui_spans(line): + """(spans, ops) for a tagged disasm line. + + ``spans`` is [[kind, text], ...] with colour tags resolved; ``ops`` is + [[start, end, n], ...], the extent of each operand in the SAME (collapsed) + coordinates the row's ``text`` uses -- which is what lets a cursor column + name the operand it is standing on. + + Unknown tags become 'text' rather than being dropped: a processor module can + emit a colour we don't classify, and losing the characters would corrupt the + line.""" + global _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL, _IDATUI_TAGINFO + import ida_lines + if _IDATUI_TAGS is None: + _IDATUI_TAGS = _idatui_tag_map() + if _IDATUI_OPND_TAGS is None: + _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() + if _IDATUI_CTL is None: + import re as _re + # One capturing split gives [text, tag, text, tag, ..., text] in a + # single C pass. A per-character python loop over the line used to be + # the most expensive thing the `heads` tool did, and a line is ~54 + # characters but only ~13 tags -- everything between two tags is already + # exactly one span's worth of text. + _IDATUI_CTL = _re.compile("([\\x01\\x02\\x03](?s:.))") + if _IDATUI_TAGINFO is None: + _IDATUI_TAGINFO = { + tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag)) + for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS)} + taginfo = _IDATUI_TAGINFO + plain_tag = ("text", None) + on, off, esc = "\x01", "\x02", "\x03" + addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) + addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) + parts = _IDATUI_CTL.split(line) + spans, stack = [], [] # stack entries: (kind, operand index|None) + kind, opnd = "text", None # state the current run of text belongs to + pend = "" + skip = 0 # characters of an address payload still due + i, n = 0, len(parts) + while i < n: + txt = parts[i] + i += 1 + if skip: + if len(txt) <= skip: + skip -= len(txt) + txt = "" + else: + txt = txt[skip:] + skip = 0 + if txt: + pend += txt + if i >= n: + break + pair = parts[i] + i += 1 + if skip: # a tag INSIDE an address payload: 2 chars + skip = skip - 2 if skip > 2 else 0 + continue + ch = pair[0] + if ch == esc: # escaped literal: keep the char it guards + pend += pair[1] + continue + tag = pair[1] + if ch == on and tag == addr_tag: + # An embedded target address, not display text: 16 hex digits that + # must not reach the screen. Deliberately NOT a span boundary. + skip = addr_len + continue + if pend: + spans.append([kind, pend, opnd]) + pend = "" + if ch == on: + stack.append((kind, opnd)) + kind, o = taginfo.get(tag, plain_tag) + if o is not None: + opnd = o # operands nest: an inner colour keeps the operand + elif stack: + kind, opnd = stack.pop() + else: + kind, opnd = "text", None + if pend: + spans.append([kind, pend, opnd]) + # Collapse IDA's column padding EXACTLY as the plain text does. A run of + # spaces can straddle two spans, so the leading space of a span is dropped + # when the previous one ended in space — otherwise the spans and `text` + # disagree about the line and the row silently loses its highlighting. + # ``" ".join(txt.split())`` splits on exactly what str.isspace() calls + # whitespace, which is what the character walk this replaces tested. + out, prev_space = [], False + for kind, txt, opnd in spans: + core = " ".join(txt.split()) + if core == txt: + # Nothing to collapse and no edge whitespace -- which is the common + # case ("mov", "rax", ", ") and skips both isspace() probes below. + prev_space = False + out.append([kind, txt, opnd]) + continue + if not core: # the span is nothing but padding + if not prev_space: + prev_space = True + out.append([kind, " ", opnd]) + continue + acc = core + if txt[0].isspace() and not prev_space: + acc = " " + acc + if txt[-1].isspace(): + acc += " " + prev_space = acc[-1] == " " + out.append([kind, acc, opnd]) + while out and out[0][1] == " ": + out.pop(0) + while out and out[-1][1] == " ": + out.pop() + if out and out[0][1].startswith(" "): + out[0][1] = out[0][1].lstrip() + if out and out[-1][1].endswith(" "): + out[-1][1] = out[-1][1].rstrip() + out = [s for s in out if s[1]] + # Operand extents, in the coordinates of the collapsed text these spans + # spell out. Adjacent spans of the same operand merge, so an operand like + # ``[rbp+var_40]`` (five differently-coloured tokens) comes back as ONE + # range -- which is the thing a cursor is inside of, and the thing a format + # change applies to. + ops, pos, cur, start = [], 0, None, 0 + for _kind, txt, opnd in out: + if opnd != cur: + if cur is not None and pos > start: + ops.append([start, pos, cur]) + cur, start = opnd, pos + pos += len(txt) + if cur is not None and pos > start: + ops.append([start, pos, cur]) + text = "".join(t for _k, t, _o in out) + trimmed = [] + for lo, hi, k in ops: # don't let a range own trailing space + while hi > lo and text[hi - 1].isspace(): + hi -= 1 + while lo < hi and text[lo].isspace(): + lo += 1 + if hi > lo: + trimmed.append([lo, hi, k]) + return [[k, t] for k, t, _o in out], trimmed + + +def _idatui_rows_digest(rows): + """A value that changes whenever any of ``rows`` would render differently. + + Covers everything a client keeps off a row: address, kind, size, the plain + text, the symbol name and the colour spans (which is what makes it exact + rather than a heuristic -- two lines can collapse to the same text and still + be coloured differently). + + Uses the interpreter's own ``hash``, deliberately. It never has to mean + anything outside this process: the client stores what a page hashed to when + it loaded it and hands the same number back to ask whether the page still + hashes to that. One worker, one process, one hash seed. + """ + acc = 0 + # The per-line render is memoised, so one spans list is shared by every row + # that says the same thing -- about 45% of them within a page. Hash each + # distinct list once and key that by identity, rather than rebuilding a + # tuple of tuples per row (which is the exact cost that was measured and + # removed from the client side for the same reason). + seen = {} + for r in rows: + sp = r.get("spans") + if sp is None: + sh = None + else: + key = id(sp) + sh = seen.get(key) + if sh is None: + sh = seen[key] = hash(tuple(map(tuple, sp))) + acc = hash((acc, r.get("ea"), r.get("kind"), r.get("size"), + r.get("text"), r.get("name"), sh)) + return acc + + +def _idatui_unknown_row(ea, size): + """One collapsed row for a run of ``size`` undefined bytes starting at + ``ea``. A single byte is rendered normally (shows its value); a longer run + collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions + of one-byte rows.""" + import ida_name + + if size <= 1: + return _idatui_head_row(ea) + row = {"ea": hex(ea), "kind": "unknown", "size": int(size), + "text": f"db {size} dup(?)"} + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +def _idatui_struct_member_rows(ea): + """Indented member rows for a struct-typed data item at ``ea`` (expansion), + or [] if it isn't a struct. Top-level fields only.""" + import ida_nalt + import ida_typeinf + import idaapi + + tif = ida_typeinf.tinfo_t() + if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): + return [] + udt = ida_typeinf.udt_type_data_t() + if not tif.get_udt_details(udt): + return [] + rows = [] + for m in udt: + off = m.begin() // 8 + try: + mtype = m.type._print() or "" + except Exception: + mtype = "" + try: + sz = int(m.type.get_size()) + if sz == idaapi.BADSIZE: + sz = 0 + except Exception: + sz = 0 + name = m.name or "" + text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") + rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, + "text": text}) + return rows + + +def _idatui_func_header_rows(ea): + """IDA-style subroutine banner rows shown just before a function's entry.""" + import ida_funcs + + name = ida_funcs.get_func_name(ea) or "sub_%X" % ea + bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 + return [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + + +def _idatui_func_footer_rows(ea, func): + """End-of-function marker shown just after a function's last item.""" + import ida_funcs + + name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea + return [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + + +def heads( + addr: Annotated[str, "Start address or name to walk from"], + count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, + offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, + end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", + back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, + annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, + expect: Annotated[str, "Digest a caller already holds: the rows are omitted when they still hash to it"] = "", +) -> dict: + """Walk item heads from ``addr`` as a flat listing: every head is rendered + (code OR data OR undefined) via generate_disasm_line and stepped with + next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data + byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's + real disassembly view. Address-paged: page forward by re-calling with + ``addr`` = the returned cursor.next; page up with ``back=true``.""" + import ida_bytes + import ida_segment + import idaapi + + count = 2000 if count > 2000 else (1 if count < 1 else count) + offset = max(int(offset), 0) + try: + start = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} + seg = ida_segment.getseg(start) + if not seg: + return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} + lo, hi = seg.start_ea, seg.end_ea + if end: + try: + hi = min(hi, parse_address(end)) + except Exception: + pass + + rows = [] + if back: + # Collect up to (count+offset) heads strictly before `start`, then take + # the window closest to `start`, returned in forward order. + walk = [] + cur = ida_bytes.prev_head(start, lo) + while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: + walk.append(cur) + cur = ida_bytes.prev_head(cur, lo) + walk.reverse() + chosen = walk[: len(walk) - offset] if offset else walk + chosen = chosen[-count:] + rows = [_idatui_head_row(e) for e in chosen] + first = chosen[0] if chosen else start + pea = ida_bytes.prev_head(first, lo) + cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} + return {"addr": str(addr), "heads": rows, "cursor": cursor} + + # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a + # flat listing must show them (IDA renders undefined as `db ?` lines, and + # navigating to an unmarked address must land ON it). Defined items advance + # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its + # end found in O(1) via next_head, which skips undefined) so a large .bss or + # gap doesn't explode into millions of one-byte rows. + def _is_unknown_f(f): + return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) + + def _run_end(e): + """End (exclusive) of the undefined run starting at ``e``.""" + nh = ida_bytes.next_head(e, hi) + return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi + + def _advance(e, f): + if _is_unknown_f(f): + return _run_end(e) + nxt = ida_bytes.get_item_end(e) + return nxt if nxt > e else e + 1 + + # The function the walk is currently inside, reused while it stays inside. + # get_func is ~0.5us and the walk asks per head; a head is nearly always in + # the same function as the one before it. Only ever consulted when ``e`` + # falls in [start_ea, end_ea), so a tail chunk elsewhere cannot be + # misattributed -- checked against get_func over 437k heads of + # bash/ls_ttl/echo with zero disagreements. + fn_cache = [None] + + def _func_at(e): + cur = fn_cache[0] + if cur is not None and cur.start_ea <= e < cur.end_ea: + return cur + cur = idaapi.get_func(e) + fn_cache[0] = cur + return cur + + def _rows_for(e, f): + if _is_unknown_f(f): + return [_idatui_unknown_row(e, _run_end(e) - e)] + func = _func_at(e) if annotate else None + at_start = func is not None and func.start_ea == e + out = [] + if at_start: + out.extend(_idatui_func_header_rows(e)) + row = _idatui_head_row(e, f) + if at_start: + row = dict(row) + row["name"] = None # the name is shown on the proc header line + elif annotate and row.get("kind") == "code" and row.get("name"): + # A code label (loc_XXX/jump target) gets its OWN line at depth 0, + # like IDA; strip it from the instruction row below. + nm = row["name"] + out.append({"ea": hex(e), "kind": "label", "size": 0, + "text": nm + ":", "name": nm}) + row = dict(row) + row["name"] = None + out.append(row) + if row.get("kind") == "data": + out.extend(_idatui_struct_member_rows(e)) # expand struct fields + if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: + out.extend(_idatui_func_footer_rows(e, func)) + return out + + ea = ida_bytes.get_item_head(start) + get_flags = ida_bytes.get_flags + for _ in range(offset): + if ea >= hi or ea == idaapi.BADADDR: + break + ea = _advance(ea, get_flags(ea)) + more = False + while ea != idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True + break + f = get_flags(ea) # once per head, not once per consumer + rows.extend(_rows_for(ea, f)) # a struct head expands into member rows + ea = _advance(ea, f) + cursor = {"next": hex(ea)} if more else {"done": True} + dig = _idatui_rows_digest(rows) + out = {"addr": str(addr), "cursor": cursor, "digest": dig, "count": len(rows)} + # ``expect`` says "I already hold a page that hashed to this". The rows are + # built either way -- generate_disasm_line is the floor and there is no way + # to know a line is unchanged without rendering it -- but pickling several + # hundred rows with their colour spans, unpickling them and rebuilding Heads + # is about 40% of what a page costs, and after a rename almost every page + # comes back identical. + # + # It carries the expected value rather than being a yes/no "digest mode" so + # that a page which HAS changed still costs one round trip: asking first and + # fetching afterwards made every changed page two. + if not (expect and str(dig) == expect): + out["heads"] = rows + return out + + +_IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") + + +_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", + "float", "stack", "default") + + +def _idatui_fmt_nibbles(): + """{format name: IDA operand-type nibble}. Built on call, not at import: + this module is injected into a file that is imported before a database is + open.""" + import ida_bytes + return { + "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH, + "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR, + "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF, + "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO, + "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP, + "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK, + "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST, + } + + +def _idatui_fmt_name(nib): + for name, v in _idatui_fmt_nibbles().items(): + if v == nib: + return name + return "default" + + +def _idatui_op_fmt(ea, n): + """The format operand ``n`` of the item at ``ea`` is currently displayed in. + + Reads the nibble IDA keeps per operand rather than guessing from the text -- + ``1`` renders identically in hex and decimal, so the rendered line cannot + answer this.""" + import ida_bytes + F = ida_bytes.get_flags(ea) + nib = (F >> ida_bytes.get_operand_type_shift(int(n))) & 0xF + return _idatui_fmt_name(nib) + + +def _idatui_op_value(ea, n): + """(value, byte width) of operand ``n``, or (None, 0) if it hasn't got one. + + The value is what decides which formats are OFFERED: a character constant + for 0x38A9 or an offset to an unmapped address are stops worth skipping.""" + import ida_bytes + import ida_ua + + F = ida_bytes.get_flags(ea) + if ida_bytes.is_code(F): + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) <= 0: + return None, 0 + try: + op = insn.ops[int(n)] + except Exception: + return None, 0 + if op.type == ida_ua.o_void: + return None, 0 + v = op.value if op.type == ida_ua.o_imm else op.addr + try: + size = int(ida_ua.get_dtype_size(op.dtype)) + except Exception: + size = 0 + return int(v), size + size = int(ida_bytes.get_item_size(ea)) + read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word, + 4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size) + if read is None: + return None, size + try: + return int(read(ea)), size + except Exception: + return None, size + + +def _idatui_printable(v): + """Whether ``v`` would actually render as a character constant. IDA accepts + op_chr on anything and then prints the number anyway, so a cycle that offers + 'char' for 0x18 has a stop where nothing visibly happens.""" + if v is None or v < 0 or v > 0xFFFFFFFF: + return False + bs, x = [], int(v) + while True: + bs.append(x & 0xFF) + x >>= 8 + if not x: + break + return all(0x20 <= b <= 0x7E or b in (9, 10, 13) for b in bs) + + +def _idatui_offset_worth(v): + """Whether 'offset' is worth OFFERING as a cycle stop for value ``v``. + + Making an offset is not free: IDA invents a dummy name at the target + (``off_18``) and that name STAYS once you cycle past it. So the ring only + stops there when the target is already something you could name -- a symbol, + a function, or an item something else references. In a PIE at base 0 half + the small constants in a function are 'mapped' (they land in the ELF + header); ``sub rsp, 18h`` is not a reference and must not offer to become + one on the way past. + + An explicit request still converts anything mapped: that's a decision, not a + keypress that happened to land here. After it, the target HAS a name, so the + ring includes the stop from then on.""" + import ida_bytes + import ida_name + + return bool(v and ida_bytes.is_mapped(v) and ida_name.get_ea_name(v)) + + +def _idatui_op_candidates(ea): + """Operand indices at ``ea`` whose display format is worth changing. + + Immediates and displacements -- the literals. Deliberately NOT: + + * branch targets (o_near/o_far), or every jump on the listing would offer to + become a bare number, on a view you navigate by label; + * memory references (o_mem), e.g. x86-64's RIP-relative ``lea rdi, name``. + IDA prints those from the reference, not from the operand's number format, + so setting one is accepted and changes nothing on screen -- a keypress + that appears to do nothing is worse than one that says it can't. + + An explicit ``n`` still reaches them; this is what a bare cursor picks.""" + import ida_bytes + import ida_ua + + F = ida_bytes.get_flags(ea) + if ida_bytes.is_data(F): + return [0] # a data item's value is operand 0 + if not ida_bytes.is_code(F): + return [] # undefined bytes: IDA refuses a format outright + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) <= 0: + return [] + want = (ida_ua.o_imm, ida_ua.o_displ) + out = [] + for i in range(len(insn.ops)): + op = insn.ops[i] + if op.type == ida_ua.o_void: + break + if op.type in want: + out.append(i) + return out + + +def _idatui_op_spans(ea, text): + """[(start, end, n)] -- where each operand sits inside ``text`` (the + whitespace-collapsed line the TUI shows), so a cursor column can name the + operand it is standing on. + + Read out of IDA's own COLOR_OPND markers on the line, which is both free + (the line is generated anyway) and exact. print_operand is kept as a + fallback for a processor module that emits no operand markers -- it agrees + with the tags where both exist, but it re-renders every operand to say so. + """ + import ida_lines + import ida_ua + + line = ida_lines.generate_disasm_line(ea, 0) + if line: + _spans, ops = _idatui_spans(line) + if ops: + return [tuple(o) for o in ops] + + out, pos = [], 0 + for n in range(8): + try: + raw = ida_ua.print_operand(ea, n) + except Exception: + raw = None + if not raw: + continue + op = " ".join(ida_lines.tag_remove(raw).split()) + if not op: + continue + i = text.find(op, pos) + if i < 0: # duplicated operand text (mov eax, eax) + i = text.find(op) + if i < 0: + continue + out.append((i, i + len(op), n)) + pos = i + len(op) + return out + + +def _idatui_line_text(ea): + import ida_lines + line = ida_lines.generate_disasm_line(ea, 0) + return " ".join(ida_lines.tag_remove(line).split()) if line else "" + + +def _idatui_op_text(ea, text, n): + """How operand ``n`` reads on the line, for a message that names it.""" + for lo, hi, i in _idatui_op_spans(ea, text): + if i == int(n): + return text[lo:hi].strip() + return "" + + +def _idatui_apply_fmt(ea, n, fmt): + """Set operand ``n``'s display format. Returns (ok, error).""" + import ida_bytes + import ida_offset + import idaapi + + n = int(n) + if fmt == "default": + return bool(ida_bytes.clr_op_type(ea, n)), "" + if fmt == "offset": + base = ida_offset.calc_offset_base(ea, n) + if base in (idaapi.BADADDR, None) or base < 0: + base = 0 + return bool(ida_offset.op_plain_offset(ea, n, base)), "" + fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec, + "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin, + "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg, + "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt) + if fn is None: + return False, (f"can't set {fmt!r} from a name alone" + if fmt in _idatui_fmt_nibbles() else + f"unknown format {fmt!r}") + return bool(fn(ea, n)), "" + + +def op_format( + addr: Annotated[str, "Address of the instruction or data item"], + mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle", + col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1, + n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1, +) -> dict: + """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal, + binary, character, or an offset to the address it names. + + The value in the bytes never changes -- only the representation IDA renders + and remembers. ``cycle``/``back`` step the stops that make sense for THIS + operand: 'char' is skipped unless the value prints as one, 'offset' unless + the target is already named, so no press is ever a no-op you have to press + again. ``show`` reports without changing anything. + + A format the ring can't hold (a stack variable, an enum) is reported in + ``warn`` on the way out, with what to do about it -- ``mode`` takes any of + the names above outright, which is also how you put one back. + + Which operand: ``n`` if given, else the one under ``col`` (a column in the + whitespace-collapsed line, as ``heads`` renders it), else the first literal + on the line.""" + import ida_bytes + + try: + ea = ida_bytes.get_item_head(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e)} + + before = _idatui_line_text(ea) + cands = _idatui_op_candidates(ea) + n = int(n) + if n < 0: + n = -1 + if int(col) >= 0: + for lo, hi, i in _idatui_op_spans(ea, before): + if not (lo <= int(col) < hi): + continue + if i in cands: + n = i + break + # The cursor IS on an operand, just not one with a format. The + # client highlights what the cursor is on, so quietly moving to + # a different operand would make that highlight a lie -- say + # which one can be changed instead. + where = before[lo:hi].strip() + alt = (f"; the literal on this line is operand {cands[0]} " + f"({_idatui_op_text(ea, before, cands[0])})" + if cands else "") + return {"addr": hex(ea), "n": i, "text": before, + "error": f"operand {i} ({where}) has no format to " + f"change{alt}"} + if n < 0: + if not cands: + F = ida_bytes.get_flags(ea) + why = ("no literal on this line to reformat" + if ida_bytes.is_code(F) or ida_bytes.is_data(F) else + "undefined bytes have no format to change -- define " + "them first ('d' makes data, 'c' makes code)") + return {"addr": hex(ea), "text": before, "error": why} + n = cands[0] + + cur = _idatui_op_fmt(ea, n) + value, width = _idatui_op_value(ea, n) + mapped = value is not None and value != 0 and ida_bytes.is_mapped(value) + # The ring is a property of the OPERAND, not of what you last pressed: every + # stop is one that changes what you see for this value, and it is the same + # ring at every step, so a lap always comes home. + choices = [f for f in _IDATUI_FMT_CYCLE + if (f != "char" or _idatui_printable(value)) + and (f != "offset" or _idatui_offset_worth(value))] + # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame + # member, not a way of writing a number, and IDA's own "is this a stack + # variable" test isn't exposed to Python here (calc_stkvar_struc_offset + # happily answers for ``[r14+8]`` too, which would put a bogus stop in the + # ring). Leaving one is reported instead, with the command that undoes it. + lossy = cur not in choices and cur != "default" + + mode = str(mode or "cycle").lower() + if mode == "show": + return {"addr": hex(ea), "n": n, "format": cur, "prev": cur, + "choices": choices, "text": before, "before": before, + "value": None if value is None else hex(value), + "width": width, "applied": False} + if mode in ("cycle", "back"): + step = 1 if mode == "cycle" else -1 + if cur in choices: + want = choices[(choices.index(cur) + step) % len(choices)] + else: + # Standing on a format the ring can't hold (an enum names a type a + # nibble doesn't record): enter the ring at its end, don't skip a + # stop working out where we "would have" been. + want = choices[0] if step > 0 else choices[-1] + else: + want = mode + if want not in _idatui_fmt_nibbles(): + return {"addr": hex(ea), "n": n, "text": before, + "error": f"unknown format {mode!r}; one of " + + ", ".join(_IDATUI_FMT_SETTABLE)} + if want == "offset" and not mapped: + return {"addr": hex(ea), "n": n, "text": before, "format": cur, + "error": (f"{'0x%x' % value if value is not None else 'this operand'}" + " isn't a mapped address -- an offset to it would" + " invent a name for nothing")} + + ok, err = _idatui_apply_fmt(ea, n, want) + if err: + return {"addr": hex(ea), "n": n, "text": before, "format": cur, + "error": err} + got = _idatui_op_fmt(ea, n) + out = {"addr": hex(ea), "n": n, "prev": cur, "format": got, + "requested": want, "applied": bool(ok), "choices": choices, + "before": before, "text": _idatui_line_text(ea), + "value": None if value is None else hex(value), "width": width} + if not ok: + out["error"] = f"IDA refused {want} on operand {n}" + elif lossy: + out["warn"] = ( + f"operand {n} was {cur} and the ring has no stop there -- " + + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else + f"{cur} names a type this can't put back, reassign it by hand")) + return out + + +_IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default") + + +def _idatui_compact(line): + """The ida-pro-mcp whitespace collapse the pseudocode is served through, so + a column in what the client SHOWS can be mapped back to Hex-Rays' line. + + DEVIATION FROM THE EXTRACTED ORIGINAL, deliberately: this used to be + ``from ida_pro_mcp.ida_mcp.utils import compact_whitespace`` inside a + try/except, with a plain ``[ \\t]{2,}`` regex as the fallback. Under Code + Mode ida_pro_mcp is not installed in the database process, so BOTH halves + of that were wrong: + + * the import failed on every call, and a failed import is never cached, so + each one re-searched the whole of sys.path -- 422 failures per pc_nums + call, which was the majority of its runtime; + * the fallback collapses runs of spaces INSIDE STRING LITERALS, which the + real function preserves. Pseudocode columns are served in these + coordinates, so a line containing a string with two spaces would have put + every literal's mark, and every reformat, on the wrong column. + + The module-level shim above is byte-identical to the original regex, so + call it directly. + """ + return compact_whitespace(line) + + +def _idatui_compact_col(plain, compact, col): + """The inverse of ``_idatui_uncompact_col``: a column in Hex-Rays' own line, + expressed in the collapsed line the client shows.""" + j = 0 + for i in range(min(int(col), len(plain))): + if j < len(compact) and plain[i] == compact[j]: + j += 1 + return j + + +def _idatui_uncompact_col(plain, compact, col): + """Map a column in the collapsed line back to the same character in the + original. The transform only ever DELETES spaces, so walking both in step + and skipping what vanished is exact.""" + i = 0 + for j in range(min(int(col), len(compact))): + c = compact[j] + while i < len(plain) and plain[i] != c: + i += 1 + i += 1 + return min(i, max(len(plain) - 1, 0)) + + +_IDATUI_LIT_CHARS = frozenset("0123456789abcdefABCDEFxXuUlL") + + +def _idatui_lit_extent(plain, x): + """The [start, end) of the literal token containing column ``x``. + + Hex-Rays says WHICH item a column belongs to, but not how wide the printed + literal is -- and it attributes neighbouring punctuation to the same item, + so ``if ( a1 > 1 )`` reports the closing paren as part of the number. The + identity comes from the ctree; the extent is the run of literal characters + around the column, which cannot reach a ``)`` or a space.""" + if x >= len(plain): + return None + if plain[x] == "'": # a character constant: '-' + end = plain.find("'", x + 1) + return (x, end + 1) if end > x else None + lo = plain.rfind("'", 0, x) + if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \ + plain[lo:x].count("'") == 1 and " " not in plain[lo:x]: + return (lo, plain.find("'", x) + 1) # inside 'c' + if plain[x] not in _IDATUI_LIT_CHARS: + return None + lo = x + while lo > 0 and plain[lo - 1] in _IDATUI_LIT_CHARS: + lo -= 1 + hi = x + while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS: + hi += 1 + if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it + lo -= 1 + return (lo, hi) + + +def _idatui_pc_nums(cf, sl): + """Every number literal on one pseudocode line, as + [{x0, x1, ea, opnum, value, nbytes, fmt}]. + + Asks Hex-Rays what each column belongs to rather than pattern-matching the + text: a regex over ``v6 = a1 - 1;`` has to guess which of those characters + are a literal, and ``v11`` looks like one.""" + import ida_bytes + import ida_hexrays + import ida_lines + import idaapi + + plain = ida_lines.tag_remove(sl.line) + out = [] + x = 0 + while x < len(plain): + ch = plain[x] + if ch not in _IDATUI_LIT_CHARS and ch != "'": + x += 1 + continue + head, item, tail = (ida_hexrays.ctree_item_t() for _ in range(3)) + if not cf.get_line_item(sl.line, x, True, head, item, tail): + x += 1 + continue + if item.citype != ida_hexrays.VDI_EXPR: + x += 1 + continue + e = item.e + if e.op != ida_hexrays.cot_num: + x += 1 + continue + extent = _idatui_lit_extent(plain, x) + if extent is None: + x += 1 + continue + nf = e.n.nf + opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum) + nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) + else int(nf.org_nbytes)) + ea = int(e.ea) + if ea == idaapi.BADADDR: + x = extent[1] + continue # synthesised: nothing to key on + nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF + # Whether this format is the USER's or Hex-Rays' own guess. The nibble + # can't say: an untouched number reads back as whatever it happens to + # be printed as, and cycling from there would skip that stop forever + # (default already looks like it) and never come back to it. + loc = ida_hexrays.operand_locator_t(ea, opnum) + user = (ida_hexrays.user_numforms_find(cf.numforms, loc) + != ida_hexrays.user_numforms_end(cf.numforms)) + out.append({"x0": extent[0], "x1": extent[1], "ea": ea, + "opnum": opnum, "value": int(e.n._value), + "nbytes": nbytes, "user": user, + "fmt": _idatui_fmt_name(nib) if user else "default", + "shown": _idatui_fmt_name(nib)}) + x = extent[1] # past this literal, not into it + return out + + +def pc_nums( + addr: Annotated[str, "Function address (or any address inside it)"], +) -> dict: + """Every number literal in a function's pseudocode, as + [{line, x0, x1, ea, opnum, value, fmt, user}]. + + One call per decompilation, so a client can show WHICH literal the cursor is + on (and reformat exactly that one) without a round trip per cursor move. + Columns are in the same collapsed coordinates the decompile tool serves its + text in, i.e. what the client actually displays.""" + import ida_hexrays + import ida_lines + import idaapi + + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": str(addr), "error": "no decompiler", "nums": []} + try: + f = idaapi.get_func(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e), "nums": []} + if f is None: + return {"addr": str(addr), "error": "no function here", "nums": []} + try: + cf = ida_hexrays.decompile(f.start_ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", + "nums": []} + if cf is None: + return {"addr": hex(f.start_ea), "error": "decompilation failed", + "nums": []} + sv = cf.get_pseudocode() + out = [] + for i in range(len(sv)): + plain = ida_lines.tag_remove(sv[i].line) + compact = _idatui_compact(plain) + for rec in _idatui_pc_nums(cf, sv[i]): + out.append({ + "line": i, + "x0": _idatui_compact_col(plain, compact, rec["x0"]), + "x1": _idatui_compact_col(plain, compact, rec["x1"]), + "ea": hex(rec["ea"]), "opnum": rec["opnum"], + "value": hex(rec["value"]), "fmt": rec["fmt"], + "shown": rec["shown"], "user": bool(rec["user"]), + }) + return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)} + + +def pc_num_format( + addr: Annotated[str, "Function address (or any address inside it)"], + mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle", + line: Annotated[int, "0-based pseudocode line index"] = -1, + col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1, + ea: Annotated[str, "Address of the number instead of line/col"] = "", + opnum: Annotated[int, "Operand number, with ``ea``"] = -1, +) -> dict: + """Change how a number is displayed in the DECOMPILATION (Hex-Rays keeps its + own number formats, per (address, operand), independent of the listing). + + Same stops as ``op_format`` minus the two C can't express: binary (no such + literal -- IDA takes the format and prints decimal anyway) and offset (it + makes the function stop decompiling). Returns the re-rendered line, and + marks the function dirty so the next decompile is the new text.""" + import ida_hexrays + import ida_lines + import idaapi + + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": str(addr), "error": "no decompiler"} + try: + f = idaapi.get_func(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e)} + if f is None: + return {"addr": str(addr), "error": "no function here"} + try: + cf = ida_hexrays.decompile(f.start_ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} + if cf is None: + return {"addr": hex(f.start_ea), "error": "decompilation failed"} + + sv = cf.get_pseudocode() + line = int(line) + target = None + if ea: + try: + want_ea = parse_address(ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": str(e)} + for i in range(len(sv)): + for rec in _idatui_pc_nums(cf, sv[i]): + if rec["ea"] == want_ea and (int(opnum) < 0 + or rec["opnum"] == int(opnum)): + target, line = rec, i + break + if target: + break + elif 0 <= line < len(sv): + nums = _idatui_pc_nums(cf, sv[line]) + if nums: + if int(col) >= 0: + plain = ida_lines.tag_remove(sv[line].line) + x = _idatui_uncompact_col(plain, _idatui_compact(plain), int(col)) + target = next((r for r in nums if r["x0"] <= x < r["x1"]), None) + target = target or nums[0] + else: + return {"addr": hex(f.start_ea), + "error": f"line {line} is outside the {len(sv)}-line decompilation"} + if target is None: + return {"addr": hex(f.start_ea), "line": line, + "text": (ida_lines.tag_remove(sv[line].line).strip() + if 0 <= line < len(sv) else ""), + "error": "no number literal on this line"} + + cur, value = target["fmt"], target["value"] + choices = [c for c in _IDATUI_PC_FMT_CYCLE + if c != "char" or _idatui_printable(value)] + # Same rule as the listing: one ring per literal, every step. A format the + # ring can't hold (an enum set in the GUI) is reported on the way out + # instead of being kept for one lap and then lost. + lossy = cur not in choices and cur != "default" + out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]), + "opnum": target["opnum"], "line": line, "prev": cur, + "format": cur, "shown": target["shown"], "choices": choices, + "value": hex(value), + "before": ida_lines.tag_remove(sv[line].line).strip()} + + mode = str(mode or "cycle").lower() + if mode == "show": + out["text"] = out["before"] + out["applied"] = False + return out + if mode in ("cycle", "back"): + step = 1 if mode == "cycle" else -1 + if cur in choices: + want = choices[(choices.index(cur) + step) % len(choices)] + else: + want = choices[0] if step > 0 else choices[-1] + else: + want = mode + if want in ("bin", "offset", "stack", "seg", "float"): + out["error"] = (f"Hex-Rays has no {want} format for a number " + f"-- set it on the listing instead") + out["text"] = out["before"] + return out + if want not in ("hex", "dec", "oct", "char", "default"): + out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, " + f"char, default") + out["text"] = out["before"] + return out + + loc = ida_hexrays.operand_locator_t(target["ea"], target["opnum"]) + it = ida_hexrays.user_numforms_find(cf.numforms, loc) + if it != ida_hexrays.user_numforms_end(cf.numforms): + # std::map::insert is a no-op on an existing key, so a format already + # set here would silently win over the new one. + ida_hexrays.user_numforms_erase(cf.numforms, it) + if want != "default": + import ida_bytes + nf = ida_hexrays.number_format_t(target["opnum"]) + nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], + target["opnum"]) + try: + nf.org_nbytes = target["nbytes"] + except Exception: + pass + ida_hexrays.user_numforms_insert(cf.numforms, loc, nf) + cf.save_user_numforms() + try: + ida_hexrays.mark_cfunc_dirty(f.start_ea) + except Exception: + pass + + out["format"] = want + out["applied"] = True + if lossy: + out["warn"] = (f"this number was {cur}, which names a type a radix " + f"can't put back -- reassign it in IDA") + try: + cf2 = ida_hexrays.decompile(f.start_ea, + flags=ida_hexrays.DECOMP_NO_CACHE) + sv2 = cf2.get_pseudocode() if cf2 is not None else None + out["text"] = (ida_lines.tag_remove(sv2[line].line).strip() + if sv2 is not None and line < len(sv2) else out["before"]) + except Exception as e: + out["text"] = out["before"] + out["warn"] = f"re-render failed: {e}" + return out + + +def decompile(addr, include_addresses=True): + """Pseudocode for the function at ``addr``, plus the objects it references. + + Faithful to the tool ida-tui was written against, and in particular to its + COST: the per-line address anchor comes from ONE ``get_line_item`` at column + 0 per line. The Code Mode port asked for the full per-column line map (what + ``decomp_map`` is for) purely to fill in that anchor, which is thousands of + ``get_line_item``+``dstr()`` calls per function instead of one per line, and + made every pseudocode open cost the same as opening the split view. + + Text is whitespace-collapsed exactly as the client displays it, because + ``pc_nums`` reports literal columns in those coordinates. + """ + import ida_bytes + import ida_hexrays + import ida_lines + import ida_name + import idaapi + + try: + ea = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "code": None, "error": str(e)} + fn = idaapi.get_func(ea) + if fn is None: + return {"addr": str(addr), "code": None, "error": f"no function at {ea:#x}"} + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": hex(int(fn.start_ea)), "code": None, "error": "no decompiler"} + failure = ida_hexrays.hexrays_failure_t() + try: + cfunc = ida_hexrays.decompile_func(fn, failure) + except Exception as e: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": f"Decompilation failed at {ea:#x}: {e}"} + if cfunc is None: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": failure.desc() or f"Decompilation failed at {ea:#x}"} + + lines = [] + for sl in cfunc.get_pseudocode(): + head = ida_hexrays.ctree_item_t() + item = ida_hexrays.ctree_item_t() + tail = ida_hexrays.ctree_item_t() + line_ea = None + if include_addresses and cfunc.get_line_item(sl.line, 0, False, head, item, tail): + parts = (item.dstr() or "").split(": ") + if len(parts) == 2: + try: + line_ea = int(parts[0], 16) + except ValueError: + line_ea = None + text = compact_whitespace(ida_lines.tag_remove(sl.line)) + lines.append(f"{text} /*{line_ea:#x}*/" if line_ea is not None else text) + + refs, seen = [], set() + + class _RefVisitor(ida_hexrays.ctree_visitor_t): + def __init__(self): + ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST) + + def visit_expr(self, e): + if e.op == ida_hexrays.cot_obj: + target = int(e.obj_ea) + if target != idaapi.BADADDR and target not in seen: + seen.add(target) + try: + raw = ida_bytes.get_strlit_contents(target, -1, 0) + text = raw.decode("utf-8", "replace") if raw else None + except Exception: + text = None + refs.append({"addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text}) + return 0 + + try: + _RefVisitor().apply_to(cfunc.body, None) + except Exception: + pass + return {"addr": hex(int(fn.start_ea)), "code": "\n".join(lines), "refs": refs} + + +def decomp_map( + addr: Annotated[str, "Function address or name"], +) -> dict: + """Per-pseudocode-line instruction coverage for the split view's region + highlight: for each line, the set of EAs the decompiler attributes to it, + swept across the line's columns via get_line_item. Shape: + {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" + import ida_hexrays + import idaapi + try: + ea = int(str(addr), 16) + except ValueError: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + func = idaapi.get_func(ea) + if not func: + return {"error": f"no function at {addr}"} + try: + cfunc = ida_hexrays.decompile(func.start_ea) + except Exception as e: # noqa: BLE001 + return {"error": f"decompile failed: {e}"} + if cfunc is None: + return {"error": "decompile failed"} + import ida_lines + # Three things this loop must not do, each measured on real functions (the 25 + # largest of bash went 68.3s -> 6.5s; echo's 60 largest 5.4s -> 0.6s, with + # byte-identical output): + # + # * allocate ctree_item_t's per COLUMN. They are SWIG objects and this is + # the innermost loop; one per call is enough, and head/tail are never + # read, so don't ask for them at all. + # * sweep the TAGGED length. ``x`` is a screen column but ``sl.line`` still + # carries IDA's colour tags, so a 23-column line was swept 124 times. + # * call dstr() per column. It formats a whole 'EA: description' string -- + # 24us a call, which is 79% of this tool. Comparing against the PREVIOUS + # column's item id is not enough: items interleave, so `foo(a, b)` flips + # call -> arg -> call -> arg and every flip re-formats an item already + # seen (106 594 calls for 15 417 lines of bash). Memoise id -> ea for the + # whole function instead: obj_id is unique within a cfunc, so the same id + # always yields the same string, and the result is deduped by ``seen`` + # anyway. Items with no ctree node (it is None) have no id to key on and + # still pay per occurrence. + item = ida_hexrays.ctree_item_t() + tag_remove = ida_lines.tag_remove + get_line_item = cfunc.get_line_item + ea_of_id = {} + lines = [] + for sl in cfunc.get_pseudocode(): + line = sl.line + eas, seen = [], set() + prev_id = None + for x in range(len(tag_remove(line)) + 1): + if not get_line_item(line, x, False, None, item, None): + continue + it = item.it + if it is not None: + oid = it.obj_id + if oid == prev_id: + continue + prev_id = oid + if oid in ea_of_id: + e = ea_of_id[oid] + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + continue + else: + oid = None + prev_id = None + # Match the /*ea*/ marker's source (decompile_function_safe): the + # item's dstr() is 'EA: description'; get_ea() reports a different ea. + e = None + dstr = item.dstr() + if dstr: + parts = dstr.split(": ", 1) + if len(parts) == 2: + try: + e = int(parts[0], 16) + except ValueError: + e = None + if oid is not None: + ea_of_id[oid] = e + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + lines.append({"ea": eas[0] if eas else None, "eas": eas}) + return {"addr": hex(func.start_ea), "lines": lines} diff --git a/idatui/rpc.py b/idatui/rpc.py index 0921b89..1262ebb 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -28,7 +28,8 @@ from typing import Any from rich.console import Console from ._sync import drain, settle -from .app import DecompView, HexView, ListingView +from . import diag +from .app import DecompView, GraphView, HexView, ListingView, ViewMode PROTO_VERSION = 1 TYPE_DELAY_MS = 35 # default per-char delay for high-level typed ops (aesthetic) @@ -38,6 +39,7 @@ _PROGRAM_METHODS = { "goto", "open", "rename", "comment", "retype", "follow", "xrefs", "symbols", "structs", "search", "select", "save", "hex", "toggle_view", "pseudocode", "disassembly", "xrefs_to", "xrefs_from", "resolve", + "define", "rename_many", "opfmt", "graph", "export", "find", } # Self-documenting method table (returned by the 'methods' verb). @@ -59,24 +61,62 @@ METHODS = { "text": "{text,delay_ms?,settle?} type a literal string into the focused input", "goto/open": "{target,delay_ms?} g-prompt to a name or 0xADDR", "rename": "{name,word?,delay_ms?} rename the token under (or 'word') the cursor", - "comment": "{text,delay_ms?} comment the current line", + "comment": "{text} comment the current line (use \\n for newlines)", "retype": "{proto,word?,delay_ms?} set the prototype/type under the cursor", "follow": "{word?} follow the reference under (or 'word') the cursor", "cursor_on": "{word,line?,occurrence?=1} place the cursor on a token", "back": "pop the nav stack", "toggle_view": "disasm <-> pseudocode", "hex": "hex view", + "graph": "{action?=show|open|close|toggle|zoom|block|entry|succ|pred," + "target?,blocks?} the control-flow graph: 'show' reports its " + "structure (blocks, edges, cursor) without touching it; the others " + "drive it. 'block' takes target=<id|0xADDR>", "xrefs": "open the xref picker", "symbols": "{query?} open the symbol palette", "structs": "open the struct editor", + "export": "{path?,types?=true} write the session's comments/names/types as " + "a markdown report -> {path,comments,names,types}", + "find": "{query,mode?=auto|text|bytes,limit?=500,regex?,case?} search the " + "WHOLE database: disassembly text, or a byte pattern with " + "wildcards (48 8b ?? c3) -> {mode,hits:[{addr,head,line,func}]}", "search": "{term,direction?=1} incremental search in the code view", "select": "{index?} choose the highlighted/nth item in the open modal", "save": "persist the .i64 (Ctrl+S)", + "trace": "{seek|goto|step,over?} navigate the execution trace (seek '!50' = percent)", "binaries": "-> project binaries {label,active,resident,indexed} (project mode)", "switch": "{binary,addr?} make another project binary active (addr also jumps)", "close": "dismiss a modal (Escape)", "move": "{dir,n?=1} fast movement (down/up/.../pagedown)", "cursor": "{line?,col?} set the code-pane cursor directly", + "define": "{kind:code|func|undef|thumb|thumbscan|data|string,target?} " + "(re)define bytes at target — the raw-image workflow", + "rename_many": "{items:[{addr,name}] | file:JSON} bulk-apply a symbol file " + "in ONE call (no typing, no navigation)", + "opfmt": "{mode?=cycle|back|show|hex|dec|oct|bin|char|offset|stack|" + "default,target?,word?,line?,col?} how the literal under the cursor is " + "DISPLAYED (IDA's 'o'); works on the listing and on pseudocode " + "numbers. 'show' reports the format and the stops without editing", +} + +#: `opfmt` modes that have a real key on the code views. Driving the key keeps +#: the pane honest (a viewer sees the same thing a human would do); the named +#: formats have no key, so those go through the view's action directly. +_OPFMT_KEYS = {"cycle": "o", "back": "O"} +_OPFMT_MODES = ("cycle", "back", "show", "hex", "dec", "oct", "bin", "char", + "offset", "stack", "default") + +# `define` kinds -> the ListingView key that runs them. Driving the real key +# keeps the pane honest (a viewer sees the same thing a human would do) and +# reuses the app's own edit worker, which reports what actually happened. +_DEFINE_KEYS = { + "code": "c", # make code (runs until flow/undecodable) + "func": "p", # make function + "undef": "u", + "thumb": "t", # flip ARM/Thumb at the cursor, then disassemble + "thumbscan": "T", # find Thumb entry pointers in a vector table + "data": "d", + "string": "a", } # Movement keys — driven fast (no typed delay) so the pane still visibly moves. @@ -94,15 +134,81 @@ _MOVE_KEYS = { # --------------------------------------------------------------------------- # def _active_widget(app): """The currently *shown* code widget (mirrors app._active).""" - if app._active == "hex": + if app.is_hex: return app.query_one(HexView) - if app._active in ("listing", "disasm"): + if app.is_graph: + return app.query_one(GraphView) + if app.is_listing: return app.query_one(ListingView) return app.query_one(DecompView) +def graph_info(app, blocks: bool = True) -> dict[str, Any]: + """Structured view of the control-flow graph: what a driver actually wants, + rather than the box-drawing characters it is rendered as.""" + gv = app.query_one(GraphView) + if gv.fc is None or gv.lay is None: + return {"open": app.is_graph, "loaded": False, + "note": "press space (or graph {action:'open'}) on a function"} + lay, fc = gv.lay, gv.fc + out: dict[str, Any] = { + "open": app.is_graph, + "loaded": True, + "func": {"name": fc.name, "ea": fc.func_ea, "entry": fc.entry}, + "zoom": gv.ZOOMS[gv._zoom], + "canvas": {"w": lay.width, "h": lay.height}, + "stats": dict(lay.stats), + "cursor": {"block": gv.cursor_node, "row": gv.cursor_row, + "ea": gv._cursor_ea(), "word": gv.word_under_cursor()}, + } + if blocks: + rows = [] + for n in lay.nodes: + b = gv._blocks.get(n.id) + rows.append({ + "id": n.id, + "start": b.start if b else None, + "end": b.end if b else None, + "insns": len(b.rows) if b else 0, + "rank": n.rank, + "box": {"x": n.x, "y": n.y, "w": n.w, "h": n.h}, + "succs": [{"id": i, "kind": k} for i, k in lay.succ.get(n.id, [])], + "preds": [{"id": i, "kind": k} for i, k in lay.pred.get(n.id, [])], + "selfloop": bool(b and any(d == n.id for d, _ in b.succs)), + }) + out["blocks"] = rows + return out + + _MODALS = ("XrefsScreen", "SymbolPalette", "StructEditor", "ConfirmScreen") +#: ``drive raw`` (and any k=v CLI) hands every param through as a *string*. +#: Handlers that did ``int(...)`` coped; the ones that compared directly blew up +#: with e.g. "'<' not supported between instances of 'int' and 'str'". Coerce the +#: known-numeric names once, centrally, instead of at every call site. +_INT_PARAMS = ("lines", "limit", "max", "n", "index", "line", "col", + "occurrence", "delay_ms", "direction", "addr", "count") +_FLOAT_PARAMS = ("timeout",) + + +def _coerce_params(params: dict[str, Any]) -> dict[str, Any]: + out = dict(params) + for k in _INT_PARAMS: + v = out.get(k) + if isinstance(v, str) and v.strip(): + try: + out[k] = int(v, 0) + except ValueError: + pass + for k in _FLOAT_PARAMS: + v = out.get(k) + if isinstance(v, str) and v.strip(): + try: + out[k] = float(v) + except ValueError: + pass + return out + def _modal_snapshot(app) -> dict[str, Any] | None: """Describe the top modal screen, if any, enough to drive it.""" @@ -131,6 +237,12 @@ def _cursor_info(app, w) -> dict[str, Any]: if isinstance(w, HexView): return {"kind": "hex", "va": (w.cursor_va() if w.model else None), "byte": w.cursor} + if isinstance(w, GraphView): + # The graph cursor is (block, row), not a line index -- reporting it as + # one would make a driver's `cursor line=` land somewhere arbitrary. + return {"kind": "graph", "ea": w._cursor_ea(), "block": w.cursor_node, + "row": w.cursor_row, "col": w.cursor_x, + "word": w.word_under_cursor(), "text": w._line_plain()} # disasm / decomp share the ColumnCursor surface word = None try: @@ -147,6 +259,14 @@ def _cursor_info(app, w) -> dict[str, Any]: "scroll_y": round(w.scroll_offset.y)} +def _where(app) -> str: + """Short 'name @ 0xea' for error messages that need to say where we ended up.""" + cur = getattr(app, "_cur", None) + if cur is None: + return "nowhere" + return f"{getattr(cur, 'name', '?')} @ {getattr(cur, 'ea', 0):#x}" + + def _readiness(app) -> dict[str, Any]: """Whether the app is drivable yet, and how far function-loading has got. (Cheap: no network — never call client.health() here.)""" @@ -191,6 +311,11 @@ def view_lines(app, lines: int | None = None) -> dict[str, Any]: if isinstance(w, HexView): return {"active": "hex", "note": "use screen() for the hex grid", "cursor": _cursor_info(app, w)} + if isinstance(w, GraphView): + return {"active": "graph", "note": "use graph() for structure, " + "screen() for the drawing", + "cursor": _cursor_info(app, w), + "graph": graph_info(app, blocks=False)} top = round(w.scroll_offset.y) height = w.size.height or 40 n = min(lines or height, max(w.total - top, 0)) @@ -224,21 +349,58 @@ def screen_text(app, fmt: str = "text") -> dict[str, Any]: return out +def place_cursor(w, line=None, col=None) -> None: + """Move a code view's cursor and BRING IT INTO VIEW. + + Setting the cursor without scrolling leaves the pane showing somewhere else + entirely, and the next verb then edits a line the operator cannot see — the + status describes one thing, the screen shows another. Every programmatic + cursor move goes through here for that reason. + """ + if line is not None: + w.cursor = max(0, min(getattr(w, "total", 1) - 1, int(line))) + if col is not None: + w.cursor_x = max(0, int(col)) + if hasattr(w, "_after_cursor_move"): + w._after_cursor_move() + if hasattr(w, "_scroll_cursor_into_view"): + w._scroll_cursor_into_view() + if hasattr(w, "_hscroll"): + w._hscroll() + w.refresh() + + def cursor_on(app, word: str, line: int | None = None, occurrence: int = 1) -> bool: """Place the cursor on the ``occurrence``-th token equal to ``word`` in the active code pane (optionally restricted to ``line``). Verified with the app's own tokenizer so 'main' won't match inside 'domain'. Disasm scan is limited to already-cached lines (what's on/near screen); decomp searches the whole body. - Returns whether it found and moved.""" + Returns whether it found and moved. + + Search starts at the VIEWPORT, not at row 0. A continuous listing is the + whole segment, so counting from the top finds an occurrence in some unrelated + function thousands of rows away -- and the cursor then lands there, off + screen, where the next verb edits something the operator cannot see. Wrapping + to the rows above keeps every match reachable; landing scrolls, so wherever + it goes is visible. + """ w = _active_widget(app) if isinstance(w, HexView): raise ValueError("cursor_on: not supported in the hex view") + if isinstance(w, GraphView): + raise ValueError("cursor_on: not supported in the graph view — use " + "graph {action:'block'} or goto") if isinstance(w, DecompView): texts = list(w._texts) else: texts = [(w._line_plain(i) or "") for i in range(getattr(w, "total", 0))] orig = (w.cursor, w.cursor_x) - rows = [line] if line is not None else range(len(texts)) + if line is not None: + rows = [line] + else: + # From the top of the viewport, then wrap round to what's above it. + top = round(w.scroll_offset.y) + rows = list(range(top, len(texts))) + list(range(0, top)) hits = 0 for i in rows: if not (0 <= i < len(texts)): @@ -250,9 +412,7 @@ def cursor_on(app, word: str, line: int | None = None, occurrence: int = 1) -> b if w.word_under_cursor() == word: hits += 1 if hits >= max(1, occurrence): - if hasattr(w, "_after_cursor_move"): - w._after_cursor_move() - w.refresh() + place_cursor(w) # scrolls: an off-screen cursor edits blind return True col = t.find(word, col + 1) w.cursor, w.cursor_x = orig # not found: leave the cursor untouched @@ -451,11 +611,85 @@ class RpcServer: return {"id": rid, "error": {"message": f"{type(e).__name__}: {msg}"}} # -- composed helpers (semantic verbs) -------------------------------- # - async def _press(self, keys, pred=None, timeout=20.0): + async def _press(self, keys, pred=None, timeout=20.0, what=""): await self.app._press_keys([str(k) for k in keys]) - await settle(self.app, pred, timeout=timeout) + ok = await settle(self.app, pred, timeout=timeout) + if pred is not None and not ok: + # Never report success for an action that did not happen: the caller + # would go on to edit whatever the *previous* location was. + raise TimeoutError( + f"{what or 'action'} did not complete within {timeout}s " + f"(still at {_where(self.app)}); retry with a larger timeout=") return snapshot(self.app) + async def _graph(self, params, timeout): + """Drive / read the control-flow graph. + + Everything goes through the real keys and the real view state, so a + driver sees exactly what a person would -- and 'show' is a pure read, + which is what you want between edits. + """ + app = self.app + action = str(params.get("action") or "show").lower() + gv = app.query_one(GraphView) + want_blocks = params.get("blocks", True) not in (False, "false", "0", 0) + + if action == "show": + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + if action in ("open", "toggle", "close"): + if action == "open" and app.is_graph: + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + if action == "close" and not app.is_graph: + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + want = "graph" if action in ("open", "toggle") and \ + not app.is_graph else None + res = await self._press( + ["space"], + (lambda: app.is_graph) if want else + (lambda: not app.is_graph), + timeout, f"graph {action}") + return {**res, "graph": graph_info(app, blocks=want_blocks)} + + if not app.is_graph: + raise ValueError(f"graph {action}: the graph is not open " + f"(graph {{action:'open'}} first)") + if action == "zoom": + before = gv._zoom + await self._press(["z"], lambda: gv._zoom != before, timeout, "graph zoom") + elif action == "entry": + await self._press(["0"], None, timeout, "graph entry") + elif action in ("succ", "pred"): + before = gv.cursor_node + await self._press(["J" if action == "succ" else "K"], + lambda: gv.cursor_node != before, timeout, + f"graph {action}") + elif action == "block": + target = params.get("target") + if target is None: + raise ValueError("graph block: need target=<block id|0xADDR>") + nid = None + s = str(target) + if s.startswith("0x") or s.startswith("0X"): + ea = int(s, 16) + b = gv.fc.block_at(ea) if gv.fc else None + if b is None: + raise ValueError(f"graph block: {s} is not in this graph") + nid = b.id + else: + nid = int(s) + if gv.lay is None or nid not in gv.lay.by_id: + raise ValueError(f"graph block: no block {nid}") + gv.cursor_node = nid + gv.cursor_row = 0 + gv.cursor_x = 0 + gv._clamp_cursor() + gv._center_cursor() + gv.refresh() + await settle(app, None, timeout=2) + else: + raise ValueError(f"graph: unknown action {action!r}") + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + async def _fill_prompt(self, open_key, input_id, value, delay_ms, clear): """Open a prompt (a keystroke), optionally clear its prefill, type the value with the typed-out delay, submit. Returns after the prompt closes.""" @@ -465,12 +699,101 @@ class RpcServer: await settle(app, lambda: app.query_one(f"#{input_id}", Input).display, timeout=10) inp = app.query_one(f"#{input_id}", Input) if not inp.display: - raise RuntimeError(f"{input_id!r} prompt did not open (word under cursor?)") + # Say *why*. The old message always blamed the word under the cursor, + # which sent readers hunting for a cursor problem when the real cause + # was usually a modal eating the opening keystroke. + modal = type(app.screen).__name__ + why = (f"modal {modal!r} has focus and ate the {open_key!r} keystroke" + if modal in _MODALS or modal != "Screen" + else "no renameable token under the cursor") + raise RuntimeError(f"{input_id!r} prompt did not open: {why}") if clear: inp.value = "" await app._press_keys(_text_to_keys(value, delay_ms)) await app._press_keys(["enter"]) + async def _rename_many(self, params: dict[str, Any], timeout: float) -> dict: + """Apply a whole symbol file in one worker call. + + The per-symbol path (goto + typed rename prompt) is the right thing for + one name a human is watching, and hopeless for the case a firmware image + always brings: hundreds of names from a loader map, an emulator's + symbols.json, or another tool's export. Each of those renames costs a + navigation (which pulls a listing page and a decompile) plus two prompt + round-trips, so 400 symbols is tens of minutes of driving and the pane + just flickers. IDA's own rename tool already takes a *list*; this hands + it the whole list, then refreshes the caches and the function table once. + """ + app = self.app + items = params.get("items") + src = params.get("file") + if isinstance(items, str): # `drive raw` hands params through as text + items = json.loads(items) + if items is None: + if not src: + raise ValueError("rename_many needs items=[{addr,name}] or file=<json>") + with open(os.path.expanduser(str(src))) as f: + items = json.load(f) + if isinstance(items, dict): # {"0x4370": "name"} is a natural shape too + items = [{"addr": k, "name": v} for k, v in items.items()] + if not isinstance(items, list) or not items: + raise ValueError("rename_many: items must be a non-empty list") + + ops, skipped = [], 0 + for it in items: + if not isinstance(it, dict): + skipped += 1 + continue + # Accept the field names symbol files actually use. + addr = next((it[k] for k in ("addr", "start", "ea", "address") + if it.get(k) is not None), None) + name = it.get("name") or it.get("label") + if addr is None or not name: + skipped += 1 + continue + ea = int(str(addr), 0) if isinstance(addr, str) else int(addr) + ops.append({"addr": hex(ea), "name": str(name)}) + if not ops: + raise ValueError("rename_many: no usable {addr,name} entries") + + overwrite = params.get("allow_overwrite", True) + if isinstance(overwrite, str): + overwrite = overwrite.lower() not in ("0", "false", "no", "") + batch = {"func": ops, "allow_overwrite": bool(overwrite)} + # The worker is blocking and single-threaded; off the event loop it goes, + # or the TUI freezes for the length of the batch. + res = await asyncio.to_thread(app.program.client.invoke, "rename", batch=batch) + summary = res.get("summary", {}) if isinstance(res, dict) else {} + failed = [r for r in (res.get("func") or []) if isinstance(r, dict) + and r.get("error")] if isinstance(res, dict) else [] + + # Names live in the IDB, but every cache in front of it is now stale -- + # including Hex-Rays', which is per-function and does NOT notice that a + # *callee* was renamed. That cache is persisted in the .i64, so without + # this a batch import leaves pseudocode calling sub_98C0 forever while + # the listing (and every readback) says memset. + try: + await asyncio.to_thread(app.program.client.invoke, "force_recompile") + except Exception: # noqa: BLE001 -- older worker without the tool + pass + app.program.bump_names() + app.program.invalidate_functions() + app._func_index = None + app._load_functions() # re-streams the function table + await settle(app, timeout=timeout) + app._dirty = True + app._status(f"renamed {summary.get('ok', 0)} symbols" + + (f", {len(failed)} failed" if failed else "") + + " (Ctrl+S to save)") + snap = snapshot(app) + snap["rename_many"] = { + "requested": len(ops), "skipped": skipped, + "ok": summary.get("ok", 0), "failed": summary.get("failed", 0), + "errors": [{"addr": r.get("addr"), "error": r.get("error")} + for r in failed[:10]], + } + return snap + def _goto_target_pred(self, target): """A predicate that holds once a goto to ``target`` has landed.""" app = self.app @@ -478,14 +801,39 @@ class RpcServer: ea = app.program.resolve(target) except Exception: # noqa: BLE001 — unknown name; caller falls back to generic return None - if app._active == "hex": + if app.is_hex: return lambda: app.query_one(HexView).cursor_va() == ea fn = app.program.function_of(ea) want = fn.addr if fn else ea return lambda: app._cur is not None and app._cur.ea == want + #: Verbs that drive the *main* app by injecting keystrokes. If a modal is on + #: top it eats those keys, so they must refuse rather than silently no-op. + _NEEDS_NO_MODAL = { + "goto", "open", "rename", "comment", "retype", "follow", "back", + "toggle_view", "hex", "save", "search", "move", "cursor", "cursor_on", + "define", "opfmt", + } + #: Modals the driver is expected to interact with (they have their own verbs). + _DRIVABLE_MODALS = {"XrefsScreen", "SymbolPalette", "StructEditor", + "ProjectPalette", "QuitScreen"} + + def _modal_kind(self) -> str | None: + scr = self.app.screen + name = type(scr).__name__ + return name if name in _MODALS or name in self._DRIVABLE_MODALS else None + async def _dispatch(self, method: str, params: dict[str, Any]) -> Any: app = self.app + params = _coerce_params(params) + if method in self._NEEDS_NO_MODAL: + modal = self._modal_kind() + if modal is not None: + raise RuntimeError( + f"modal {modal!r} is on top and will swallow this verb's " + f"keystrokes; dismiss it first (close) or use its own verb " + f"(select/symbols/xrefs). Note: a binary with no entry " + f"function can land in the symbol palette on startup.") if method in (None, "ping"): module = None try: @@ -497,9 +845,25 @@ class RpcServer: if method == "methods": return METHODS if method == "quit": + # Route through the same teardown a human gets, so a dirty database + # is written instead of dropped. `app.exit()` alone skips the dirty + # check entirely, and the caller (pane stop) then kills the pane -- + # which used to destroy a whole session's annotations. + dirty = list(app._dirty_labels()) + save = params.get("save", True) + if isinstance(save, str): + save = save.lower() not in ("0", "false", "no", "") + + def _go(): + if dirty and save: + app._on_quit_choice("save") # saves, then exits + else: + app._on_quit_choice("discard") + # answer first, then tear down (so this response still gets written) - asyncio.get_running_loop().call_later(0.2, app.exit) - return {"ok": True, "quitting": True} + asyncio.get_running_loop().call_later(0.2, _go) + return {"ok": True, "quitting": True, "saving": bool(dirty and save), + "dirty": dirty} if method in _PROGRAM_METHODS and app.program is None: raise ValueError("not ready: still connecting / loading functions") @@ -531,6 +895,49 @@ class RpcServer: return functions(app, params.get("filter"), int(params.get("limit", 50))) # -- projects ------------------------------------------------------ # + if method == "diag": + # What has been swallowed lately. A driver that got "success" and a + # pane showing nothing has no other way to ask; inside a + # full-screen TUI there is nowhere for a traceback to go. + if params.get("clear"): + diag.clear() + return {"cleared": True} + return {"recent": diag.recent(int(params.get("n", 10))), + "log": os.environ.get("IDATUI_LOG") or None} + + if method == "trace": + tc = app.trace_ctl + if tc.trace is None: + raise ValueError("no trace loaded (launch with --trace FILE)") + t = tc.trace + if "seek" in params: + v = params["seek"] + # "!50" seeks a percentage, like Tenet's timestamp shell. + if isinstance(v, str) and v.startswith("!"): + idx = int(float(v[1:]) * (t.length - 1) / 100.0) + else: + idx = int(str(v).replace(",", ""), 0) if isinstance(v, str) else int(v) + tc.seek(idx) + elif "goto" in params: # first execution of an address/name + tgt = params["goto"] + ea = (int(str(tgt), 0) if str(tgt).lower().startswith("0x") + else app.program.resolve(str(tgt))) + first = t.first_execution(ea) + if first is None: + raise ValueError(f"{tgt} never executed in this trace") + tc.seek(first) + elif "step" in params: + n = int(params.get("step") or 1) + over = bool(params.get("over")) + for _ in range(abs(n)): + (tc.step_over if over else tc.step)(1 if n > 0 else -1) + await settle(app, timeout=float(params.get("timeout", 20.0))) + snap = snapshot(app) + snap["trace"] = {"idx": tc.t, "length": t.length, + "pc": hex(t.ip(tc.t)), + "changed": sorted(t.changed(tc.t))} + return snap + if method == "binaries": if app._project is None: raise ValueError("not a project session (launch with --project)") @@ -609,15 +1016,106 @@ class RpcServer: target = str(params.get("target", "")) pred = self._goto_target_pred(target) await self._fill_prompt("g", "goto", target, delay, clear=False) - await settle(app, pred, timeout=timeout) + ok = await settle(app, pred, timeout=timeout) + if pred is not None and not ok: + # A goto that silently "succeeds" without moving is worse than an + # error: on a big database the listing build can outrun the + # default timeout, and every subsequent rename/comment then lands + # on the function the caller *used* to be looking at. + raise TimeoutError( + f"goto {target!r} did not land within {timeout}s " + f"(still at {_where(app)}); retry with a larger timeout=") return snapshot(app) + if method == "define": + kind = str(params.get("kind", "code")).lower() + if kind not in _DEFINE_KEYS: + raise ValueError( + f"unknown define kind {kind!r}; one of " + f"{', '.join(sorted(_DEFINE_KEYS))}") + target = params.get("target") + if target not in (None, ""): + # Land on the address first. A raw image is mostly *undefined*, + # so the target usually has no name and no function — the goto + # predicate can't be address-based, only "we moved". + await self._fill_prompt("g", "goto", str(target), delay, + clear=False) + await settle(app, timeout=timeout) + if app.is_hex: + # backslash leaves hex for the code view (which may be decomp). + await self._press(["backslash"], + lambda: not app.is_hex, timeout, + "leave the hex view") + if app.is_decomp: + # These bindings live on the listing; in the decompiler the key + # would be swallowed or do something else entirely. + await self._press(["tab"], lambda: app.is_listing, + timeout, "switch to the listing") + if not app.is_listing: + raise RuntimeError( + f"define needs the listing view, but the active pane is " + f"{app._active!r}") + snap = await self._press([_DEFINE_KEYS[kind]], timeout=timeout, + what=f"define {kind}") + snap["define"] = {"kind": kind, "status": snap.get("status", "")} + return snap + + if method == "opfmt": + mode = str(params.get("mode", "cycle")).lower() + if mode not in _OPFMT_MODES: + raise ValueError(f"unknown opfmt mode {mode!r}; one of " + f"{', '.join(_OPFMT_MODES)}") + target = params.get("target") + if target not in (None, ""): + await self._fill_prompt("g", "goto", str(target), delay, + clear=False) + await settle(app, timeout=timeout) + if app.is_hex: + await self._press(["backslash"], lambda: not app.is_hex, + timeout, "leave the hex view") + view = _active_widget(app) + if isinstance(view, HexView): + raise RuntimeError("opfmt needs a code view, not the hex view") + if params.get("word"): + # Land the column on the literal first: WHICH operand gets + # reformatted is decided by where the cursor is. + if not cursor_on(app, str(params["word"]), params.get("line"), + int(params.get("occurrence", 1) or 1)): + raise RuntimeError( + f"{params['word']!r} is not on screen in this view, so " + f"there is no literal to reformat") + await drain(app) + elif params.get("line") is not None or params.get("col") is not None: + place_cursor(view, params.get("line"), params.get("col")) + await drain(app) + before = _where(app) + if mode in _OPFMT_KEYS: + snap = await self._press([_OPFMT_KEYS[mode]], timeout=timeout, + what=f"opfmt {mode}") + else: + view.focus() + view.action_op_format(mode) + await settle(app, timeout=timeout) + snap = snapshot(app) + snap["opfmt"] = {"mode": mode, "at": before, + "status": snap.get("status", "")} + return snap + + if method == "rename_many": + return await self._rename_many(params, timeout) + if method == "rename": await self._fill_prompt("n", "rename", str(params["name"]), delay, clear=True) await settle(app, timeout=timeout) return snapshot(app) if method == "comment": - await self._fill_prompt("semicolon", "comment", str(params["text"]), delay, + # Comments can be long; skip the per-char delay so the agent isn't + # blocked for seconds watching the typing animation. Also: the + # prompt is single-line, so literal newlines (0x0a) get swallowed by + # the Input widget. The app's _do_comment converts the two-char + # sequence '\n' into a real newline for IDA, so we escape here. + ctext = str(params["text"]).replace("\n", "\\n") + await self._fill_prompt("semicolon", "comment", ctext, 0, clear=True) await settle(app, timeout=timeout) return snapshot(app) @@ -628,7 +1126,8 @@ class RpcServer: if method == "follow": depth = len(app._nav) - return await self._press(["enter"], lambda: len(app._nav) > depth, timeout) + return await self._press(["enter"], lambda: len(app._nav) > depth, + timeout, "follow") if method == "back": return await self._press(["escape"], timeout=timeout) if method == "toggle_view": @@ -639,25 +1138,34 @@ class RpcServer: if app._active != before: return True # Fallback case: a tab toward pseudocode on a function Hex-Rays - # can't decompile snaps `_active` back to disasm (see + # can't decompile lands back on the LISTING (see # App._apply_decomp), so `_active` never changes and the naive # `_active != before` predicate would block for the full # timeout. Treat "requested decomp but it's known-failed" as # settled (the decompile is cached, so this is cheap). cur = app._cur - if before == "disasm" and cur is not None: + if before == ViewMode.LISTING and cur is not None: try: return app.program.decompile(cur.ea).failed except Exception: # noqa: BLE001 return False return False - return await self._press(["tab"], _toggled, timeout) + return await self._press(["tab"], _toggled, timeout, "toggle_view") if method == "hex": - return await self._press(["backslash"], lambda: app._active == "hex", timeout) + # Backslash TOGGLES the hex view, so the predicate has to be "the + # mode flipped", not "we are in hex". Waiting for is_hex meant the + # call that LEAVES hex could never be satisfied and always timed + # out -- a driver could open the hex view but never close it. + was_hex = app.is_hex + return await self._press(["backslash"], lambda: app.is_hex != was_hex, + timeout, "hex") + if method == "graph": + return await self._graph(params, timeout) if method == "xrefs": return await self._press( - ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", timeout) + ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", + timeout, "xrefs") if method == "symbols": await app._press_keys(["ctrl+n"]) await settle(app, lambda: type(app.screen).__name__ == "SymbolPalette", timeout=10) @@ -668,7 +1176,47 @@ class RpcServer: return snapshot(app) if method == "structs": return await self._press( - ["ctrl+t"], lambda: type(app.screen).__name__ == "StructEditor", timeout) + ["ctrl+t"], lambda: type(app.screen).__name__ == "StructEditor", + timeout, "structs") + if method == "find": + from . import search as _search + q = str(params.get("query", "")) + forced = params.get("mode") + forced = None if forced in (None, "auto") else str(forced) + mode, cleaned = _search.classify(q, forced) + if mode == _search.BYTES: + problem = _search.pattern_problem(cleaned) + if problem: + raise ValueError(f"find: {problem}") + cleaned = _search.normalise_pattern(cleaned) + hits, err, truncated = await asyncio.to_thread( + app.program.search, cleaned, mode, + limit=int(params.get("limit", 500)), + regex=bool(params.get("regex")), case=bool(params.get("case"))) + if err: + raise ValueError(f"find: {err}") + return {"mode": mode, "query": cleaned, "truncated": truncated, + "hits": [{"addr": hex(h.addr), "head": hex(h.head), + "line": h.line, "func": h.func, + "seg": h.seg} for h in hits]} + if method == "export": + # Deliberately NOT driven through the prompt: this is the one verb + # whose whole point is the file it leaves behind, and a driver needs + # the path back, not a screenshot of a prompt closing. + from . import findings + path = params.get("path") + app.journal.load(app.program) + app.journal.flush(app.program) + out, f = await asyncio.to_thread( + findings.export, app.program, app._open_path or "", + str(path) if path else None, + types=bool(params.get("types", True)), journal=app.journal) + app._status(f"exported findings → {out}", priority=True) + await drain(app) + return {"path": out, "comments": len(f.comments), + "names": len(findings._user_names(f)), + "types": len(f.types), "functions": f.n_functions, + "bytes": os.path.getsize(out) if os.path.exists(out) else 0} if method == "close": return await self._press(["escape"], timeout=timeout) if method == "save": @@ -713,13 +1261,7 @@ class RpcServer: w = _active_widget(app) if isinstance(w, HexView): raise ValueError("cursor: not supported in the hex view (use goto)") - if "line" in params and params["line"] is not None: - w.cursor = max(0, min(getattr(w, "total", 1) - 1, int(params["line"]))) - if "col" in params and params["col"] is not None: - w.cursor_x = max(0, int(params["col"])) - if hasattr(w, "_after_cursor_move"): - w._after_cursor_move() - w.refresh() + place_cursor(w, params.get("line"), params.get("col")) await drain(app) return snapshot(app) diff --git a/idatui/rpcclient.py b/idatui/rpcclient.py index e0a5602..4231d62 100644 --- a/idatui/rpcclient.py +++ b/idatui/rpcclient.py @@ -28,7 +28,15 @@ class RpcClient: #: Default read timeout (s). Bounds any single call so a slow/hung server #: (e.g. Hex-Rays grinding on an undecompilable function) can't block the #: CLI forever. Override via ctor or the IDATUI_RPC_TIMEOUT env var. - DEFAULT_TIMEOUT = 90.0 + #: + #: 90s was too tight on real firmware: a comment on a 42k-line flat listing + #: (one segment, no function boundaries to limit the rebuild) took 26-106s, + #: so the client reported "no response ... server busy or the op is hung" + #: for edits that had in fact been applied. A driver that believes a + #: successful edit failed is worse than a slow one -- it redoes the work, or + #: "fixes" something that was never broken. Real hangs still get caught, + #: just later. + DEFAULT_TIMEOUT = 300.0 def __init__(self, sock_path: str, timeout: float | None = None): self.path = sock_path diff --git a/idatui/search.py b/idatui/search.py new file mode 100644 index 0000000..feb7e3d --- /dev/null +++ b/idatui/search.py @@ -0,0 +1,112 @@ +"""What did the user mean by that query? (Ctrl+F.) + +Database-wide search comes in two kinds -- **text** through the disassembly and +**bytes** through the image -- and asking people to pick a mode before they type +is a tax on every search. So the query decides, and the rule has to be +conservative in one specific direction: a hex-looking word is often real text +(``add``, ``dec``, ``dead``, ``beef``, ``cafe`` are all valid hex AND things you +would search for), while nobody types ``48 8b ?? c3`` meaning prose. + +Hence: bytes only when the query is unambiguously a byte pattern -- several +whitespace/comma separated tokens that are all hex pairs or wildcards, or any +query containing a ``?``. Everything else is text, and the two explicit prefixes +(``hex:`` / ``text:``) settle any argument, as does F2 in the palette. + +Pure: no IDA, no Textual, so ``tests/test_search.py`` runs it offline. +""" + +from __future__ import annotations + +import re + +TEXT = "text" +BYTES = "bytes" + +#: One token of a byte pattern: a hex pair, a wildcard nibble ("8?"), or a bare +#: "?" standing for a whole byte. IDA's find_bytes accepts all three. +_TOKEN = re.compile(r"^(?:[0-9A-Fa-f?]{2}|\?)$") + +#: A quoted literal inside a pattern ('"Hello", 0'), which IDA also accepts. +_QUOTED = re.compile(r'"[^"]*"') + + +def looks_like_bytes(query: str) -> bool: + """True when ``query`` can only sensibly be a byte pattern.""" + q = (query or "").strip() + if not q: + return False + if _QUOTED.search(q): + return True + tokens = [t for t in re.split(r"[\s,]+", q) if t] + if not all(_TOKEN.match(t) for t in tokens): + return False + # A single token is ambiguous ("ff" is also a word); a wildcard never is. + return len(tokens) > 1 or "?" in q + + +def probably_meant_bytes(query: str) -> bool: + """True for a query that is *shaped* like bytes but does not parse. + + ``48 zz c3`` is a typo in a byte pattern, and treating it as a text search + answers "no match" -- the most misleading thing a search can say, because it + is indistinguishable from "those bytes are not in this binary". Every token + being byte-sized is the tell; ``add ff`` (a three-letter token) is not, and + stays text. + """ + tokens = [t for t in re.split(r"[\s,]+", (query or "").strip()) if t] + if len(tokens) < 2 or any(len(t) > 2 for t in tokens): + return False + return any(_TOKEN.match(t) for t in tokens) + + +def classify(query: str, forced: str | None = None) -> tuple[str, str]: + """Return ``(mode, cleaned_query)``. + + An explicit ``hex:``/``bytes:``/``text:`` prefix wins, then ``forced`` (the + palette's F2), then the shape of the query. + """ + q = (query or "").strip() + low = q.lower() + for prefix, mode in (("hex:", BYTES), ("bytes:", BYTES), ("text:", TEXT)): + if low.startswith(prefix): + return (mode, q[len(prefix):].strip()) + if forced in (TEXT, BYTES): + return (forced, q) + if looks_like_bytes(q) or probably_meant_bytes(q): + return (BYTES, q) + return (TEXT, q) + + +def normalise_pattern(pattern: str) -> str: + """Tidy a byte pattern for IDA: single spaces, commas as separators. + + ``48 8B?? C3``, ``48,8b,??,c3`` and ``48 8b ?? c3`` are the same search; + people paste all three (the middle one out of a signature file). + """ + q = (pattern or "").strip() + if _QUOTED.search(q): + return q # a quoted literal owns its own spacing + q = q.replace(",", " ") + # "488B??C3" -- a bare hex run with no separators at all. + if " " not in q and len(q) > 2 and len(q) % 2 == 0: + q = " ".join(q[i:i + 2] for i in range(0, len(q), 2)) + return " ".join(q.split()) + + +def pattern_problem(pattern: str) -> str | None: + """A human explanation if this cannot be a byte pattern, else ``None``. + + Checked before the round trip, because IDA's own message for a bad pattern + is empty about half the time. + """ + q = normalise_pattern(pattern) + if not q: + return "type some bytes, e.g. 48 8b ?? c3" + if _QUOTED.search(q): + return None + tokens = [t for t in q.split() if t] + bad = [t for t in tokens if not _TOKEN.match(t)] + if bad: + return (f"{bad[0]!r} is not a byte: use hex pairs, ? wildcards " + 'or a "quoted string"') + return None diff --git a/idatui/trace.py b/idatui/trace.py new file mode 100644 index 0000000..931f918 --- /dev/null +++ b/idatui/trace.py @@ -0,0 +1,495 @@ +"""Reading Tenet execution traces. + +A Tenet trace is a line-per-instruction delta log:: + + rax=0x3c,rbx=0x0,...,rip=0x7ffff6faaae0 # full state on the first line + rip=0x7ffff6faaae4 # then only what changed + r9=0x7ffff6762e60,rip=0x7ffff6faaae9,mw=0x7ffff6f9b7c8:08c177f6ff7f0000 + +Registers that changed, the PC on every line, and every memory access *with its +bytes*. That is enough to reconstruct any register or any memory address at any +point in time, forwards or backwards, which is the whole trick. + +This is our own reader, not a port. The reference implementation +(``~/dev/tenet/tenet-original/plugins/tenet/trace/``) packs the trace into +segments with compressed address/mask tables, which earns its keep for its Qt +timeline; we need different queries and would rather own the ~400 lines than +inherit 3700. It is differential-tested against that implementation +(``tests/test_trace_vs_tenet.py``) so "our own" doesn't quietly mean "different". + +The index is built around the query the UI actually asks, which the reference +answers one address at a time: **which timestamps executed this set of +addresses**. A listing row is one address, but a pseudocode line covers many, so +``by_ip`` maps address -> timestamps and set queries are unions of those. +""" + +from __future__ import annotations + +import array +import bisect +import os +import re +from dataclasses import dataclass, field + +#: Tenet packs its register delta into a uint32, so a trace arch may name at +#: most 32 registers. Ours is discovered from the trace instead of declared, +#: but the cap is worth knowing when a trace looks short of registers. +MAX_REGISTERS = 32 + +_MEM_RE = re.compile(r"^m(r|w|rw)$") + + +@dataclass +class TraceInfo: + """The sidecar ``<prefix>.info`` written by our QEMU tracer. + + Optional: a trace from another tracer has none, and everything here can be + recovered or guessed from the log itself. + """ + + arch: str = "" + mode: str = "" + binary: str = "" + start_code: int = 0 + end_code: int = 0 + entry_code: int = 0 + traced: str = "" + + @classmethod + def load(cls, path: str) -> "TraceInfo | None": + try: + with open(path) as f: + raw = dict( + ln.strip().split("=", 1) for ln in f if "=" in ln) + except OSError: + return None + def num(k): + try: + return int(raw.get(k, "0"), 0) + except ValueError: + return 0 + return cls(arch=raw.get("arch", ""), mode=raw.get("mode", ""), + binary=raw.get("binary", ""), start_code=num("start_code"), + end_code=num("end_code"), entry_code=num("entry_code"), + traced=raw.get("traced", "")) + + +@dataclass +class MemOp: + """One memory access made by one instruction. + + ``addr`` is the address the TRACE recorded — not slid onto the database. + Most accesses are stack or heap, which have no counterpart in the database + at all, and applying the image's relocation to a stack pointer produces a + nonsense address (it went negative in testing). Only addresses inside the + image can be translated, and the caller knows when that applies. + """ + + addr: int + data: bytes + write: bool + + @property + def end(self) -> int: + return self.addr + len(self.data) + + +@dataclass +class Trace: + """An indexed Tenet trace. + + Timestamps are indices into the executed-instruction sequence: 0 is the + first instruction, ``length - 1`` the last. + """ + + path: str = "" + info: TraceInfo | None = None + #: PC per timestamp. + ips: array.array = field(default_factory=lambda: array.array("Q")) + #: name -> (timestamps of change, value at each change). A register's value + #: at time t is the last change at or before t; the first line carries a + #: full state dump, so every register has an entry at 0. + reg_at: dict[str, tuple[array.array, array.array]] = field(default_factory=dict) + #: address -> timestamps that executed it (ascending, by construction). + by_ip: dict[int, array.array] = field(default_factory=dict) + #: memory accesses, parallel arrays indexed by access number. + mem_idx: array.array = field(default_factory=lambda: array.array("I")) + mem_addr: array.array = field(default_factory=lambda: array.array("Q")) + mem_write: bytearray = field(default_factory=bytearray) + mem_off: array.array = field(default_factory=lambda: array.array("Q")) + mem_len: array.array = field(default_factory=lambda: array.array("H")) + mem_blob: bytearray = field(default_factory=bytearray) + #: first access number of each timestamp, plus a final sentinel. + mem_row: array.array = field(default_factory=lambda: array.array("I")) + #: applied so trace addresses line up with the database (see ``rebase``). + slide: int = 0 + #: accesses ordered by address, built on first memory query. + _mem_order: list | None = None + _mem_starts: list = field(default_factory=list) + _mem_maxlen: int = 0 + + # -- construction ------------------------------------------------------ # + @classmethod + def load(cls, path: str, progress=None, limit: int = 0) -> "Trace": + """Parse a text trace. ``progress(lines)`` is called every 50k lines.""" + t = cls(path=os.path.abspath(path)) + base = path[:-6] if path.endswith(".0.log") else os.path.splitext(path)[0] + t.info = TraceInfo.load(base + ".info") + regs: dict[str, tuple[array.array, array.array]] = {} + ips, by_ip = t.ips, t.by_ip + n = 0 + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + ip = None + t.mem_row.append(len(t.mem_idx)) + for part in line.split(","): + key, _, val = part.partition("=") + if not val: + continue + key = key.strip().lower() + if _MEM_RE.match(key): + addr_s, _, data_s = val.partition(":") + try: + addr = int(addr_s, 16) + data = bytes.fromhex(data_s) + except ValueError: + continue + # 'mrw' is one access that both reads and writes; record + # the write, which is what the memory state follows. + t.mem_idx.append(n) + t.mem_addr.append(addr) + t.mem_write.append(0 if key == "mr" else 1) + t.mem_off.append(len(t.mem_blob)) + t.mem_len.append(len(data)) + t.mem_blob += data + continue + try: + v = int(val, 16) + except ValueError: + continue + slot = regs.get(key) + if slot is None: + slot = regs[key] = (array.array("I"), array.array("Q")) + slot[0].append(n) + slot[1].append(v) + ip = v if key in ("rip", "eip", "pc") else ip + if ip is None: + # No PC on this line: the format says always emit it, and + # without it the line cannot be placed. Carry the previous + # one rather than dropping the instruction. + ip = ips[-1] if ips else 0 + ips.append(ip) + where = by_ip.get(ip) + if where is None: + where = by_ip[ip] = array.array("I") + where.append(n) + n += 1 + if progress is not None and n % 50000 == 0: + progress(n) + if limit and n >= limit: + break + t.mem_row.append(len(t.mem_idx)) + t.reg_at = regs + return t + + # -- basics ------------------------------------------------------------ # + def __len__(self) -> int: + return len(self.ips) + + @property + def length(self) -> int: + return len(self.ips) + + @property + def registers(self) -> list[str]: + """Register names in the trace, PC last (the order tracers emit).""" + return sorted(self.reg_at, key=lambda r: (r in ("rip", "eip", "pc"), r)) + + @property + def pc_name(self) -> str: + for n in ("rip", "eip", "pc"): + if n in self.reg_at: + return n + return "" + + def ip(self, idx: int) -> int: + """PC at ``idx``, in DATABASE addresses (slide applied).""" + return self.ips[idx] + self.slide + + def raw_ip(self, idx: int) -> int: + return self.ips[idx] + + # -- register state ----------------------------------------------------- # + def register(self, name: str, idx: int) -> int | None: + """Value of ``name`` at ``idx``, or None if the trace never set it.""" + slot = self.reg_at.get(name.lower()) + if slot is None: + return None + idxs, vals = slot + i = bisect.bisect_right(idxs, idx) - 1 + return vals[i] if i >= 0 else None + + def register_state(self, idx: int) -> dict[str, int]: + return {n: v for n in self.reg_at + if (v := self.register(n, idx)) is not None} + + def changed(self, idx: int) -> set[str]: + """Registers written BY the instruction at ``idx`` (what the line said). + + Used to highlight what an instruction actually did, which is the reason + a delta trace is readable at all. + """ + out = set() + for n, (idxs, _vals) in self.reg_at.items(): + i = bisect.bisect_left(idxs, idx) + if i < len(idxs) and idxs[i] == idx: + out.add(n) + return out + + def last_write(self, name: str, idx: int) -> int | None: + """Timestamp of the write that produced ``name``'s value at ``idx``. + + "Which instruction set this register?" — the question that motivates a + trace explorer in the first place. + """ + slot = self.reg_at.get(name.lower()) + if slot is None: + return None + i = bisect.bisect_right(slot[0], idx) - 1 + return slot[0][i] if i >= 0 else None + + def next_write(self, name: str, idx: int) -> int | None: + slot = self.reg_at.get(name.lower()) + if slot is None: + return None + i = bisect.bisect_right(slot[0], idx) + return slot[0][i] if i < len(slot[0]) else None + + # -- memory ------------------------------------------------------------- # + def memory_ops(self, idx: int) -> list[MemOp]: + """Accesses made by the instruction at ``idx``.""" + if not (0 <= idx < len(self.mem_row) - 1): + return [] + lo, hi = self.mem_row[idx], self.mem_row[idx + 1] + out = [] + for k in range(lo, hi): + off, ln = self.mem_off[k], self.mem_len[k] + out.append(MemOp(addr=self.mem_addr[k], + data=bytes(self.mem_blob[off:off + ln]), + write=bool(self.mem_write[k]))) + return out + + # -- memory state ------------------------------------------------------- # + def _mem_index(self) -> None: + """Order the accesses by address, once. + + Queries ask "what was at this window at time t", so the accesses that + matter are the few touching that window — not the tens of thousands in + the trace. Sorting by address makes those a bisect away; sorting by time + (the order they arrive in) would mean scanning everything per repaint. + """ + if self._mem_order is not None: + return + order = sorted(range(len(self.mem_addr)), key=lambda k: self.mem_addr[k]) + self._mem_order = order + self._mem_starts = [self.mem_addr[k] for k in order] + self._mem_maxlen = max(self.mem_len) if len(self.mem_len) else 0 + + def memory_raw(self, addr: int, length: int, + idx: int | None = None) -> tuple[bytes, bytes]: + """Memory at a TRACE address (no slide). + + The stack lives here. Measured on two real traces, 0% of memory accesses + fall inside the image — every one is stack or heap — so a query that + insists on database addresses can't answer the question anyone actually + has about memory in a trace. + """ + return self.memory(addr + self.slide, length, idx) + + def memory(self, addr: int, length: int, + idx: int | None = None) -> tuple[bytes, bytes]: + """``(data, known)`` for ``length`` bytes at ``addr`` as of ``idx``. + + ``known`` is a byte-per-byte mask: a trace only says what it saw, so a + byte nobody read or wrote is genuinely unknown and must not be drawn as + zero. That distinction is the whole value of reading memory from a trace + rather than from the database — the database has the file's bytes, the + trace has what was actually there at that instant. + + Reads count as evidence, not just writes: an instruction reading a byte + reveals what it held then. + """ + if idx is None: + idx = self.length - 1 + length = max(int(length), 0) + out, known = bytearray(length), bytearray(length) + if not length or not len(self.mem_idx): + return bytes(out), bytes(known) + self._mem_index() + raw = addr - self.slide + best = [-1] * length + import bisect as _b + lo = _b.bisect_left(self._mem_starts, raw - self._mem_maxlen) + hi = _b.bisect_right(self._mem_starts, raw + length - 1) + for pos in range(lo, hi): + k = self._mem_order[pos] + t = self.mem_idx[k] + if t > idx: + continue + a, ln = self.mem_addr[k], self.mem_len[k] + s, e = max(a, raw), min(a + ln, raw + length) + if s >= e: + continue + off = self.mem_off[k] + for b in range(s, e): + j = b - raw + # >= not >: several accesses can share a timestamp (an + # instruction that reads and writes), and the later entry on the + # line is the one that stands. + if t >= best[j]: + best[j] = t + out[j] = self.mem_blob[off + (b - a)] + known[j] = 1 + return bytes(out), bytes(known) + + def memory_writes(self, addr: int, length: int) -> list[int]: + """Timestamps that WROTE any byte of ``[addr, addr+length)``.""" + return self._mem_touch(addr, length, writes=True) + + def memory_accesses(self, addr: int, length: int) -> list[int]: + """Timestamps that read or wrote any byte of the range.""" + return self._mem_touch(addr, length, writes=False) + + def _mem_touch(self, addr: int, length: int, writes: bool) -> list[int]: + if not len(self.mem_idx) or length <= 0: + return [] + self._mem_index() + raw = addr - self.slide + import bisect as _b + lo = _b.bisect_left(self._mem_starts, raw - self._mem_maxlen) + hi = _b.bisect_right(self._mem_starts, raw + length - 1) + out = set() + for pos in range(lo, hi): + k = self._mem_order[pos] + a, ln = self.mem_addr[k], self.mem_len[k] + if a + ln <= raw or a >= raw + length: + continue + if writes and not self.mem_write[k]: + continue + out.add(self.mem_idx[k]) + return sorted(out) + + # -- execution queries (what painting is built on) ---------------------- # + def executions(self, ea: int) -> array.array: + """Every timestamp that executed ``ea`` (database address).""" + return self.by_ip.get(ea - self.slide, array.array("I")) + + def executions_between(self, ea: int, lo: int, hi: int) -> list[int]: + ts = self.executions(ea) + a = bisect.bisect_left(ts, lo) + b = bisect.bisect_right(ts, hi) + return list(ts[a:b]) + + def hits(self, eas) -> dict[int, int]: + """{address: execution count} for a set of addresses. + + The set form is the point: painting a listing row needs one address, but + a pseudocode line covers many, and asking per-address would mean a + lookup per instruction per repaint. + """ + out = {} + for ea in eas: + ts = self.by_ip.get(ea - self.slide) + if ts: + out[ea] = len(ts) + return out + + def prev_ips(self, idx: int, n: int) -> list[int]: + """Addresses executed in the ``n`` steps before ``idx`` (nearest first). + + A trail, not all of history: showing every address the trace ever + touched says almost nothing on a loop-heavy program, whereas the last + few dozen steps say how you GOT here. + """ + lo = max(idx - n, 0) + return [self.ips[i] + self.slide for i in range(idx - 1, lo - 1, -1)] + + def next_ips(self, idx: int, n: int) -> list[int]: + """Addresses executed in the ``n`` steps after ``idx`` (nearest first).""" + hi = min(idx + n + 1, self.length) + return [self.ips[i] + self.slide for i in range(idx + 1, hi)] + + def trail(self, idx: int, n: int = 96) -> dict[int, str]: + """{address: 'now' | 'past' | 'future'} around ``idx``. + + Where an address appears on both sides — a loop body, which is most of + them — the nearer side wins, because that's the one that explains the + step you are about to take or just took. + """ + out: dict[int, str] = {} + for k, ea in enumerate(self.next_ips(idx, n)): + out.setdefault(ea, "future") + for k, ea in enumerate(self.prev_ips(idx, n)): + prev = out.get(ea) + if prev is None: + out[ea] = "past" + elif prev == "future": + # Same distance rule as above, resolved by which loop found it + # first would be arbitrary; compare real distances instead. + fwd = next((i for i, a in enumerate(self.next_ips(idx, n)) if a == ea), n) + if k < fwd: + out[ea] = "past" + if 0 <= idx < self.length: + out[self.ips[idx] + self.slide] = "now" + return out + + def first_execution(self, ea: int) -> int | None: + ts = self.executions(ea) + return ts[0] if ts else None + + def next_execution(self, ea: int, idx: int) -> int | None: + ts = self.executions(ea) + i = bisect.bisect_right(ts, idx) + return ts[i] if i < len(ts) else None + + def prev_execution(self, ea: int, idx: int) -> int | None: + ts = self.executions(ea) + i = bisect.bisect_left(ts, idx) - 1 + return ts[i] if i >= 0 else None + + # -- lining the trace up with the database ------------------------------ # + def rebase(self, db_addresses) -> int: + """Find the slide between trace addresses and database addresses. + + A traced process is relocated (ASLR, or a PIE base the database doesn't + share): our echo trace runs at 0x7ffff6faa000 while the database has the + same code at 0x2490. Nothing lines up until this is solved, so it is not + optional. + + Page offsets survive relocation — only whole pages move — so the low 12 + bits of an instruction address are invariant. Bucket the database's + addresses by those bits, and for each trace address the candidate slides + are the differences to database addresses in its bucket. The slide that + the most instructions agree on wins. + """ + buckets: dict[int, list[int]] = {} + for a in db_addresses: + buckets.setdefault(a & 0xFFF, []).append(a) + if not buckets: + return 0 + votes: dict[int, int] = {} + # A sample is enough and keeps this O(1)-ish on a 10M trace; unique + # addresses, because a hot loop shouldn't outvote the rest of the code. + for ea in list(self.by_ip)[:4096]: + for cand in buckets.get(ea & 0xFFF, ()): + votes[cand - ea] = votes.get(cand - ea, 0) + 1 + if not votes: + return 0 + best, n = max(votes.items(), key=lambda kv: kv[1]) + return best if n > 1 else 0 + + def apply_slide(self, slide: int) -> None: + self.slide = int(slide) diff --git a/idatui/trace_ctl.py b/idatui/trace_ctl.py new file mode 100644 index 0000000..8072801 --- /dev/null +++ b/idatui/trace_ctl.py @@ -0,0 +1,421 @@ +"""Execution-trace navigation: the controller behind `--trace FILE`. + +`idatui.trace` is the *reader* (parse the Tenet delta log, answer questions about +it). This is the half that used to live in `IdaTui`: the trace's position in +time, the trail painted onto the listing/pseudocode/hex, and every key that +moves through it (`[` `]` `{` `}` `<` `>` `W`). + +Why it's a plain object and not a mixin: Textual only merges ``BINDINGS`` from +``DOMNode`` subclasses, so a mixin's bindings are silently dropped. The keys +therefore stay declared on ``IdaTui``, whose ``action_*`` methods are one-line +delegates into here. Same reason the ``@work(thread=True)`` entry point stays on +the app -- Textual's worker machinery wants a ``DOMNode`` host. + +The controller owns the trace state. ``IdaTui`` keeps forwarding properties +(``app._trace``, ``app._t``, ``app._trail_map``...) because the pilot suite and +the RPC layer read them by those names; see ``IdaTui._trace``. +""" +from __future__ import annotations + +import bisect +import os +from typing import TYPE_CHECKING + +from . import diag + +if TYPE_CHECKING: # pragma: no cover + from .app import IdaTui + +_app_mod = None + + +def _views(): + """The widget classes, imported lazily. + + ``app.py`` imports this module at the top, so importing it back at module + scope would be a cycle. Nothing here needs the classes until a trace is + actually driven, by which point ``app`` is fully imported. + """ + global _app_mod + if _app_mod is None: + from . import app as _m + _app_mod = _m + return _app_mod + + +class TraceController: + """Where we are in the trace, and everything that moves us.""" + + def __init__(self, app: "IdaTui", path: str = "") -> None: + self.app = app + self.path = path or "" # the Tenet trace to explore, if any + self.trace = None # the loaded Trace, once analysed + self.t = 0 # current timestamp in that trace + self.trail_map = [] # decomp_map for trail_map_ea + self.trail_map_ea = None + self.trail_line_of: dict[int, int] = {} # ea -> pseudocode line + self.trail_eas: list[int] = [] # sorted keys of trail_line_of + self.trail_span = None # ea span of that function + self.pending_line = None # step waiting on a re-decompile + + @property + def armed(self) -> bool: + """A trace was asked for but hasn't been parsed yet.""" + return bool(self.path) and self.trace is None + + def _need_trace(self) -> bool: + """Complain once, in one voice, rather than at four call sites.""" + if self.trace is None: + self.app._status("no trace loaded (--trace FILE)") + return False + return True + + # -- loading ------------------------------------------------------------ # + def load(self) -> None: + """Parse the trace and line it up with the database. + + Runs after the function index exists: rebasing needs the database's + addresses, and without it nothing in the trace matches anything on + screen (our echo trace runs at 0x7ffff6faa000; the database has that + code at 0x2000). + + Called on a worker thread, so every touch of the UI hops back. + """ + from .trace import Trace + app = self.app + path = self.path + try: + def note(n): + app.call_from_thread( + app._status, f"trace: {n:,} instructions\u2026") + trace = Trace.load(path, progress=note) + except OSError as e: + app.call_from_thread(app._status, f"trace: {e}") + return + if not trace.length: + app.call_from_thread( + app._status, f"trace: {os.path.basename(path)} is empty") + return + idx = app._func_index + addrs = [f.addr for f in idx.all_loaded()] if idx is not None else [] + slide = trace.rebase(addrs) + trace.apply_slide(slide) + hit = sum(1 for f in (idx.all_loaded() if idx else []) + if trace.executions(f.addr)) + app.call_from_thread(self.ready, trace, slide, hit) + + def ready(self, trace, slide: int, hit: int) -> None: + app = self.app + self.trace = trace + self.t = 0 + dock = app.query_one(_views().TraceDock) + dock.display = True + dock.show(trace, 0) + where = (f"rebased {slide:+#x}" if slide else "no rebase needed") + app._status(f"trace: {trace.length:,} instructions, {hit} functions " + f"touched ({where})", priority=True) + self.seek(0, follow=True) + + # -- trace navigation --------------------------------------------------- # + def seek(self, idx: int, follow: bool = True) -> None: + """Move to timestamp ``idx``; ``follow`` takes the code view with it.""" + app = self.app + t = self.trace + if t is None or not t.length: + return + # A seek invalidates any navigation still in flight. They run in workers + # and finish out of order: the trace's opening seek lands on the entry + # point, takes a while, and used to arrive AFTER later seeks — dragging + # the cursor back to _start while the trace was elsewhere, permanently. + # + # Bumped HERE and not in _goto_ea. Doing it for every navigation is the + # more general rule ("the last thing you asked for wins") but it also + # lets an ordinary follow be dropped by whatever navigates next, and the + # only evidence I have is about seeks. Narrow fix for the measured bug. + app._nav_seq += 1 + self.t = max(0, min(int(idx), t.length - 1)) + app.query_one(_views().TraceDock).show(t, self.t) + self.paint_trail() + if not follow: + return + pc = t.ip(self.t) + if app._split and self.seek_split(pc): + return + # Stay in whichever view you're reading. Without prefer_decomp a step + # from the pseudocode navigates to an address, which opens the listing — + # so stepping through C threw you out of C on the first keypress. + app._goto_ea(pc, push=False, + prefer_decomp=(app.is_decomp)) + + def seek_split(self, pc: int) -> bool: + """Put BOTH panes on ``pc``. True if handled. + + Normal navigation moves one pane and gives the companion a band, never a + cursor — that rule exists so the two can't chase each other. A trace step + isn't navigation though: time is a single global position, and both views + are showing the same instant, so both cursors belong on it. + + The scroll anchoring is unchanged: after placing the cursors, the usual + _sync_split still bands the companion and aligns it to the driver's + screen row, so the eye tracks straight across. + """ + app = self.app + M = _views() + lst = app.query_one(M.ListingView) + if lst.model is None: + return False + row = lst.model.ensure_ea(pc) + if row is None or row < 0: + return False # not in this listing (other segment): full nav + lst.cursor = row + lst._scroll_cursor_into_view() + + # Has execution actually left the decompiled function? Ask the map the + # trail painting keeps, which is keyed to what the decompiler currently + # HOLDS. _split_range comes from the guarded async path and lags, so a + # stale one made every step look like a function change: the decompiler + # bounced main -> PLT stub -> main, each bounce costing a synchronous + # 769-line map fetch on the UI thread. + span = self.trail_span + inside = (pc in self.trail_line_of + or (span is not None and span[0] <= pc <= span[1])) + if not inside: + self.pending_line = pc + app._resync_decomp_async(pc) + return True + self.place_decomp_at(pc) + app._sync_split(app._active) + return True + + def place_decomp_at(self, pc: int) -> None: + """Move the pseudocode cursor to the line covering ``pc``. + + Uses the map the trail painting already keeps (keyed to the decompiler's + CURRENTLY loaded function), not the split view's _split_ea2line. That one + is refreshed by a guarded async path — it drops a result if _cur moved + while it was in flight — and a burst of steps moves _cur constantly, so + during stepping it is frequently a map of the function you just left. + """ + app = self.app + dec = app.query_one(_views().DecompView) + line = None + if self.trail_map_ea == dec.loaded_ea and self.trail_line_of: + # EXACT match only. The decompiler doesn't attribute every + # instruction to a line (about half of main's aren't), and the + # tempting fallback — the nearest mapped instruction at or before + # the pc — is unsound: C lines are not monotonic in address, so + # 0x24a8 early in main resolved to line 708, "sub_2040();", near the + # end. A cursor that jumps to an unrelated statement is worse than + # one that waits; the trail still marks where we are. + line = self.trail_line_of.get(pc) + if line is None: + line = app._split_ea2line.get(pc) + if line is None: + line = dec.line_for_ea(pc) + if line is not None: + dec.goto(line, dec.cursor_x) + + # -- painting ----------------------------------------------------------- # + def paint_trail(self) -> None: + """Push the execution trail into the code views. + + Recomputed per seek rather than per repaint: it's ~200 lookups, and a + repaint happens far more often than a step. + """ + app = self.app + M = _views() + t = self.trace + if t is None: + return + hx = app._try_view(M.HexView) # None until it's mounted + if hx is not None: + hx.trace, hx.trace_idx = t, self.t + if hx.display: + hx.refresh() + trail = t.trail(self.t) + lst = app._try_view(M.ListingView) + if lst is not None: + lst.trail = trail + lst.refresh() + self.paint_trail_decomp(trail) + + def paint_trail_decomp(self, trail: dict) -> None: + """Map the instruction trail onto pseudocode lines. + + This is the thing Tenet can't do: it paints disassembly, because that's + where a trace's addresses live. We already have decomp_map (built for + the split view) saying which instructions each pseudocode line covers, + so the same trail lands on C. + + A line covers many instructions, so it takes the strongest kind present: + 'now' wins over 'past' wins over 'future' — if the instruction you are + standing on is part of this line, this line is where you are. + """ + app = self.app + dec = app._try_view(_views().DecompView) + if dec is None: + return + ea = dec.loaded_ea + if not dec.display or ea is None or app.program is None: + dec.trail = {} + return + if self.trail_map_ea != ea: + # One index, built once per decompiled function and shared with the + # split view (_apply_split_map fills the same fields). decomp_map is + # an RPC and stepping is interactive, so paying it per keystroke — + # or twice, once for each of two parallel maps — would be felt. + try: + app._apply_split_map(ea, app.program.decomp_map(ea)) + except Exception as e: # noqa: BLE001 + # The pseudocode simply stops being painted with the trail, with + # nothing on screen to say why. + diag.note(f"trail: decomp_map({ea:#x})", e) + self.trail_map, self.trail_map_ea = [], ea + self.trail_line_of, self.trail_eas = {}, [] + self.trail_span = None + rank = {"future": 0, "past": 1, "now": 2} + lines: dict[int, str] = {} + for i, eas in enumerate(self.trail_map or []): + best = None + for a in eas: + k = trail.get(a) + if k is not None and (best is None or rank[k] > rank[best]): + best = k + if best is not None: + lines[i] = best + dec.trail = lines + dec.refresh() + pend, self.pending_line = self.pending_line, None + if pend is not None and app._split: + # The function was still decompiling when the step happened; land + # now that its line map exists. + self.place_decomp_at(pend) + app._sync_split(app._active) + + def adopt_map(self, ea: int, m: list, ea2line: dict, span) -> None: + """Take the line map the split view just built. + + ONE index, shared with the split view: the trace path used to keep a + parallel copy of exactly this, fetched separately and keyed differently, + which is how the two ended up describing different functions. + """ + self.trail_map, self.trail_map_ea = m, ea + self.trail_line_of = dict(ea2line) + self.trail_eas = sorted(self.trail_line_of) + self.trail_span = span + + # -- stepping ----------------------------------------------------------- # + def step(self, delta: int) -> None: + if not self._need_trace(): + return + self.seek(self.t + delta) + + def step_over(self, direction: int) -> None: + """Step over a call by following the stack pointer. + + A call pushes, so the callee runs with SP BELOW where we started; + stepping until SP comes back up lands after the call returns. Cheaper + and more robust than recognising call instructions per architecture, + which is what the mode makes it: if this instruction doesn't call + anything, SP is already >= the start and it degenerates to one step. + """ + if not self._need_trace(): + return + t = self.trace + sp_name = "rsp" if "rsp" in t.reg_at else ("esp" if "esp" in t.reg_at else "sp") + sp0 = t.register(sp_name, self.t) + i = self.t + direction + limit = 200000 # a runaway search must not hang the UI + while 0 <= i < t.length and limit > 0: + sp = t.register(sp_name, i) + if sp0 is None or sp is None or sp >= sp0: + break + i += direction + limit -= 1 + self.seek(max(0, min(i, t.length - 1))) + + def seek_hit(self, direction: int) -> None: + """Seek to the next/previous time the focused view's subject was touched. + + Two different questions with one pair of keys, because the answer to + "which thing?" is already on screen: in a code view it's the instruction + under the cursor ("when else did this run?"), in hex it's the byte under + the cursor ("who else touched this?"). + """ + app = self.app + M = _views() + if not self._need_trace(): + return + t = self.trace + if app.is_hex: + hx = app._try_view(M.HexView) + va = hx.cursor_va() if hx is not None else None + if va is None: + return + stamps = t.memory_accesses(va, 1) + what = f"access to {va:#x}" + else: + view = app._active_code_view() + if isinstance(view, M.DecompView): + # A C line is not one address, so ask about the whole statement: + # "when else did this line run?" is the question, and it's the + # union of its instructions' executions. Falling back to the + # line's single /*ea*/ marker would answer a narrower question + # and often no question at all, since most lines have no marker. + line = view.cursor + eas = [] + if (self.trail_map_ea == view.loaded_ea + and 0 <= line < len(self.trail_map or [])): + eas = list(self.trail_map[line]) + if not eas: + one = view._line_ea(line) + eas = [one] if one is not None else [] + if not eas: + app._status("this line has no instructions to seek on", + priority=True) + return + stamps = sorted({x for e in eas for x in t.executions(e)}) + what = f"execution of C line {line + 1}" + else: + ea = view._cursor_ea() if view is not None else None + if ea is None: + app._status("no address on this line", priority=True) + return + stamps = list(t.executions(ea)) + what = f"execution of {ea:#x}" + if not stamps: + app._status(f"no {what} in this trace", priority=True) + return + if direction > 0: + i = bisect.bisect_right(stamps, self.t) + else: + i = bisect.bisect_left(stamps, self.t) - 1 + if not (0 <= i < len(stamps)): + edge = "last" if direction > 0 else "first" + app._status(f"already at the {edge} {what} " + f"({len(stamps)} in the trace)", priority=True) + return + self.seek(stamps[i]) + app._status(f"{what}: {i + 1} of {len(stamps)} @ t={stamps[i]:,}", + priority=True) + + def seek_reg_write(self) -> None: + """W: which instruction set each register to its current value.""" + app = self.app + if not self._need_trace(): + return + t = self.trace + rows = [] + for name in t.registers: + v = t.register(name, self.t) + if v is None: + continue + rows.append((name, v, t.last_write(name, self.t), + t.next_write(name, self.t))) + if rows: + app.push_screen(_views().RegWriteScreen(rows, self.t), + self._on_reg_write_chosen) + + def _on_reg_write_chosen(self, idx) -> None: # type: ignore[no-untyped-def] + if idx is not None: + self.seek(int(idx)) diff --git a/idatui/worker.py b/idatui/worker.py deleted file mode 100644 index a4e3509..0000000 --- a/idatui/worker.py +++ /dev/null @@ -1,233 +0,0 @@ -"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor. - -Opens ONE database in-process (on the main thread, as idalib requires) and -serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed -pickle. Same tool implementations as the MCP path (we call -``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are -byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no -supervisor / HTTP / JSON / 50KB-truncation machinery. - - python -m idatui.worker <sock_path> <binary_path> - -The socket only appears once the database is open + analyzed, so a client can -poll ``connect()`` to know when the worker is ready. Requests are served -serially on the main thread (idalib is single-threaded; every tool runs inline -through its own execute_sync, which is a no-op on the main thread). - -Protocol (both directions length-prefixed: 4-byte big-endian len + pickle): - request = (tool_name: str, kwargs: dict) - response = (ok: bool, result_or_error) - tool_name == "__shutdown__" ends the worker. -""" -from __future__ import annotations - -import os -import pickle -import socket -import struct -import sys -import uuid - - -# --------------------------------------------------------------------------- # -# framing -# --------------------------------------------------------------------------- # -def _recvn(sock: socket.socket, n: int) -> bytes | None: - buf = bytearray() - while len(buf) < n: - chunk = sock.recv(n - len(buf)) - if not chunk: - return None - buf += chunk - return bytes(buf) - - -def send(sock: socket.socket, obj) -> None: - data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) - sock.sendall(struct.pack(">I", len(data)) + data) - - -def recv(sock: socket.socket): - hdr = _recvn(sock, 4) - if hdr is None: - return None - (n,) = struct.unpack(">I", hdr) - body = _recvn(sock, n) - return None if body is None else pickle.loads(body) - - -# --------------------------------------------------------------------------- # -# worker -# --------------------------------------------------------------------------- # -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 - (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__))) - patch = os.path.join(repo, "server", "patch_server.py") - if not os.path.exists(patch): - return - try: - spec = importlib.util.spec_from_file_location("_idatui_patch", patch) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types - mod.main() - except Exception as e: # noqa: BLE001 -- tools may already be present - sys.stderr.write(f"idatui: tool injection skipped: {e}\n") - - -def _has_database(binpath: str) -> bool: - """Whether IDA already has a database for ``binpath``. - - IDA names it ``<file>.i64`` (keeping the extension), but a database made - from ``foo.bin`` can also appear as ``foo.i64`` depending on how it was - created — check both, because guessing wrong here means re-passing load - switches to an existing database, which fails the open. - """ - return (os.path.exists(binpath + ".i64") - or os.path.exists(os.path.splitext(binpath)[0] + ".i64")) - - -def _open_and_register(binpath: str, load_args: str = ""): - """Open the DB (main thread) then import ida-pro-mcp so every @tool registers - against this live database. Returns (tools_dict, module_name, save_fn). - - ``load_args`` is passed to IDA as command-line switches, which is the only - way to tell it how to read a headerless blob: a raw firmware image has no - format to detect, so without ``-p<processor>`` it loads as metapc at 0 and - finds nothing. Ignored once a database exists — the .i64 already records how - it was loaded, and re-passing conflicting switches is how you corrupt one. - """ - _ensure_tools_injected() # before any ida_pro_mcp import - import idapro - idapro.enable_console_messages(False) - args = load_args or None - if args and _has_database(binpath): - # The .i64 already records how this image was loaded. Passing the - # switches again on reopen makes IDA fail outright (rc != 0) — the load - # options belong to the FIRST open only. - args = None - if idapro.open_database(binpath, run_auto_analysis=True, - args=args): # nonzero == failure - if args: - # With load switches in play they are the likeliest culprit by far: - # IDA refuses an unknown -p name with no diagnostic of its own, so - # saying "the database is locked" here sends people hunting a - # problem they don't have. - raise RuntimeError( - f"failed to open {binpath} with load options {args!r}: IDA " - f"rejected them \u2014 an unknown processor name is the usual " - f"cause (see tools/verify_procs.py for the valid ones)") - raise RuntimeError( - f"failed to open {binpath}: the .i64 is likely held by a running " - f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash " - f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)") - import ida_auto - ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp) - - # importing the package registers all api_*/patched tools against MCP_SERVER - from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 - - import ida_nalt - module = os.path.basename(ida_nalt.get_root_filename() or binpath) - - def save(): - import idc - try: - idc.save_database(idc.get_idb_path(), 0) - except Exception: # noqa: BLE001 - import ida_loader, ida_pro # noqa: WPS433 - ida_loader.save_database(idc.get_idb_path(), 0) - - return MCP_SERVER.tools.methods, module, save - - -def serve(sockpath: str, binpath: str, load_args: str = "") -> None: - tools, module, save = _open_and_register(binpath, load_args) - sid = uuid.uuid4().hex[:8] - - def dispatch(name: str, args: dict): - args = dict(args) - args.pop("database", None) # single-DB worker: no session routing - # session-management shims (were the supervisor's job): - if name in ("idb_open",): - return {"success": True, - "session": {"session_id": sid, "module": module, - "input_path": binpath}} - if name in ("idb_save", "save"): - save() - return {"success": True} - if name in ("server_health", "ping", "health", "state"): - return {"module": module, "ok": True, "session_id": sid} - if name in ("idb_list",): - return {"sessions": [{"session_id": sid, "module": module, - "input_path": binpath}]} - fn = tools.get(name) - if fn is None: - raise KeyError(f"unknown tool: {name!r}") - result = fn(**args) - # Match the MCP server's structuredContent: a dict passes through, any - # other return (list/scalar) is wrapped as {"result": ...}. domain.py - # parses that exact shape (e.g. lookup_funcs -> payload["result"]). - return result if isinstance(result, dict) else {"result": result} - - try: - os.unlink(sockpath) - except OSError: - pass - srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - srv.bind(sockpath) - srv.listen(8) - try: - while True: - conn, _ = srv.accept() - try: - while True: - req = recv(conn) - if req is None: - break - name, args = req - if name == "__shutdown__": - return - try: - send(conn, (True, dispatch(name, args))) - except Exception as e: # noqa: BLE001 -- report, keep serving - send(conn, (False, f"{type(e).__name__}: {e}")) - except (ConnectionError, OSError): - pass - finally: - conn.close() - finally: - try: - import idapro - idapro.close_database(save=False) - except Exception: # noqa: BLE001 - pass - try: - os.unlink(sockpath) - except OSError: - pass - - -def main(argv=None) -> None: - argv = argv if argv is not None else sys.argv[1:] - if len(argv) < 2: - sys.stderr.write( - "usage: python -m idatui.worker <sock> <binary> [ida-load-args]\n") - raise SystemExit(2) - try: - serve(argv[0], argv[1], argv[2] if len(argv) > 2 else "") - except SystemExit: - raise - except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1 - import traceback - sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n") - traceback.print_exc() - sys.stderr.flush() - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/idatui/worker_client.py b/idatui/worker_client.py deleted file mode 100644 index 79a6db9..0000000 --- a/idatui/worker_client.py +++ /dev/null @@ -1,234 +0,0 @@ -"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own -idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's -HTTP/JSON transport. - -It exposes exactly the surface the app/domain use on the client -(``call``/``call_envelope``/``connect``/``set_db``/``resolve_db``/ -``list_sessions``/``health``/``keepalive``/``close``) and returns byte-identical -payloads (the worker calls the same tool functions), so ``domain.py`` and the -app are unchanged — you just construct a WorkerClient instead of an IDAClient. - -Concurrency: the app fires calls from several worker threads over one client; -the worker is single-threaded, so calls are serialized under a lock (the worker -processes one tool at a time anyway — and at ~50us/call that's free). -""" -from __future__ import annotations - -import os -import socket -import subprocess -import sys -import threading -import time -import uuid -from typing import Any - -from .errors import IDAToolError, IDAConnectionError, Session -from .worker import recv as _recv -from .worker import send as _send - -_WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") -_worker_python_cache: str | None = None - - -def _find_worker_python() -> str: - """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily - the TUI's python. On a typical box the TUI runs under a venv that has textual - + idalib but not ida_pro_mcp, while the system python has idalib + - ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" - global _worker_python_cache - if _worker_python_cache: - return _worker_python_cache - override = os.environ.get("IDATUI_WORKER_PYTHON") - candidates = [override] if override else [] - candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] - for py in candidates: - if not py or not os.path.exists(py): - continue - try: - r = subprocess.run([py, "-c", "import ida_pro_mcp"], - capture_output=True, timeout=30) - if r.returncode == 0: - _worker_python_cache = py - return py - except Exception: # noqa: BLE001 - continue - return sys.executable # last resort; the worker will report the real error - - -class _NoopKeepAlive: - """The worker is ours and never idles out, so keepalive is a no-op.""" - - def __init__(self) -> None: - self.beats = self.failures = 0 - - def start(self): - return self - - def stop(self) -> None: - pass - - -class WorkerClient: - def __init__(self, binary_path: str, *, ttl: int = 0, - python: str | None = None, load_args: str = "") -> None: - self._bin = os.path.abspath(os.path.expanduser(binary_path)) - self._load_args = load_args or "" # IDA switches for a headerless blob - self._python = python or _find_worker_python() - tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" - self._sock_path = f"/tmp/idatui-worker-{tag}.sock" - self._log_path = f"/tmp/idatui-worker-{tag}.log" - self._proc: subprocess.Popen | None = None - self._sock: socket.socket | None = None - self._sid = uuid.uuid4().hex[:8] - self._lock = threading.Lock() # serialize socket use - self._spawn_lock = threading.Lock() - - # -- lifecycle --------------------------------------------------------- # - def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": - """Spawn the worker (opens + analyzes the DB) and connect once ready.""" - with self._spawn_lock: - if self._sock is not None: - return self - if self._proc is None or self._proc.poll() is not None: - # run worker.py as a SCRIPT (not -m idatui.worker) so we don't - # import the textual-dependent idatui package __init__ under the - # IDA python, which usually has no textual. - argv = [self._python, _WORKER_PY, self._sock_path, self._bin] - if self._load_args: - argv.append(self._load_args) - self._proc = subprocess.Popen( - argv, - stdout=open(self._log_path, "wb"), - stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, - ) - deadline = time.time() + timeout - t0 = time.time() - while time.time() < deadline: - try: - s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - s.connect(self._sock_path) - self._sock = s - return self - except OSError: - if self._proc.poll() is not None: - raise IDAConnectionError( - f"worker exited (code {self._proc.returncode}): " - f"{self._log_tail()} [full log: {self._log_path}]") - if progress: - progress(f"auto-analyzing {os.path.basename(self._bin)}… " - f"({int(time.time() - t0)}s)") - time.sleep(0.2) - raise IDAConnectionError("worker did not become ready in time") - - @property - def pid(self) -> int | None: - """The worker process id (for memory accounting), or None if not spawned.""" - return self._proc.pid if self._proc is not None else None - - def close(self, grace: float = 20.0) -> None: - """Shut the worker down cleanly. - - After ``__shutdown__`` the worker still has to ``close_database()``, which - re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. - Signalling it before that finishes is what leaves databases wedged, so - wait out the grace period first and only escalate if it really is stuck. - """ - with self._lock: - s = self._sock - self._sock = None - if s is not None: - try: - _send(s, ("__shutdown__", {})) - except Exception: # noqa: BLE001 - pass - try: - s.close() - except Exception: # noqa: BLE001 - pass - if self._proc is not None: - try: - self._proc.wait(timeout=grace) # let it close the DB properly - except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck - try: - self._proc.terminate() - self._proc.wait(timeout=5) - except Exception: # noqa: BLE001 - try: - self._proc.kill() - except Exception: # noqa: BLE001 - pass - - # -- the call surface -------------------------------------------------- # - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - if self._sock is None: - self.connect() - with self._lock: - s = self._sock - if s is None: - raise IDAConnectionError("worker connection is closed") - try: - _send(s, (tool, args)) - reply = _recv(s) - except (OSError, ConnectionError) as e: - self._sock = None - raise IDAConnectionError(f"worker transport failed: {e}") from e - if reply is None: - self._sock = None - raise IDAConnectionError("worker closed the connection") - ok, payload = reply - if not ok: - raise IDAToolError(tool, str(payload)) - return payload - - def call_envelope(self, tool: str, *, timeout: float | None = None, - **args) -> dict: - # domain.decompile() reads result.structuredContent — mirror that shape. - return {"result": {"structuredContent": self.call(tool, timeout=timeout, - **args)}} - - # -- session shims (single-DB worker) --------------------------------- # - def set_db(self, db: str | None) -> None: - if db: - self._sid = db - - def resolve_db(self) -> str: - return self._sid - - def list_sessions(self) -> list[Session]: - return [Session(session_id=self._sid, - filename=os.path.basename(self._bin), - input_path=self._bin, is_active=True)] - - def health(self) -> dict: - try: - return self.call("server_health") - except IDAToolError: - return {"module": os.path.basename(self._bin), "ok": True} - - def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: - return _NoopKeepAlive() - - def _log_tail(self, n: int = 400) -> str: - """Last meaningful line(s) of the worker log (skip IDA's licence banner), - so a startup crash surfaces the real cause instead of just 'code 1'.""" - try: - with open(self._log_path, encoding="utf-8", errors="replace") as f: - lines = [ln.strip() for ln in f if ln.strip()] - except OSError: - return "(no worker log)" - # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash - for ln in reversed(lines): - if ln.startswith("WORKER-FATAL:"): - return ln[len("WORKER-FATAL:"):].strip()[-n:] - skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") - meaningful = [ln for ln in lines - if not any(s in ln.lower() for s in skip)] - return " | ".join((meaningful or lines)[-3:])[-n:] - - # context manager parity with IDAClient - def __enter__(self) -> "WorkerClient": - return self.connect() - - def __exit__(self, *exc) -> None: - self.close() |