diff options
Diffstat (limited to 'idatui')
| -rw-r--r-- | idatui/__init__.py | 9 | ||||
| -rw-r--r-- | idatui/__main__.py | 6 | ||||
| -rw-r--r-- | idatui/app.py | 6028 | ||||
| -rw-r--r-- | idatui/client.py | 678 | ||||
| -rw-r--r-- | idatui/domain.py | 1060 | ||||
| -rw-r--r-- | idatui/drive.py | 92 | ||||
| -rw-r--r-- | idatui/errors.py | 76 | ||||
| -rw-r--r-- | idatui/formats.py | 143 | ||||
| -rw-r--r-- | idatui/graph.py | 715 | ||||
| -rw-r--r-- | idatui/highlight.py | 31 | ||||
| -rw-r--r-- | idatui/index.py | 206 | ||||
| -rw-r--r-- | idatui/kittygfx.py | 271 | ||||
| -rw-r--r-- | idatui/launch.py | 181 | ||||
| -rw-r--r-- | idatui/pane.py | 570 | ||||
| -rw-r--r-- | idatui/pool.py | 235 | ||||
| -rw-r--r-- | idatui/project.py | 353 | ||||
| -rw-r--r-- | idatui/rpc.py | 589 | ||||
| -rw-r--r-- | idatui/rpcclient.py | 10 | ||||
| -rw-r--r-- | idatui/trace.py | 495 | ||||
| -rw-r--r-- | idatui/tui.py | 45 | ||||
| -rw-r--r-- | idatui/worker.py | 233 | ||||
| -rw-r--r-- | idatui/worker_client.py | 234 |
22 files changed, 10927 insertions, 1333 deletions
diff --git a/idatui/__init__.py b/idatui/__init__.py index a52ba8b..2cbdde8 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -1,7 +1,7 @@ -"""idatui — a minimal keyboard-first TUI for IDA Pro over ida-pro-mcp (idalib).""" +"""idatui — a minimal keyboard-first TUI for IDA Pro, driving idalib via a +private unix-socket worker (idatui.worker / WorkerClient).""" -from .client import ( - IDAClient, +from .errors import ( IDAError, IDAConnectionError, IDATimeoutError, @@ -10,7 +10,6 @@ from .client import ( IDAToolError, IDASessionError, Session, - KeepAlive, ) from .domain import ( Program, @@ -36,7 +35,6 @@ __all__ = [ "Decompilation", "LIST_PAGE", "DISASM_BLOCK", - "IDAClient", "IDAError", "IDAConnectionError", "IDATimeoutError", @@ -45,5 +43,4 @@ __all__ = [ "IDAToolError", "IDASessionError", "Session", - "KeepAlive", ] diff --git a/idatui/__main__.py b/idatui/__main__.py index 29f5cf4..0090ace 100644 --- a/idatui/__main__.py +++ b/idatui/__main__.py @@ -1,6 +1,6 @@ -"""``python -m idatui`` -> client self-check (until the TUI app lands).""" +"""``python -m idatui`` -> the one-shot launcher (open a binary in the TUI).""" import sys -from .client import _main +from .launch import main -raise SystemExit(_main(sys.argv[1:])) +raise SystemExit(main(sys.argv[1:])) diff --git a/idatui/app.py b/idatui/app.py index 1db32ff..fc19fa5 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -20,46 +20,108 @@ import asyncio import os import re import subprocess -from dataclasses import dataclass +import time +from dataclasses import dataclass, field +from rich.align import Align from rich.segment import Segment from rich.style import Style from rich.text import Text from textual import work from textual.app import App, ComposeResult from textual.binding import Binding -from textual.containers import Horizontal, Vertical +from textual.command import DiscoveryHit, Hit, Provider +from textual.containers import Grid, Horizontal, Vertical, VerticalScroll from textual.geometry import Region, Size from textual.message import Message from textual.reactive import reactive +from textual.theme import Theme from textual.screen import ModalScreen from textual.scroll_view import ScrollView from textual.strip import Strip from textual.widgets import ( - DataTable, Footer, Input, OptionList, Static, TextArea, + DataTable, Input, OptionList, Static, TextArea, ) from textual.widgets.option_list import Option +from . import graph +from . import kittygfx from .highlight import highlight_c -from .client import IDAClient, IDAToolError -from .domain import DisasmModel, Func, Program, Struct +from .errors import IDAToolError, IDAConnectionError +from .worker_client import WorkerClient +from .domain import DisasmModel, Func, Head, ListingModel, Program, Struct # Styles for the disassembly listing. -_S_ADDR = Style(color="grey58") -_S_LABEL = Style(color="yellow", bold=True) -_S_INSN = Style(color="white") -_S_MNEM = Style(color="cyan") -_S_OPBYTES = Style(color="grey50") # raw opcode bytes column -_S_CURSOR = Style(bgcolor="grey30") -_S_DIM = Style(color="grey42", italic=True) +_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") +_S_UNK = Style(color="#7c8b9e", italic=True) # undefined bytes in the flat listing +_S_MEMBER = Style(color="#93aee0") +_S_SEP = Style(color="#5e6875") # function boundary separators / banners +_S_FUNCHDR = Style(color="#7aa2f7", bold=True) # 'name proc'/'endp' headers + +_LST_INDENT = " " # one depth level: function names sit at level 0, code at 1 +_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 + +# 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({ + "db", "dw", "dd", "dq", "dt", "byte", "word", "dword", "qword", "tbyte", + "offset", "short", "near", "far", "ptr", "dup", "cs", "ds", "es", "fs", + "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="#b58900", color="black") # the current match -_S_NAME_MATCH = Style(bgcolor="#b58900", color="black") # filter match in a name -_S_WORD = Style(bgcolor="#264f78") # identifier under the cursor +_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="grey37") # pseudocode line-number gutter -_S_LINENO_CUR = Style(color="grey66", bold=True) # gutter on the cursor line +_S_LINENO = Style(color="#626c7a") # pseudocode line-number gutter +_S_LINENO_CUR = Style(color="#c3cad3", bold=True) # gutter on the cursor line +_S_DECOMP_SPIN = Style(color="#d0a215", bold=True) # 'decompiling' spinner glyph +_S_DECOMP_WAIT = Style(color="#7c8b9e", italic=True) # 'decompiling' label +_S_DECOMP_DOTS = Style(color="#626c7a") # trailing ellipsis +_S_LINK = Style(bgcolor="#233044") # split view: rows linked to the other pane's cursor # Hex-Rays appends a `/*0xEA*/` address marker to each pseudocode line (we fetch # with include_addresses so we have a per-line anchor). Matched here to extract @@ -69,6 +131,48 @@ _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.""" + + label: str + program: object | None = None + func_index: object | None = None + nav: list = field(default_factory=list) + cur: object | None = None + active: str = "listing" + split: bool = False + filter_term: str = "" + dirty: bool = False + + +@dataclass +class ViewAnchor: + """Where the user is looking, expressed in ADDRESSES. + + Every path that rebuilds a model must round-trip through this. Row indices + do NOT survive a rebuild: defining code collapses four undefined byte rows + into one instruction row, undefining does the reverse, and a rename can add + or remove banner rows above a function. Anything that remembers an index + puts the user somewhere else afterwards, which reads as "the edit jumped my + screen" or, worse, "the edit didn't apply". + + ``flash`` travels with it because the same rebuild also decides what the + status bar says: the reload writes its own status when it lands, so an edit + that doesn't hand its message over here gets silently overwritten. + """ + + view: str = "listing" + ea: int | None = None # cursor address + 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 class NavEntry: ea: int name: str @@ -79,6 +183,8 @@ class NavEntry: dec_cursor_x: int = 0 # pseudocode column dec_scroll_y: int = -1 # pseudocode viewport top (-1 = derive) dec_scroll_x: int = 0 # pseudocode horizontal scroll + is_region: bool = False # not inside a function (flat listing view) + view: str = "listing" # which code view to restore this entry in class SearchRequested(Message): @@ -132,6 +238,35 @@ class RetypeRequested(Message): self.name = name +class EditItemRequested(Message): + """The disasm view asks to change item structure (IDA c/d/u/p) at the line + under the cursor. ``kind`` is 'code' | 'func' | 'undef'.""" + + def __init__(self, view, kind: str) -> None: # type: ignore[no-untyped-def] + super().__init__() + self.view = view + self.kind = kind + + +class MakeDataRequested(Message): + """The listing view asks to define typed data (IDA 'd') at the current head.""" + + def __init__(self, view) -> None: # type: ignore[no-untyped-def] + super().__init__() + 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).""" @@ -206,6 +341,27 @@ def _word_bounds(text: str, x: int) -> tuple[int, int]: return (s, e) +def _word_occurrences(text: str, word: str) -> list[tuple[int, int]]: + """[start, end) spans of every WHOLE-word occurrence of ``word`` in ``text`` + (for highlight-all-occurrences of the token under the cursor).""" + if not word or not text: + return [] + isw = lambda c: c.isalnum() or c == "_" # noqa: E731 + out: list[tuple[int, int]] = [] + n = len(word) + i = 0 + while True: + j = text.find(word, i) + if j < 0: + break + before_ok = j == 0 or not isw(text[j - 1]) + after_ok = j + n >= len(text) or not isw(text[j + n]) + if before_ok and after_ok: + out.append((j, j + n)) + i = j + n + return out + + def _overlay_over(strip: Strip, ranges: list[tuple[int, int]], style: Style) -> Strip: """Like _overlay_ranges but ``style`` OVERRIDES the existing cell styles (used for the cursor cell / word, which must win over the line background).""" @@ -255,9 +411,21 @@ class ColumnCursor: Binding("dollar_sign", "col_end", "eol", show=False), ] + _hl_word: str | None = None # token to highlight across all visible lines + def _line_plain(self, idx: int) -> str | None: raise NotImplementedError + def _refresh_hl(self) -> None: + """Recompute the highlight-all token from the word under the cursor; if it + changed, repaint the whole viewport (occurrences elsewhere changed).""" + w = self.word_under_cursor() + if not (w and len(w) >= 2 and (w[0].isalpha() or w[0] == "_")): + w = None + if w != self._hl_word: + self._hl_word = w + self.refresh() + def _hscroll(self) -> None: pass @@ -271,6 +439,7 @@ class ColumnCursor: self.cursor_x = max(0, min(max(len(plain) - 1, 0), self.cursor_x + dx)) self._hscroll() _refresh_lines(self, self.cursor) + self._refresh_hl() def action_col_left(self) -> None: self._move_x(-1) @@ -282,12 +451,14 @@ class ColumnCursor: self.cursor_x = 0 self._hscroll() _refresh_lines(self, self.cursor) + self._refresh_hl() def action_col_end(self) -> None: plain = self._line_plain(self.cursor) or "" self.cursor_x = max(len(plain) - 1, 0) self._hscroll() _refresh_lines(self, self.cursor) + self._refresh_hl() def action_col_word(self, direction: int) -> None: plain = self._line_plain(self.cursor) or "" @@ -311,6 +482,7 @@ class ColumnCursor: self.cursor_x = max(i, 0) self._hscroll() _refresh_lines(self, self.cursor) + self._refresh_hl() def word_under_cursor(self) -> str | None: plain = self._line_plain(self.cursor) @@ -338,6 +510,7 @@ class ColumnCursor: self._scroll_cursor_into_view() self._hscroll() self.refresh() + self._refresh_hl() self._after_cursor_move() def on_click(self, event) -> None: # type: ignore[no-untyped-def] @@ -549,6 +722,7 @@ class SearchMixin: self.scroll_to(y=max(self.cursor - self._visible_height() // 2, 0), animate=False) self._hscroll() # bring the match column into horizontal view self.refresh() + self._refresh_hl() def clear_search(self) -> None: self._term = "" @@ -568,8 +742,15 @@ class SearchMixin: # --------------------------------------------------------------------------- # # Virtualized disassembly view # --------------------------------------------------------------------------- # -class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True): - """A line-virtualized disassembly listing for a single function.""" +# --------------------------------------------------------------------------- # +# 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 ``DisasmModel`` + which is bounded to one function. Backed by ``ListingModel`` (the ``heads`` + server tool). Used for non-function regions and raw segment browsing. + """ BINDINGS = [ Binding("j,down", "cursor_down", "Down", show=False), @@ -578,10 +759,23 @@ class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True 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("tab,shift+tab", "app.toggle_view", "Pseudocode", priority=True), + Binding("home", "col_home", "bol", show=False), + Binding("shift+home", "col_insn_home", "insn start", show=False), + Binding("end", "col_end", "eol", show=False), + Binding("G,ctrl+end", "goto_bottom", "Bottom", show=False), + Binding("ctrl+home", "goto_top", "Top", show=False), + # `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, @@ -591,125 +785,272 @@ class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True cursor_x = reactive(0, repaint=False) class CursorMoved(Message): - """Posted when the disasm cursor moves; carries the instruction ea.""" + """Posted when the listing cursor moves; carries the head ea.""" def __init__(self, index: int, ea: int | None) -> None: super().__init__() self.index = index self.ea = ea + class Scrolled(Message): + """Posted when the viewport scrolls (wheel/scrollbar) — the cursor need + not have moved, so the split view can still follow along.""" + def __init__(self) -> None: super().__init__() - self.model: DisasmModel | None = None + self.model: ListingModel | None = None self.total = 0 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]]] = {} - 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) + self._op_mode = 1 # opcode column: 0=off, 1=limited, 2=full ('o' cycles) + self._op_w = 0 # char width of the hex-bytes field (excl. gap) + 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: + return self.model.get(idx) if self.model is not None else None + + @staticmethod + def _name_prefix(h: Head) -> str: + return f"{h.name} " if h.name else "" - 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: + 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.""" + 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) + + @staticmethod + def _span_segments(h: Head, fallback: Style): + """Segments for a row's disassembly text. + + Uses IDA's own token classification when the worker supplied it; falls + back to the old mnemonic/rest split so an older worker (or a row whose + spans didn't match the text) still renders. + """ + 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 + cursor/search offsets line up.""" + if self._op_mode == 0 or self._op_w <= 0: return "" - raw = line.raw or b"" - return " ".join(f"{b:02X}" for b in raw).ljust(self._op_w) + " " + return self._op_bytes_text(h).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: + h = self._head(idx) + if h is None: return None - s = f"{line.ea:08X} " + self._op_field(line) - if line.label: - s += f"{line.label}: " - return s + line.text + # Function headers and code labels sit at depth 0 (with the address); + # everything else is indented one level (opcode+text included). + if h.kind in ("funchdr", "label"): + return f"{h.ea:08X} {h.text}" + base = f"{h.ea:08X} " + _LST_INDENT + if h.kind == "sep": + return base + h.text + extra = _LST_INDENT if h.kind == "member" else "" + return base + self._op_field(h) + extra + self._name_prefix(h) + h.text + + 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 + prefix so the column lines up with what's rendered.""" + h = self._head(idx) + if h is None: + return 0 + if h.kind in ("funchdr", "label"): + return len(f"{h.ea:08X} ") + base = f"{h.ea:08X} " + _LST_INDENT + if h.kind == "sep": + return len(base) + extra = _LST_INDENT if h.kind == "member" else "" + return len(base + self._op_field(h) + extra + self._name_prefix(h)) + + def action_col_insn_home(self) -> None: + """shift+home: jump to the start of the instruction text, skipping the + address + opcode-bytes gutter (home/0 still go to the true line start).""" + self.cursor_x = self._insn_col(self.cursor) + self._hscroll() + _refresh_lines(self, self.cursor) + self._refresh_hl() # -- public API -------------------------------------------------------- # - def load(self, model: DisasmModel, name: str, cursor: int = 0, - cursor_x: int = 0, scroll_y: int | None = None) -> None: + 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 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._pending_focus = focus # token to land the cursor column on + self._pending_op = None 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 - while off < total: - lines = model.lines(off, DisasmModel.BLOCK, prefetch=False) - if not lines: - break - texts.extend(self._fmt(ln) for ln in lines) - off += len(lines) - self._search_texts = texts - self.app.call_from_thread(done) - - @work(thread=True, exclusive=True, group="disasm-prime") + @work(thread=True, exclusive=True, group="listing-prime") def _prime(self) -> None: model = self.model if model is None: return - total = model.total() + # Load just enough to render the viewport around the cursor, so the + # listing appears immediately even on a huge segment; the rest streams + # in via _grow. (load_all here would blank the pane for seconds.) 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) + model.ensure(self.cursor + height + 2 * ListingModel.PAGE) + self.app.call_from_thread(self._on_primed, len(model), model.complete) + if not model.complete: + self._grow() - def _on_primed(self, total: int) -> None: + def _on_primed(self, total: int, complete: bool) -> None: + if self.model is None: + return 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 + self.cursor = max(0, min(self.cursor, max(total - 1, 0))) + self._update_op_w() # finalize column layout before locating the token + # Land the cursor on the requested token's column (e.g. the ref an xref + # jump targets), now that the row is loaded and renderable. + if self._pending_focus: + plain = self._line_plain(self.cursor) + if plain: + occ = _word_occurrences(plain, self._pending_focus) + 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))) else: self._scroll_cursor_into_view() self._pending_scroll_y = None + self._hscroll() # bring the cursor column into horizontal view + self.refresh() + self.post_message(ListingView.CursorMoved(self.cursor, self._cursor_ea())) + + def _update_op_w(self) -> bool: + """Recompute the opcode-column width from the widest code head seen (capped + in 'limited' mode). Returns True if it changed.""" + w = 0 + if self.model is not None and self._op_mode != 0: + mx = self.model.max_raw_len() + if mx > 0: + if self._op_mode == 1: # limited: cap bytes, +1 for the ellipsis + shown = min(mx, _OP_LIMIT) + w = shown * 3 - 1 + (1 if mx > _OP_LIMIT else 0) + else: # full + w = mx * 3 - 1 + if w != self._op_w: + self._op_w = w + return True + return False + + def action_toggle_opcodes(self) -> None: + # cycle: off -> limited -> full -> off + self._op_mode = (self._op_mode + 1) % 3 + self._update_op_w() + self._ranges = {} # column layout changed -> stale match offsets + self._clamp_x() + self.refresh() + self._app_status("opcodes: " + {0: "off", 1: f"limited ({_OP_LIMIT} bytes)", + 2: "full"}[self._op_mode]) + + @work(thread=True, exclusive=True, group="listing-grow") + def _grow(self) -> None: + """Stream the rest of the segment's heads in the background, growing the + virtual size as they land so the scrollbar/paging catch up.""" + model = self.model + if model is None: + return + since = 0 + while not model.complete: + if model.load_next_page() == 0: + break + if self.model is not model: # a new load() replaced us + return + since += 1 + if since >= 4: # throttle repaints on huge segments + since = 0 + self.app.call_from_thread(self._grew, len(model)) + self.app.call_from_thread(self._grew, len(model)) + + def _grew(self, total: int) -> None: + if self.model is None or total <= self.total: + return + self.total = total + self.virtual_size = Size(0, total) + self._update_op_w() # more code streamed in -> widen the op column self.refresh() + # -- search hooks ----------------------------------------------------- # + def _search_line_count(self) -> int: + return self.total + + def _search_line_text(self, i: int) -> str | None: + return self._line_plain(i) + + def _search_ensure(self, done) -> None: + # Search needs the whole segment loaded to find every match. Kick off a + # SINGLE background load (guarded so per-keystroke updates don't spawn + # one worker each — an exclusive worker would cancel+restart on every + # keypress and never finish) and queue the callbacks until it lands. + model = self.model + if model is None or model.complete: + done() + return + self._search_pending.append(done) + if not self._search_loading: + self._search_loading = True + self._search_load_all() + + @work(thread=True, group="listing-search-load") + def _search_load_all(self) -> None: + model = self.model + if model is not None: + model.load_all() + self.app.call_from_thread(self._search_loaded) + + def _search_loaded(self) -> None: + self._search_loading = False + if self.model is not None: + self._grew(len(self.model)) + pending, self._search_pending = self._search_pending, [] + for cb in pending: + cb() + # -- rendering --------------------------------------------------------- # def render_line(self, y: int) -> Strip: model = self.model @@ -717,85 +1058,108 @@ class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True 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: + h = model.get(idx) + if h is None: strip = Strip([Segment(f" {idx:>8} …", _S_DIM)]) + elif h.kind == "sep": + strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(_LST_INDENT + h.text, _S_SEP)]) + elif h.kind == "funchdr": + # depth-0: address + 'name proc'/'endp' (no indent) + strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(h.text, _S_FUNCHDR)]) + elif h.kind == "label": + # depth-0: address + 'loc_XXX:' on its own line + strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(h.text, _S_LABEL)]) else: - segs: list[Segment] = [Segment(f"{line.ea:08X} ", _S_ADDR)] - op = self._op_field(line) + # depth-1: address, one indent, then opcode+text + segs: list[Segment] = [Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(_LST_INDENT, _S_INSN)] + op = self._op_field(h) 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)) + if h.kind == "member": + segs.append(Segment(_LST_INDENT, _S_MEMBER)) + if h.name: + segs.append(Segment(f"{h.name} ", _S_LABEL)) + if h.kind == "member": + segs.append(Segment(h.text, _S_MEMBER)) + else: + 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)) - if is_cursor: - strip = _cursor_decorate(strip, self._line_plain(idx) or "", self.cursor_x) - return strip.adjust_cell_length(width, _S_INSN) + if self._hl_word and plain: + occ = _word_occurrences(plain, self._hl_word) + if occ: + 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 _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 + def _cursor_operand(self, idx: int) -> tuple[int, int] | None: + """Screen columns of the operand under the cursor on row ``idx``, or None. - @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) + 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 - def _settle_op_w(self) -> None: - if self._update_op_w(): - self._search_texts = None # layout changed -> stale offsets + # -- split-view link highlight ---------------------------------------- # + def set_link(self, rows) -> None: + rows = set(rows) if rows else set() + if rows != self._link_rows: + self._link_rows = rows 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 reveal(self, row: int) -> None: + """Scroll ``row`` into view without moving the cursor (companion pane).""" + height = self._visible_height() + top = round(self.scroll_offset.y) + if row < top or row >= top + height: + self.scroll_to(y=max(row - height // 3, 0), animate=False) - 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 + def align(self, row: int, screen_row: int) -> None: + """Scroll so ``row`` sits at viewport offset ``screen_row`` — keeps this + (companion) pane visually level with the driver's cursor in split view. + Clamps at the ends, so alignment is best-effort near the edges.""" + top = max(0, min(row - max(screen_row, 0), max(self.total - 1, 0))) + if top != round(self.scroll_offset.y): + self.scroll_to(y=top, animate=False) - @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) + def watch_scroll_y(self, old_value: float, new_value: float) -> None: + super().watch_scroll_y(old_value, new_value) + if round(old_value) != round(new_value): + self.post_message(ListingView.Scrolled()) # -- navigation -------------------------------------------------------- # def _visible_height(self) -> int: @@ -818,19 +1182,58 @@ class DisasmView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self._clamp_x() self._scroll_cursor_into_view() if round(self.scroll_offset.y) != before: - self.refresh() # scrolled: the whole viewport shifted + self.refresh() else: - _refresh_lines(self, old, self.cursor) # only the two changed rows - self.post_message(DisasmView.CursorMoved(self.cursor, self._cursor_ea())) + _refresh_lines(self, old, self.cursor) + self.post_message(ListingView.CursorMoved(self.cursor, self._cursor_ea())) def _after_cursor_move(self) -> None: - self.post_message(DisasmView.CursorMoved(self.cursor, self._cursor_ea())) + self.post_message(ListingView.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 + h = self._head(self.cursor) + return h.ea if h else None + + def _next_ea(self) -> int | None: + h = self._head(self.cursor + 1) + return h.ea if h 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_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) def action_cursor_down(self) -> None: self._move(1) @@ -863,6 +1266,32 @@ def _refresh_lines(view, *indices: int) -> None: # --------------------------------------------------------------------------- # # Decompiler (pseudocode) view # --------------------------------------------------------------------------- # +class _DecompLoading(Static): + """Animated cover shown over the pseudocode pane while a (re)decompile runs. + A braille spinner + label, dim over the grayed-out code — fits the muted TUI + palette (no ASCII-bar noise).""" + + _FRAMES = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏" + + def __init__(self, **kwargs) -> None: + self._i = 0 + super().__init__(self._label(), **kwargs) + + def _label(self) -> Text: + t = Text(justify="center") + t.append(self._FRAMES[self._i], _S_DECOMP_SPIN) + t.append(" decompiling", _S_DECOMP_WAIT) + t.append("…", _S_DECOMP_DOTS) + return t + + def on_mount(self) -> None: + self.set_interval(1 / 12, self._tick) + + def _tick(self) -> None: + self._i = (self._i + 1) % len(self._FRAMES) + self.update(self._label()) + + class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True): """Read-only, line-virtualized Hex-Rays pseudocode with Pygments C highlighting. Lines are highlighted once at load and cached as Strips, so @@ -871,14 +1300,25 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True BINDINGS = [ Binding("tab,shift+tab", "app.toggle_view", "Disasm", priority=True), + Binding("f5", "app.toggle_view", "Decompile", priority=True), + Binding("L", "app.continuous_here", "Listing"), 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), + # Home/End move the cursor along the line (as in the listing); top and + # bottom of the function move to the ctrl+ pair (G still works too). + Binding("home", "col_home", "bol", show=False), + Binding("shift+home", "col_code_home", "code start", show=False), + 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, @@ -895,6 +1335,10 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self.index = index self.ea = ea + class Scrolled(Message): + """Posted when the viewport scrolls (wheel/scrollbar) — the cursor need + not have moved, so the split view can still follow along.""" + def __init__(self) -> None: super().__init__(id="decomp") self.loaded_ea: int | None = None @@ -902,13 +1346,45 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self._texts: list[str] = [] 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 @@ -982,19 +1458,27 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True else: scroll_y = top width = max(self.size.width - self._gutter, 1) - sx = round(self.scroll_offset.x) - if self.cursor_x < sx: - scroll_x = self.cursor_x - elif self.cursor_x >= sx + width: + # Always baseline the horizontal scroll at 0 for a jump, then scroll + # right only if the target column falls outside the viewport — never + # keep a stale horizontal offset from wherever we were before. + sx = 0 + if self.cursor_x >= sx + width: scroll_x = self.cursor_x - width + 1 else: scroll_x = sx self._apply_scroll(min(max(scroll_y, 0), max(total - 1, 0)), max(scroll_x, 0)) + # `cursor` is reactive(repaint=False) and _apply_scroll only repaints via + # call_after_refresh, so a jump that lands in the SAME viewport (Esc back + # to another spot in the function already on screen) moved the cursor + # with nothing to redraw it — the pane kept showing the old highlight + # until the next keypress. Repaint here; _move does the same via + # _refresh_lines. + self.refresh() self._after_cursor_move() def get_loading_widget(self): # type: ignore[override] # Shown (grayed, centered) while a (re)decompile is in flight. - return Static("― decompiling… ―", classes="decomp-loading") + return _DecompLoading(classes="decomp-loading") def _line_ea(self, idx: int) -> int | None: return self._line_eas[idx] if 0 <= idx < len(self._line_eas) else None @@ -1023,20 +1507,84 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True return Strip([Segment(" " * gw, _S_LINENO)]).adjust_cell_length(width) return Strip.blank(width) x = round(self.scroll_offset.x) + linked = idx == self._link_line 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: + occ = _word_occurrences(self._texts[idx], self._hl_word) + if occ: + 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) + code = base.crop(x, x + code_w).adjust_cell_length( + code_w, _S_LINK if linked else None) if gw <= 0: return code style = _S_LINENO_CUR if idx == self.cursor else _S_LINENO + if linked: + style = style + _S_LINK gutter = Strip([Segment(f"{idx + 1:>{gw - 1}} ", style)]) return Strip.join([gutter, code]).adjust_cell_length(width) - # -- navigation (mirrors DisasmView) ---------------------------------- # + # -- split-view link highlight ---------------------------------------- # + def set_link(self, line: int | None) -> None: + if line != self._link_line: + self._link_line = line + self.refresh() + + def reveal(self, line: int) -> None: + """Scroll ``line`` into view without moving the cursor (companion pane).""" + height = self._visible_height() + top = round(self.scroll_offset.y) + if line < top or line >= top + height: + self.scroll_to(y=max(line - height // 3, 0), animate=False) + + def align(self, line: int, screen_row: int) -> None: + """Scroll so ``line`` sits at viewport offset ``screen_row`` — keeps this + (companion) pane visually level with the driver's cursor in split view.""" + top = max(0, min(line - max(screen_row, 0), max(len(self._strips) - 1, 0))) + if top != round(self.scroll_offset.y): + self.scroll_to(y=top, animate=False) + + def watch_scroll_y(self, old_value: float, new_value: float) -> None: + super().watch_scroll_y(old_value, new_value) + 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.""" + text = self._line_plain(self.cursor) or "" + self.cursor_x = len(text) - len(text.lstrip()) if text.strip() else 0 + self._hscroll() + _refresh_lines(self, self.cursor) + self._refresh_hl() + + def line_for_ea(self, ea: int) -> int | None: + """The pseudocode line whose marker ea is the largest <= ``ea`` (the C + line that best covers an instruction address).""" + best, best_ea = None, -1 + for i, e in enumerate(self._line_eas): + if e is not None and best_ea < e <= ea: + best, best_ea = i, e + return best + + # -- navigation -------------------------------------------------------- # def _visible_height(self) -> int: return max(self.size.height, 1) @@ -1060,6 +1608,7 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self.refresh() else: _refresh_lines(self, old, self.cursor) + self._refresh_hl() self._after_cursor_move() def action_cursor_down(self) -> None: @@ -1078,9 +1627,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True # --------------------------------------------------------------------------- # # Hex view (raw bytes of the loaded image, VA-addressed, virtualized) # --------------------------------------------------------------------------- # -_S_HEX = Style(color="grey74") -_S_ASCII = Style(color="#6a9955") -_S_FOFF = Style(color="#c586c0") # file-offset column (distinct from the VA) +_S_HEX = Style(color="#c3cad3") +_S_ASCII = Style(color="#9ece6a") +_S_FOFF = Style(color="#d8a657") class HexView(ScrollView, can_focus=True): @@ -1103,6 +1652,7 @@ class HexView(ScrollView, can_focus=True): Binding("enter", "to_code", "To code"), Binding("escape", "leave", "Back"), Binding("tab,shift+tab", "app.toggle_view", "Code", priority=True), + Binding("f5", "app.toggle_view", "Decompile", priority=True), ] cursor = reactive(0, repaint=False) @@ -1124,6 +1674,11 @@ 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 -------------------------------------------------------- # def load(self, model, va: int | None = None) -> None: # type: ignore[no-untyped-def] @@ -1165,12 +1720,16 @@ class HexView(ScrollView, can_focus=True): def _apply_scroll(self, y: int) -> None: y = max(0, y) - self.scroll_to(y=y, animate=False) + if round(self.scroll_offset.y) != y: + self._internal_top = y # cursor-driven scroll incoming; don't follow it def _fix(yy: int = y) -> None: + if round(self.scroll_offset.y) != yy: + self._internal_top = yy self.scroll_to(y=yy, animate=False) self.refresh(layout=True) + self.scroll_to(y=y, animate=False) self.call_after_refresh(_fix) def _scroll_to_cursor(self, center: bool = False) -> None: @@ -1187,6 +1746,28 @@ class HexView(ScrollView, can_focus=True): top = max(0, min(top, max(self.total - 1, 0))) self._apply_scroll(top) + def watch_scroll_y(self, old_value: float, new_value: float) -> None: + super().watch_scroll_y(old_value, new_value) + if round(old_value) == round(new_value): + return + nt = round(new_value) + if self._internal_top is not None and nt == self._internal_top: + self._internal_top = None # our cursor-driven scroll: cursor already placed + return + # A user scroll (wheel / scrollbar): drag the cursor by the same row delta + # so its screen position stays frozen (it points at a new byte in place). + self._internal_top = None + self._shift_cursor(nt - round(old_value)) + + def _shift_cursor(self, rows: int) -> None: + if not rows or self.model is None or self.model.size == 0: + return + new = max(0, min(self.cursor + rows * 16, self.model.size - 1)) + if new != self.cursor: + self.cursor = new + self.refresh() # cursor is reactive(repaint=False) + self.post_message(HexView.Moved(self.cursor_va())) + # -- navigation -------------------------------------------------------- # def _move(self, delta: int) -> None: if self.model is None or self.model.size == 0: @@ -1219,6 +1800,39 @@ class HexView(ScrollView, can_focus=True): def action_leave(self) -> None: self.post_message(HexView.Leave()) + # -- mouse ------------------------------------------------------------- # + def _byte_at_x(self, x: int) -> int: + """Map a content column to a byte index 0..15, across the hex and ASCII + panes. Matches ``render_line``'s layout: addr(9) + file-offset(10) + 16 + hex cells of 3 cols (with a 1-col gap before byte 8), then ' |' + ASCII.""" + HEX, ASCII = 19, 70 + if x < HEX: # clicked the address/offset gutter -> row start + return 0 + if x < HEX + 49: # hex byte region + rel = x - HEX + if rel >= 24: # collapse the 1-col gap between the two halves + rel -= 1 + return min(rel // 3, 15) + if x < ASCII: # the ' |' separator -> last byte of the row + return 15 + return min(x - ASCII, 15) # ASCII pane (and anything past it) + + def on_click(self, event) -> None: # type: ignore[no-untyped-def] + if self.model is None or self.model.size == 0: + return + off = event.get_content_offset(self) + if off is None: + return + self.focus() + row = round(self.scroll_offset.y) + off.y + x = round(self.scroll_offset.x) + off.x + new = row * 16 + self._byte_at_x(x) + self.cursor = max(0, min(self.model.size - 1, new)) + self.refresh() # cursor is reactive(repaint=False); repaint the highlight + self.post_message(HexView.Moved(self.cursor_va())) + if event.chain >= 2: # double-click == place cursor + jump to code + self.post_message(HexView.ToCode(self.cursor_va())) + # -- rendering --------------------------------------------------------- # def _ensure_window(self, top: int) -> None: if self.model is None: @@ -1250,6 +1864,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] = [ @@ -1264,15 +1887,29 @@ class HexView(ScrollView, can_focus=True): if i == 8: segs.append(Segment(" ", _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 + segs.append(Segment(f"{val:02X} ", st)) else: segs.append(Segment(" ", _S_HEX)) segs.append(Segment(" |", _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)) @@ -1281,6 +1918,822 @@ class HexView(ScrollView, can_focus=True): # --------------------------------------------------------------------------- # +# 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 + + +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: + for k, c in enumerate(s): + self.put(i + k, c, style) + + def restyle(self, a: int, b: int, style: Style) -> None: + """Merge ``style`` over the cells in [a, b) (keeps the characters).""" + for i in range(max(a, 0), min(b, self.width)): + self.st[i] = self.st[i] + style + + 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.put(left, graph.BOX["tl"], bs) + for i in range(1, w - 1): + out.put(left + i, graph.BOX["h"], bs) + out.put(left + w - 1, 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.put(left, graph.BOX["bl"], bs) + for i in range(1, w - 1): + out.put(left + i, graph.BOX["h"], bs) + out.put(left + w - 1, graph.BOX["br"], bs) + return + out.put(left, graph.BOX["v"], bs) + out.put(left + w - 1, graph.BOX["v"], bs) + for i in range(1, w - 1): + out.put(left + i, " ", _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): @@ -1303,8 +2756,10 @@ class XrefsScreen(ModalScreen): BINDINGS = [Binding("escape", "close", "Close")] - def __init__(self, label: str, items: list[tuple[int, str]], + def __init__(self, label: str, items: list[tuple[object, str]], preselect: int = 0) -> None: + # payload is an int address, or (binary, address) for a caller in another + # project binary; dismiss() hands it back untouched. super().__init__() self._label = label self._items = items @@ -1339,10 +2794,18 @@ def _fuzzy(name: str, q: str): matched character indices (for highlighting).""" if not q: return (0.0, ()) + if not name: # defensive: never assume a symbol has a name + return None nl = name.lower() + # Case-insensitive means BOTH sides: the name was lowered but the query + # wasn't, so a single capital could never match and any query containing one + # returned nothing at all. Invisible on lowercase C symbols (main, strlen), + # fatal on libraries that capitalise — "PEM_read_bio" found 0 of 10093 + # functions in libcrypto while the backend resolved it fine. + ql = q.lower() pos: list[int] = [] i = 0 - for ch in q: + for ch in ql: j = nl.find(ch, i) if j < 0: return None @@ -1350,9 +2813,9 @@ def _fuzzy(name: str, q: str): i = j + 1 span = pos[-1] - pos[0] score = -(span * 2.0) - pos[0] - len(name) * 0.01 - if q in nl: + if ql in nl: score += 50.0 - if nl.startswith(q): + if nl.startswith(ql): score += 100.0 return (score, tuple(pos)) @@ -1364,17 +2827,28 @@ class SymbolPalette(ModalScreen): Binding("escape", "close", "Close"), Binding("down,ctrl+n", "cursor_down", show=False), Binding("up,ctrl+p", "cursor_up", show=False), + # F2, not ctrl+a: the focused Input binds "home,ctrl+a" so it would never + # reach us. Function keys are untouched by Input. + Binding("f2", "scope", "This binary / whole project", show=False), ] LIMIT = 200 + #: Project scope nearly always saturates its cap (every binary contributes), + #: where a local filter usually returns a handful. Keep the list short so + #: arrowing through it stays snappy; narrow further by typing. + PROJECT_LIMIT = 60 - def __init__(self, funcs: list[Func]) -> None: + def __init__(self, funcs: list[Func], index=None, binary=None) -> None: super().__init__() self._funcs = funcs - self._results: list[Func] = [] + self._index = index # ProjectIndex, when this is a project + self._binary = binary # label of the binary we're currently in + self._project_scope = False + #: (binary|None, addr, name) — binary is None for a local hit + self._results: list[tuple] = [] def compose(self) -> ComposeResult: with Vertical(id="pal-box"): - yield Static(" symbols", id="pal-title") + yield Static(" symbols", id="pal-title", markup=False) yield Input(placeholder="fuzzy find symbol… ↑↓ select · Enter open · Esc close", id="pal-input") yield OptionList(id="pal-list") @@ -1391,35 +2865,735 @@ class SymbolPalette(ModalScreen): event.stop() self.action_choose() + def action_scope(self) -> None: + if self._index is None: + return # not a project: nothing else to search + self._project_scope = not self._project_scope + self._apply(self.query_one("#pal-input", Input).value.strip()) + def _apply(self, query: str) -> None: - if query: + # rows: (binary|None, addr, name, match positions) + if self._project_scope and self._index is not None: + from .index import KIND_FUNC + # The trigram index already guarantees every hit CONTAINS the query, + # so ranking only has to order them — an exact-substring rank (match + # position, then name length) costs a find() per row instead of a + # full fuzzy pass, and fetching 3x the display limit rather than 10x + # keeps the per-keystroke work down on a big project. + hits = self._index.search(query, kind=KIND_FUNC, + limit=self.PROJECT_LIMIT * 3) if query else [] + q = query.lower() + scored = [] + for h in hits: + at = h.text.lower().find(q) + scored.append((at if at >= 0 else 1 << 30, len(h.text), h.text, h)) + # Sort on an explicit key: the same symbol name in two binaries ties + # on (position, length, text), and a bare sort() would then fall + # through to comparing Hit objects, which aren't orderable. + scored.sort(key=lambda t: (t[0], t[1], t[2], t[3].binary, t[3].addr)) + rows = [(h.binary, h.addr, h.text, + tuple(range(at, at + len(q))) if at < (1 << 30) else ()) + for at, _, _, h in scored[:self.PROJECT_LIMIT]] + elif query: scored = [] for f in self._funcs: m = _fuzzy(f.name, query) if m is not None: scored.append((m[0], m[1], f)) scored.sort(key=lambda t: (-t[0], t[2].name)) - rows = [(f, pos) for _, pos, f in scored[:self.LIMIT]] + rows = [(None, f.addr, f.name, pos) for _, pos, f in scored[:self.LIMIT]] else: - rows = [(f, ()) for f in self._funcs[:self.LIMIT]] - self._results = [f for f, _ in rows] + rows = [(None, f.addr, f.name, ()) for f in self._funcs[:self.LIMIT]] + self._results = [(b, a, n) for b, a, n, _ in rows] ol = self.query_one(OptionList) ol.clear_options() opts = [] - for f, pos in rows: + for binary, addr, name, pos in rows: label = Text() - label.append(f"{f.addr:08X} ", _S_ADDR) - nm = Text(f.name) + if binary: + label.append(f"{binary:<14.14} ", _S_LABEL) + label.append(f"{addr:08X} ", _S_ADDR) + nm = Text(name) for p in pos: - if p < len(f.name): + if p < len(name): nm.stylize(_S_NAME_MATCH, p, p + 1) label.append_text(nm) opts.append(Option(label)) ol.add_options(opts) if self._results: ol.highlighted = 0 + scope = "project" if self._project_scope else "this binary" + cap = self.PROJECT_LIMIT if self._project_scope else self.LIMIT + 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}") + + 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._results): + b, a, _ = self._results[i] + self.dismiss((b, a)) + + def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: + if 0 <= event.option_index < len(self._results): + b, a, _ = self._results[event.option_index] + self.dismiss((b, a)) + + def action_close(self) -> None: + self.dismiss(None) + + +def _str_display(text: str, limit: int = 200) -> str: + """One-line, printable rendering of a string literal for the browser: escape + the common control chars, drop the rest, and clip long bodies.""" + out = (text.replace("\\", "\\\\").replace("\n", "\\n") + .replace("\r", "\\r").replace("\t", "\\t")) + out = "".join(ch if ch.isprintable() else "." for ch in out) + return out[:limit] + ("\u2026" if len(out) > limit else "") + + +class StringsPalette(ModalScreen): + """Every string in the binary (IDA's Shift+F12), filterable; Enter jumps to + it in the unified listing.""" + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + Binding("f2", "scope", "This binary / whole project", show=False), + ] + LIMIT = 500 + #: Project scope saturates its cap (every binary contributes); keep the list + #: short so arrowing stays snappy. + PROJECT_LIMIT = 60 + + def __init__(self, strings: list, index=None, binary=None) -> None: + super().__init__() + # Pre-render + pre-lower once: filtering runs on every keystroke and a + # big binary has tens of thousands of strings. + self._rows = [(s, d, d.lower()) + for s in strings for d in (_str_display(s.text),)] + self._index = index # ProjectIndex, when this is a project + self._binary = binary + self._project_scope = False + #: (binary|None, addr, display text) — binary is None for a local hit + self._results: list[tuple] = [] + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box"): + yield Static(" strings", id="pal-title", markup=False) + yield Input(placeholder="filter strings\u2026 \u2191\u2193 select \u00b7 " + "Enter jump \u00b7 Esc close", id="pal-input") + yield OptionList(id="pal-list") + + def on_mount(self) -> None: + self._apply("") + self.query_one("#pal-input", Input).focus() + + def on_input_changed(self, event: Input.Changed) -> None: + event.stop() # don't leak to the app's #search/#filter handlers + self._apply(event.value.strip()) + + def on_input_submitted(self, event: Input.Submitted) -> None: + event.stop() + self.action_choose() + + def action_scope(self) -> None: + if self._index is None: + return # not a project: nothing else to search + self._project_scope = not self._project_scope + self._apply(self.query_one("#pal-input", Input).value.strip()) + + def _apply(self, query: str) -> None: + q = query.lower() + # rows: (binary|None, addr, length, display text, match offset) + if self._project_scope and self._index is not None: + from .index import KIND_STRING + hits = self._index.search(query, kind=KIND_STRING, + limit=self.PROJECT_LIMIT * 3) if query else [] + rows = [] + for h in hits: + disp = _str_display(h.text) + rows.append((h.binary, h.addr, len(h.text), disp, + disp.lower().find(q))) + # the index already guarantees a match, so ranking only orders them: + # earliest match, then shortest, with a stable (binary, addr) tiebreak + # (a literal shared by two binaries would otherwise be unordered). + rows.sort(key=lambda r: (r[4] if r[4] >= 0 else 1 << 30, + r[2], r[0], r[1])) + rows = rows[:self.PROJECT_LIMIT] + else: + rows = [] + for s, disp, low in self._rows: + hit = low.find(q) if q else -1 + if q and hit < 0: + continue + rows.append((None, s.addr, s.length, disp, hit)) + if len(rows) >= self.LIMIT: + break + self._results = [(b, a, d) for b, a, _, d, _ in rows] + ol = self.query_one(OptionList) + ol.clear_options() + opts = [] + for binary, addr, length, disp, hit in rows: + label = Text() + if binary: + label.append(f"{binary:<14.14} ", _S_LABEL) + label.append(f"{addr:08X} ", _S_ADDR) + label.append(f"{length:>5} ", _S_DIM) + body = Text(disp) + if hit >= 0: + body.stylize(_S_NAME_MATCH, hit, hit + len(q)) + label.append_text(body) + opts.append(Option(label)) + ol.add_options(opts) + if self._results: + ol.highlighted = 0 + scope = "project" if self._project_scope else "this binary" + cap = self.PROJECT_LIMIT if self._project_scope else self.LIMIT + 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}") + + 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._results): + b, a, _ = self._results[i] + self.dismiss((b, a)) + + def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: + if 0 <= event.option_index < len(self._results): + b, a, _ = self._results[event.option_index] + self.dismiss((b, a)) + + def action_close(self) -> None: + self.dismiss(None) + + +#: The keyboard cheatsheet (F1). Grouped by task rather than by widget, which is +#: what makes it readable; keep it in step with the BINDINGS above it. +_HELP = ( + ("Navigate", ( + ("Enter", "follow the symbol under the cursor"), + ("Esc", "back (navigation history)"), + ("g", "goto address or symbol"), + ("Ctrl+N", "find symbol (fuzzy)"), + ("\"", "strings browser"), + ("x", "cross-references to the symbol"), + ("L", "continuous listing at the cursor"), + ("Ctrl+O", "switch binary (projects)"), + )), + ("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"), + ("B", "cycle the opcode-bytes column"), + ("Ctrl+B", "show/hide the names pane"), + ("Ctrl+T", "structs / types editor"), + ("Ctrl+P", "command palette"), + )), + ("Move", ( + ("j / k", "down / up"), + ("Ctrl+D / Ctrl+U", "half page down / up"), + ("PgDn / PgUp", "page down / up"), + ("Ctrl+Home / Ctrl+End", "top / bottom (G also)"), + ("Home / End", "start / end of line"), + ("Shift+Home", "start of the instruction / code"), + ("h / l", "column left / right"), + ("w / b", "word forward / back"), + )), + ("Edit", ( + ("n", "rename"), + ("y", "set type (prototype, local or global)"), + (";", "comment"), + ("c", "make code"), + ("p", "make function"), + ("d", "make data"), + ("a", "make string"), + ("u", "undefine"), + ("o / O", "literal format: hex/dec/bin/char/offset"), + ("Ctrl+S", "save the database"), + )), + ("Search", ( + ("/", "search forward (repeat to continue)"), + ("?", "search backward"), + ("N", "previous match"), + ("Ctrl+Y", "copy the current line"), + ("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).""" + + BINDINGS = [ + Binding("s", "save", "Save & quit"), + Binding("d", "discard", "Discard & quit"), + Binding("escape,c", "cancel", "Cancel"), + ] + + def __init__(self, labels: list[str]) -> None: + super().__init__() + self._labels = labels + + 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") + 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", + id="quit-help") + + def action_save(self) -> None: + self.dismiss("save") + + def action_discard(self) -> None: + self.dismiss("discard") + + def action_cancel(self) -> None: + self.dismiss(None) + + +class HelpScreen(ModalScreen): + """F1: the keyboard cheatsheet, replacing the permanent footer.""" + + 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 + + def compose(self) -> ComposeResult: + # Fluid: as many columns as the terminal can hold. Textual CSS has no + # media queries, so the split is computed here from the real width. + 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) + # 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. + with VerticalScroll(id="help-body"): + with Horizontal(id="help-cols"): + for c in range(cols): + chunk = _HELP[c * per:(c + 1) * per] + if not chunk: + continue + with Vertical(classes="help-col"): + for title, rows in chunk: + card = Static(self._card(rows), + classes="help-card", markup=False) + card.border_title = title + yield card + yield Static("Esc · F1 · H to close", id="help-foot") + + @staticmethod + def _section_widths() -> list[int]: + """Rendered width of each section's card (keys right-aligned per card).""" + out = [] + for _, rows in _HELP: + kw = max(len(k) for k, _ in rows) + out.append(max(kw + 2 + len(d) for _, d in rows) + HelpScreen._CARD_PAD) + return out + + @classmethod + def _columns(cls, avail: int) -> int: + """Most columns that actually fit in ``avail``. + + Sizing off the widest section would let one long row (Move's + "Ctrl+Home / Ctrl+End") inflate every column and cost a column that would + otherwise fit. Each column hugs its own content, so measure the real + layout: chunk the sections and sum the per-chunk maxima. + """ + ws = cls._section_widths() + n = len(ws) + for cols in range(min(n, 4), 1, -1): + per = -(-n // cols) + chunks = [ws[c * per:(c + 1) * per] for c in range(cols)] + total = sum(max(c) for c in chunks if c) + (cols - 1) + if total <= avail: + return cols + return 1 + + @staticmethod + def _card(rows) -> Text: # NB: not _render — that's a Widget internal + """One section's keys. The key column is sized per section, so a card of + short keys stays narrow instead of padding out to the global maximum.""" + width = max(len(k) for k, _ in rows) + out = Text() + for i, (key, desc) in enumerate(rows): + if i: + out.append("\n") + out.append(f"{key:>{width}}", _S_MNEM) + out.append(f" {desc}", _S_INSN) + return out + + def action_close(self) -> None: + 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"): + yield Static(f" registers at t={self._idx:,} \u2014 Enter seeks to the " + f"write, f seeks forward", id="pal-title", markup=False) + 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. + + IDA's own answer to an unidentified file is a dialog; ours is this. Without + it the fallback is x86 at address 0, which doesn't fail — it analyses to + nothing, and you're left wondering why a firmware image has no functions. + + Returns ``{"processor": str, "base": int}``, or ``{}`` to load it the way + IDA would have anyway (that IS the right answer sometimes: our sniff only + recognises formats we're sure about, so it says "unknown" for things IDA + can in fact handle). + """ + + 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), + ] + + def __init__(self, path: str, size: int = 0) -> None: + super().__init__() + self._path = path + # NOT self._size: that is Textual's own backing field for outer_size, and + # assigning an int to it crashes the layout with a bewildering + # "'int' object has no attribute 'region'" from deep inside _set_dirty. + self._nbytes = size + self._results: list[tuple[str, str]] = [] + + 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) + 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) + yield Input(placeholder="filter processors\u2026", id="pal-input") + yield OptionList(id="pal-list") + yield Input(placeholder="load address, e.g. 0x8000000 (blank = 0)", + id="load-base") + yield Static(" Enter accept \u00b7 Tab base address \u00b7 " + "Esc load as IDA would", id="load-help", markup=False) + + def on_mount(self) -> None: + self._apply("") + self.query_one("#pal-input", Input).focus() + + def focus_next(self, selector="*"): # type: ignore[override] + """Tab moves between the two things you TYPE into. + + DOM order would stop at the option list on the way, which is + arrow-driven and has nothing to type — and the address you meant to + enter goes into whichever box happened to have focus. Typing an address + into the processor filter is then taken as a processor name, IDA rejects + it, and the open fails; that is a bad enough outcome to be worth + overriding Tab for. + """ + inp = self.query_one("#pal-input", Input) + base = self.query_one("#load-base", Input) + (inp if self.focused is base else base).focus() + return self.focused + + def focus_previous(self, selector="*"): # type: ignore[override] + return self.focus_next() + + def on_input_changed(self, event: Input.Changed) -> None: + event.stop() + if event.input.id == "pal-input": + self._apply(event.value.strip()) + + def on_input_submitted(self, event: Input.Submitted) -> None: + event.stop() + self.action_choose() + + def _apply(self, query: str) -> None: + q = query.lower() + rows = [(name, desc) for name, desc in self._all + if not q or q in name.lower() or q in desc.lower()] + # An unlisted processor is still valid: IDA has 73 modules and this + # offers 20, so a typed name that matches nothing is taken literally + # rather than refused. + if not rows and q: + rows = [(query.strip(), "use this processor name as typed")] + self._results = rows + ol = self.query_one(OptionList) + ol.clear_options() + opts = [] + for name, desc in rows: + label = Text() + label.append(f" {name:<12}", _S_LABEL) + label.append(desc, _S_DIM) + opts.append(Option(label)) + ol.add_options(opts) + if rows: + ol.highlighted = 0 self.query_one("#pal-title", Static).update( - f" symbols: {len(self._results)}" + ("+" if len(self._results) == self.LIMIT else "")) + f" unrecognised file \u2014 processor? ({len(rows)})") def action_cursor_down(self) -> None: ol = self.query_one(OptionList) @@ -1434,12 +3608,117 @@ class SymbolPalette(ModalScreen): def action_choose(self) -> None: ol = self.query_one(OptionList) i = ol.highlighted + if i is None or not (0 <= i < len(self._results)): + self.dismiss({}) + return + raw = self.query_one("#load-base", Input).value.strip() + base = 0 + if raw: + try: + base = int(raw, 0) + except ValueError: + self.query_one("#load-help", Static).update( + f" {raw!r} is not an address \u2014 try 0x8000000") + self.query_one("#load-base", Input).focus() + return + if base % 16: + # IDA's -b is in paragraphs, so an unaligned base can't be + # expressed and would quietly load somewhere else. + self.query_one("#load-help", Static).update( + f" {base:#x} must be 16-byte aligned") + self.query_one("#load-base", Input).focus() + return + self.dismiss({"processor": self._results[i][0], "base": base}) + + def action_close(self) -> None: + self.dismiss({}) + + +class ProjectPalette(ModalScreen): + """The project's binaries; Enter switches to one. Shows which are resident + (a live worker, so switching is instant) vs cold (needs an open).""" + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + ] + + def __init__(self, entries: list[dict]) -> None: + super().__init__() + self._entries = entries + self._results: list[dict] = [] + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box"): + yield Static(" binaries", id="pal-title", markup=False) + yield Input(placeholder="filter binaries\u2026 \u2191\u2193 select \u00b7 " + "Enter switch \u00b7 Esc close", id="pal-input") + yield OptionList(id="pal-list") + + def on_mount(self) -> None: + self._apply("") + self.query_one("#pal-input", Input).focus() + + def on_input_changed(self, event: Input.Changed) -> None: + event.stop() + self._apply(event.value.strip()) + + def on_input_submitted(self, event: Input.Submitted) -> None: + event.stop() + self.action_choose() + + def _apply(self, query: str) -> None: + q = query.lower() + rows = [e for e in self._entries + if not q or q in e["label"].lower() or q in e["source"].lower()] + self._results = rows + ol = self.query_one(OptionList) + ol.clear_options() + opts = [] + for e in rows: + label = Text() + label.append("\u25b8 " if e["active"] else " ", + _S_MNEM if e["active"] else _S_DIM) + label.append(f"{e['label']:<22}", _S_LABEL) + if e["resident"]: + mb = e.get("memory_mb") or 0 + label.append(f"resident {mb:>4}MB ", _S_MNEM) + elif e["analysed"]: + label.append("analysed ", _S_DIM) + else: + label.append("not opened ", _S_DIM) + if e["pinned"]: + label.append("pin ", _S_ADDR) + label.append(e["source"], _S_DIM) + opts.append(Option(label)) + ol.add_options(opts) + if rows: + # Land on the binary you're already in, so the switcher opens where + # 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)}") + + 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: + i = self.query_one(OptionList).highlighted if i is not None and 0 <= i < len(self._results): - self.dismiss(self._results[i].addr) + self.dismiss(self._results[i]["label"]) def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: if 0 <= event.option_index < len(self._results): - self.dismiss(self._results[event.option_index].addr) + self.dismiss(self._results[event.option_index]["label"]) def action_close(self) -> None: self.dismiss(None) @@ -1456,13 +3735,16 @@ class ConfirmScreen(ModalScreen): Binding("escape,n", "cancel", "Cancel"), ] - def __init__(self, message: str) -> None: + def __init__(self, message: str, note: str = "") -> None: super().__init__() self._message = message + self._note = note def compose(self) -> ComposeResult: with Vertical(id="confirm-box"): - yield Static(self._message, id="confirm-msg") + yield Static(self._message, id="confirm-msg", markup=False) + if self._note: + yield Static(self._note, id="confirm-note", markup=False) yield Static("[Enter/y] confirm [Esc/n] cancel", id="confirm-help") def action_confirm(self) -> None: @@ -1472,6 +3754,179 @@ class ConfirmScreen(ModalScreen): self.dismiss(False) +_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 +_logo_cells: tuple[int, int] | None = None + + +def logo_cells() -> 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. + """ + 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 + return _logo_cells +_logo_cache: object = False # False == not yet loaded (None == absent/unreadable) + + +def _load_logo() -> "Text | None": + """The ANSI-art splash logo (logo.ans) as a Rich Text, or None if missing. + Loaded once; cursor show/hide escapes are stripped so Rich sees only SGR.""" + global _logo_cache + if _logo_cache is not False: + return _logo_cache # type: ignore[return-value] + try: + with open(_LOGO_PATH, encoding="utf-8", errors="replace") as f: + data = f.read() + data = re.sub(r"\x1b\[\?25[lh]", "", data) # drop cursor hide/show + _logo_cache = Text.from_ansi(data.strip("\n"), no_wrap=True) + except OSError: + _logo_cache = None + return _logo_cache # type: ignore[return-value] + + +class LoadingScreen(ModalScreen): + """Startup overlay shown while the binary is opened + analyzed, so a slow + load (big binary) isn't just empty panes and dead air. The app updates the + note line with progress and dismisses it once we land on a function.""" + + BINDINGS = [Binding("escape", "hide", "Hide")] + + def __init__(self, title: str, note: str = "opening\u2026") -> None: + super().__init__() + self._title = title + self._note = note + self._image = False # drawing the real image, not the block art + self._last_place = 0.0 # throttles re-anchoring after a repaint + + def _fits(self, rows: int) -> bool: + """Room for the art plus the title/note/help lines and box chrome.""" + sz = self.app.size + return sz.height >= rows + 9 and sz.width >= 64 + + def compose(self) -> ComposeResult: + with Vertical(id="loading-box"): + # 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. + cols, rows = logo_cells() + kittygfx.log(f"compose: supported={kittygfx.supported()} " + f"app.size={self.app.size} cells={cols}x{rows} " + f"fits={self._fits(rows)}") + if kittygfx.supported() and self._fits(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. + yield Static(Align.center(logo), id="loading-logo") + yield Static(f"\u23f3 loading {self._title}", id="loading-title") + yield Static(self._note, id="loading-note") + yield Static("first open of a big binary can take a while \u00b7 " + "Esc to hide", id="loading-help") + + def update_note(self, text: str) -> None: + try: + 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 + cols, rows = logo_cells() + 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() + + +class BusyScreen(ModalScreen): + """A tiny blocking overlay for a short async step (e.g. gathering xrefs) so + the user can't fire more actions into a half-finished operation. Esc cancels + (dismisses with "cancel"); a programmatic dismiss carries no result.""" + + BINDINGS = [Binding("escape", "cancel", "Cancel")] + + def __init__(self, message: str) -> None: + super().__init__() + self._message = message + + def compose(self) -> ComposeResult: + with Vertical(id="busy-box"): + yield Static(self._message, id="busy-msg") + yield Static("Esc to cancel", id="busy-help") + + def action_cancel(self) -> None: + self.dismiss("cancel") + + # --------------------------------------------------------------------------- # # Struct editor (C-style local type editor) # --------------------------------------------------------------------------- # @@ -1708,16 +4163,131 @@ class StructEditor(ModalScreen): # --------------------------------------------------------------------------- # # The app # --------------------------------------------------------------------------- # +#: The app's own theme. Textual's default (textual-dark) paints every accent and +#: border in #ffa62b, a neon orange that fights the muted VS Code/Solarized +#: palette the code views already use (amber #b58900 for matches, #6a9955 green, +#: #264f78 blue). This keeps the chrome in the same family as the content: +#: desaturated blue-greys, amber for emphasis, blue reserved for focus. +IDATUI_THEME = Theme( + name="idatui", + dark=True, + background="#12161c", # deep blue-black, softer than pure black + surface="#181d25", # views + panel="#212832", # dialogs, status bar, gutters + foreground="#d6d9de", + primary="#5aa0d6", # focus / links: the one cool accent + secondary="#2f5d82", + accent="#d0a215", # the same amber as a search match — one meaning + warning="#c9762f", # burnt orange — distinct from accent, reads as care + error="#ff5f5f", # already used for failure text + success="#6a9955", # already used for comments +) + + +class IdaCommands(Provider): + """Fills the Ctrl+P command palette with real ida-tui actions instead of the + stock Textual system commands (change theme / take screenshot / …). + + App-level actions run directly; cursor-scoped ones (rename/xrefs/comment/…) + are dispatched to the active code view via ``IdaTui._palette_action``.""" + + def _commands(self): + app = self.app + va = app._palette_action # dispatch to the focused code view + return ( + ("Goto address / symbol…", "jump to an address or name (g)", + app.action_goto), + ("Find symbol…", "fuzzy function finder (Ctrl+N)", app.action_symbols), + ("Strings…", "browse every string in the binary (\")", + app.action_strings), + ("Switch binary…", "another binary in the project (Ctrl+O)", + app.action_switch_binary), + ("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)", + lambda: va("xrefs")), + ("Back", "navigation history (Esc)", app.action_back), + ("Toggle disassembly / pseudocode", "decompile / listing (F5, Tab)", + app.action_toggle_view), + ("Continuous listing here", "flat segment listing (L)", + app.action_continuous_here), + ("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)", + lambda: va("retype")), + ("Add comment…", "comment at the cursor (;)", lambda: va("comment")), + ("Define code", "make code at the cursor (c)", lambda: va("define_code")), + ("Create function", "define a function at the cursor (p)", + lambda: va("define_func")), + ("Make data", "define a data item at the cursor (d)", + lambda: va("make_data")), + ("Make string", "define a string at the cursor (a)", + lambda: va("make_string")), + ("Undefine", "undefine the item at the cursor (u)", + lambda: va("undefine")), + ("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), + ("Filter functions…", "glob-filter the function list (/)", + app.action_filter), + ("Toggle names pane", "function-list sidebar (Ctrl+B)", + app.action_toggle_functions), + ("Save database (.i64)", "persist changes (Ctrl+S)", app.action_save), + ("Quit", "exit ida-tui (q)", app.action_quit), + ) + + async def discover(self): + for title, help_text, cb in self._commands(): + yield DiscoveryHit(title, cb, help=help_text) + + async def search(self, query: str): + matcher = self.matcher(query) + for title, help_text, cb in self._commands(): + score = matcher.match(title) + if score > 0: + yield Hit(score, matcher.highlight(title), cb, help=help_text) + + class IdaTui(App): + COMMANDS = {IdaCommands} # replace the stock system-commands palette + CSS = """ Screen { layout: vertical; } #panes { height: 1fr; } #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; @@ -1734,6 +4304,10 @@ class IdaTui(App): height: 1; border: none; padding: 0 1; background: $accent-darken-2; color: $text; } + #makedata { + height: 1; border: none; padding: 0 1; + background: $secondary-darken-2; color: $text; + } #goto { height: 1; border: none; padding: 0 1; background: $primary-darken-3; color: $text; @@ -1741,24 +4315,58 @@ class IdaTui(App): #status { height: 1; background: $panel; color: $text; padding: 0 1; } .decomp-loading { width: 100%; height: 100%; content-align: center middle; - background: $panel-darken-1; color: $text-muted; text-style: italic bold; + 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; } + #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: $text; padding: 0 1; } + #xref-title { dock: top; height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } #xref-list { height: auto; max-height: 100%; } - SymbolPalette { align: center middle; } + /* every #pal-box palette centres, not just the symbol one */ + SymbolPalette, StringsPalette, ProjectPalette, + LoadOptionsScreen, RegWriteScreen { align: center middle; } + /* 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: $text; padding: 0 1; } + #pal-title { dock: top; height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } #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 + was clipped off the bottom, which read as "Tab does nothing". */ + 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; } #se-panes { height: 1fr; } #se-left { width: 38; border-right: solid $accent; } #se-right { width: 1fr; } - #se-title, #se-hint { height: 1; background: $accent; color: $text; padding: 0 1; } + #se-title, #se-hint { height: 1; background: $accent; color: $background; text-style: bold; padding: 0 1; } #se-list { height: 1fr; } #se-edit { height: 1fr; border: none; } #se-status { height: 1; background: $panel-darken-2; color: $text-muted; padding: 0 1; } @@ -1767,6 +4375,19 @@ class IdaTui(App): background: $panel; 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-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-msg { height: 1; text-style: bold; } + #busy-help { height: 1; color: $text-muted; margin-top: 1; } """ BINDINGS = [ @@ -1774,6 +4395,29 @@ class IdaTui(App): Binding("ctrl+n", "symbols", "Symbols"), Binding("ctrl+t", "structs", "Structs"), 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), + # 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), @@ -1782,33 +4426,91 @@ class IdaTui(App): Binding("escape", "back", "Back"), ] - def __init__(self, url: str, db: str | None, keepalive: bool = True, - open_path: str | None = None, rpc_path: str | None = None) -> None: + def __init__(self, open_path: str | None = None, keepalive: bool = True, + rpc_path: str | None = None, ttl: int = 1800, + project=None, load_args: str = "", trace_path: str = "") -> None: super().__init__() - self._url = url - self._db = db + # Project mode is additive: with no project this is the plain + # single-binary app, unchanged. + self._project = project + self._pool = None + self._binary: str | None = None # active project binary (label) + self._states: dict[str, BinaryState] = {} + self._pending_restore = None # entry to reopen after a switch + self._goto_after_switch = None # cross-binary search hit to land on + self._hops: list[str] = [] # binaries a navigation crossed FROM + 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) + 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._title = (os.path.basename(open_path) if open_path else "") + self._trace_path = trace_path or "" # Tenet execution trace to explore + self._trace = None # the loaded Trace, once analysed + self._trail_map = [] # decomp_map for _trail_map_ea + self._trail_map_ea = None + self._pending_trace_line = None # step waiting on a re-decompile + self._trail_line_of: dict[int, int] = {} # ea -> pseudocode line + self._trail_span = None # ea span of that function + self._trail_eas: list[int] = [] # sorted keys of _trail_line_of + self._t = 0 # current timestamp in that trace self._do_keepalive = keepalive self._rpc_path = rpc_path self._rpc = None - self.client: IDAClient | None = None + self.client: WorkerClient | None = None self.program: Program | None = None + self._loading_screen: LoadingScreen | None = None self._ka = None self._nav: list[NavEntry] = [] self._func_index = None # the unfiltered FunctionIndex (source of truth) + self._did_auto_land = False # startup jump-to-main/picker fires once self._filter_term = "" self._pending_filter = "" self._filter_timer = None self._sort_col = 0 # 0=addr, 1=name, 2=size self._sort_reverse = False - self._pref = "decomp" # preferred code view (changed by Tab) - self._active = "decomp" # currently shown view (falls back on decomp fail) + # 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._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 self._hex_pending_ea: int | None = None self._cur: NavEntry | None = None + self._pending_focus_name: str | None = None # token to land the cursor on + self._decomp_return: NavEntry | None = None # listing to return to from F5 + self._busy_screen: BusyScreen | None = None + self._xref_active = False # an xref gather is in flight (blocks re-entry) + 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) self._xref_focus_name: str | None = None self._dirty = False @@ -1818,13 +4520,22 @@ class IdaTui(App): fp = FunctionsPanel(id="left") fp.display = False # overlay-first: reveal the docked pane with Ctrl+B yield fp - dis = DisasmView() - dis.display = False - yield dis - yield DecompView() # pseudocode is the default view + # 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 @@ -1841,23 +4552,227 @@ class IdaTui(App): ti.display = False ti.can_focus = False yield ti + mdi = Input(id="makedata") + mdi.display = False + mdi.can_focus = False + yield mdi gi = Input(id="goto") gi.display = False gi.can_focus = False yield gi - yield Static("connecting…", id="status") - yield Footer() + # 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 — + # which is why [listing] and [pseudocode] never actually rendered. + yield Static("connecting\u2026", id="status", markup=False) def on_mount(self) -> None: + self.register_theme(IDATUI_THEME) + self.theme = IDATUI_THEME.name # Keep the hidden command input out of the focus chain until summoned. inp = self.query_one("#func-filter", Input) inp.can_focus = False # The names pane is an overlay now (Ctrl+N); focus the default code view # (pseudocode) so app bindings work before anything is opened. - self.query_one(DecompView).focus() - self._connect() + self.query_one(ListingView).focus() 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. + if self._project is not None: + ref = self._pending_load_ref() + if ref is not None: + self._ask_load_options(ref.source, label=ref.label) + return + elif self._should_ask_load_options(): + self._ask_load_options(self._open_path) + return + # Show a loading overlay immediately so a slow open/analysis (big binary) + # isn't just dead air behind empty panes; dismissed once we land. + self._loading_screen = LoadingScreen(self._loading_title()) + self.push_screen(self._loading_screen) + self._connect() + + def _should_ask_load_options(self) -> bool: + """Ask only when nobody has already answered, and only when it matters. + + Skipped when: options came from the command line or the project (the + user already said); a database exists (the answer is recorded in it, and + re-passing switches fails the open); or the file is a format IDA + recognises, which is nearly always. + """ + if not self._open_path: + return False + if self._load_args: + return False + from .formats import needs_load_options + if os.path.exists(self._open_path + ".i64") or os.path.exists( + os.path.splitext(self._open_path)[0] + ".i64"): + return False + return needs_load_options(self._open_path) + + def action_load_options(self) -> None: + """Ctrl+L: re-open this binary with different load options. + + The database IDA already built has the old processor and base baked into + it and takes precedence over any switches, so re-loading means throwing + it away. That destroys names and comments, hence the confirmation — but + for the case this exists for (a blob loaded as the wrong architecture, + which analysed to nothing) there is nothing to lose and no other way + forward. + """ + if not self._can_reload(): + 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} " + 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) + + def _on_reload_confirmed(self, yes) -> None: # type: ignore[no-untyped-def] + if not yes: + return + path, label = self._open_path, None + if self._project is not None and self._binary is not None: + 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() + self._reset_for_reload() + self._load_args = "" + if label is not None and self._project is not None: + self._project.set_load(label, processor="", base=0) + i = self._project._refs.index(self._project.by_label(label)) + self._project._entries[i].pop("processor", None) + self._project._entries[i].pop("base", None) + self._project.save() + self._pending_switch = None + self._ask_load_options(path, label=label) + + def _release_worker(self) -> None: + if self._pool is not None and self._binary is not None: + try: + self._pool.evict(self._binary, save=False) + except Exception: # noqa: BLE001 + pass + elif self.client is not None: + try: + self.client.close() + except Exception: # noqa: BLE001 + pass + 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 + self._cur = None + self._nav = [] + self._did_auto_land = False + self._pending_restore = None + self._split = False + self.query_one(DecompView).loaded_ea = None + + def _retry_load_options(self) -> None: + """Re-ask after IDA refused what we told it.""" + path, label = self._open_path, None + if self._project is not None and self._binary is not None: + ref = self._project.by_label(self._binary) + if ref is not None: + path, label = ref.source, ref.label + # Clear the rejected answer or _pending_load_ref would see the + # binary as already described and never ask again. + self._project.set_load(label, processor="", base=0) + self._project._entries[self._project._refs.index(ref)].pop( + "processor", None) + self._project.save() + self._load_args = "" + if path: + self._status("those load options were rejected \u2014 try again") + self._ask_load_options(path, label=label) + + def _pending_load_ref(self, label: str | None = None): # type: ignore[no-untyped-def] + """The project binary about to be opened, if it needs describing. + + Checked against the SOURCE: staging may not have happened yet, and the + question is about the bytes, not where they were copied to. + """ + if self._project is None: + return None + label = label or self._binary or self._project.refs[0].label + ref = self._project.by_label(label) + if ref is None or ref.load_args: + return None + 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 + from .formats import needs_load_options + return ref if needs_load_options(ref.source) else None + + def _ask_load_options(self, path: str, label: str | None = None) -> None: + try: + size = os.path.getsize(path) + except OSError: + size = 0 + self._load_for_label = label + self.push_screen(LoadOptionsScreen(path, size), self._on_load_options) + + def _on_load_options(self, choice) -> None: # type: ignore[no-untyped-def] + from .formats import load_args + choice = choice or {} + label, self._load_for_label = self._load_for_label, None + proc, base = choice.get("processor", ""), int(choice.get("base", 0) or 0) + if proc: + if label is not None and self._project is not None: + # Persist it: the answer belongs to the binary, not to this run. + self._project.set_load(label, proc, base) + else: + self._load_args = load_args(proc, base) + self._status(f"loading as {proc} @ {base:#x}") + if self._pending_switch is not None: + label2, self._pending_switch = self._pending_switch, None + self._switch_binary(label2) + return + self._loading_screen = LoadingScreen(self._loading_title()) + self.push_screen(self._loading_screen) + self._connect() + + def _loading_title(self) -> str: + return os.path.basename(self._open_path) if self._open_path else "database" + + def _dismiss_loading(self) -> None: + ls = self._loading_screen + self._loading_screen = None + if ls is not None: + try: + ls.dismiss() + except Exception: # noqa: BLE001 -- already popped + pass def _start_rpc(self) -> None: from .rpc import RpcServer @@ -1874,8 +4789,45 @@ class IdaTui(App): await self._rpc.stop() # -- status helper ----------------------------------------------------- # - def _status(self, text: str) -> None: - self.query_one("#status", Static).update(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: + self.query_one("#status", Static).update(text) + except Exception: # noqa: BLE001 -- status bar transiently unavailable + pass + if self._loading_screen is not None: # mirror progress into the overlay + self._loading_screen.update_note(text) # -- clipboard --------------------------------------------------------- # def _copy(self, text: str) -> int: @@ -1898,43 +4850,129 @@ class IdaTui(App): pass return len(text) + # -- 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.""" + from textual.worker import WorkerFailed + orig = error.error if isinstance(error, WorkerFailed) else error + if isinstance(orig, IDAConnectionError): + self._on_connection_lost() + return + super()._handle_exception(error) + + def _on_connection_lost(self) -> None: + if self._reconnecting: + return + self._reconnecting = True + self._conn_screen = LoadingScreen( + "the analysis server", note="connection lost \u2014 reconnecting\u2026") + self.push_screen(self._conn_screen) + self._reconnect() + + def _conn_note(self, text: str) -> None: + if self._conn_screen is not None: + self._conn_screen.update_note(text) + + def _dismiss_conn(self) -> None: + cs = self._conn_screen + self._conn_screen = None + if cs is not None: + try: + cs.dismiss() + except Exception: # noqa: BLE001 -- already popped (Esc) + pass + + @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. + 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.connect(progress=lambda m: self.app.call_from_thread( + self._conn_note, m)) + except Exception as e: # noqa: BLE001 + self.app.call_from_thread(self._reconnect_failed, str(e)) + return + self.app.call_from_thread(self._after_reconnect, client, Program(client)) + + def _after_reconnect(self, client: "WorkerClient", program: "Program") -> None: + self.client = client + self.program = program + self._reconnecting = False + self._dismiss_conn() + self._status("reconnected \u2014 reloading\u2026") + self._load_functions() # rebuild the function index against the new client + cur = self._cur + if cur is not None: # refresh the current view with the new program + self._open_entry(cur, push=False) + + def _reconnect_failed(self, why: str) -> None: + self._reconnecting = False + self._conn_note(f"reconnect failed: {why} \u2014 retry on next action, or 'q'") + self._status(f"reconnect failed: {why}") + # -- connection + initial load ---------------------------------------- # @work(thread=True, exclusive=True, group="connect") def _connect(self) -> None: try: - client = IDAClient(self._url, db=self._db) - client.connect() - if self._open_path is not None: - self.app.call_from_thread(self._status, f"opening {self._open_path}…") - import os - path = os.path.abspath(os.path.expanduser(self._open_path)) - # Moderate idle TTL: the keepalive heartbeat (below) keeps the - # session alive while the TUI runs; once it exits the worker - # idles out and frees its slot (avoids piling up to max-workers). - res = client.call("idb_open", input_path=path, - idle_ttl_sec=1800, timeout=1800.0) - if not (isinstance(res, dict) and res.get("success")): - err = res.get("error") if isinstance(res, dict) else res - self.app.call_from_thread(self._status, f"open failed: {err}") - return - client.set_db(res["session"]["session_id"]) - elif self._db is None: - client.set_db(client.resolve_db()) - health = client.health() - module = health.get("module", "?") + client = self._open_worker_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. + # it's reclaimed after the TUI closes. (No-op for the worker.) self._ka = client.keepalive(interval=120.0).start() program = Program(client) except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"connect failed: {e}") + self.app.call_from_thread(self._dismiss_loading) + # A load we described ourselves and IDA refused: offer the dialog + # again rather than leaving an empty app with an error in the status + # bar. Getting the processor wrong is an ordinary mistake and should + # cost one more keypress, not a restart. + if "load options" in str(e): + self.app.call_from_thread(self._retry_load_options) return self.client = client self.program = program self.app.call_from_thread(self._status, f"{module} — 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 + 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.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) + client.connect(progress=lambda m: self.app.call_from_thread( + self._status, m)) + return client + @work(thread=True, exclusive=True, group="load-funcs") def _load_functions(self) -> None: assert self.program is not None @@ -1942,7 +4980,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) @@ -1950,14 +4987,178 @@ 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_path and self._trace is None: + self._load_trace() # needs the index above: rebasing reads it + + @work(thread=True, exclusive=True, group="prewarm") + def _prewarm_provider(self) -> None: + """Warm the binary this one leans on hardest, once we're idle. + + Not "the next in the list" — phase 3 tells us something better. The + binary providing the most of this one's imports is where a follow is + most likely to take you, so paying its startup now is the switch you + would otherwise wait for. Refuses to evict anything (see pool.prewarm), + so at a tight budget this simply does nothing. + """ + if self._pool is None or self._index is None or self._binary is None: + return + try: + imps, _ = self.program.linkage() + except Exception: # noqa: BLE001 + return + if not imps: + return + from collections import Counter + votes: Counter = Counter() + for name in {i.name for i in imps}: + for h in self._index.providers(name, exclude=self._binary): + votes[h.binary] += 1 + resident = set(self._pool.resident()) + cand = next((b for b, _ in votes.most_common() if b not in resident), None) + if cand is None: + return + n = votes[cand] + try: + if self._pool.prewarm(cand): + self.app.call_from_thread( + self._status, f"pre-warmed {cand} (provides {n} imports)") + except Exception: # noqa: BLE001 -- speculative work must never surface + pass + + @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.""" + if self._index is None or self._project is None or self._binary is None: + return + ref = self._project.by_label(self._binary) + if ref is None or not self._index.is_stale(self._binary, ref.source): + return + from .index import KIND_EXPORT, KIND_FUNC, KIND_IMPORT, KIND_STRING + idx = self._func_index + entries = [(KIND_FUNC, f.addr, f.name) for f in (idx.all_loaded() if idx else [])] + try: + entries += [(KIND_STRING, s.addr, s.text) + for s in self.program.strings()] + except Exception: # noqa: BLE001 -- symbols alone are still worth indexing + pass + try: + 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 + pass + try: + n = self._index.reindex(self._binary, entries, source=ref.source) + except Exception as e: # noqa: BLE001 + self.app.call_from_thread(self._status, f"indexing failed: {e}") + return + self.app.call_from_thread( + self._status, f"indexed {self._binary}: {n} symbols, strings + linkage") + self._prewarm_provider() + + # -- initial landing --------------------------------------------------- # + #: function names tried (in order) as the startup landing spot + ENTRY_NAMES = ("main", "_main", "wmain", "WinMain", "wWinMain") + + def _auto_land(self) -> None: + """On startup, once functions are loaded and nothing is open yet, jump to + main() if it exists; else open the symbol picker so you're never staring + at a blank pane. Runs once (guarded by ``_cur``) and never steals focus + from a user who already navigated.""" + goto = self._goto_after_switch + if goto is not None: # arrived here from a project-wide search hit + self._goto_after_switch = None + self._did_auto_land = True + self._dismiss_loading() + self._goto_ea(goto, push=True) + return + entry = self._pending_restore + if entry is not None: # switched back to a binary we'd already explored + self._pending_restore = None + self._did_auto_land = True + self._dismiss_loading() + self._open_entry(entry, push=False) + return + if self._cur is not None or self._did_auto_land or self._func_index is None: + return + self._did_auto_land = True + self._dismiss_loading() # binary is up; hand the screen back to the views + fn = self._entry_func() + if fn is not None: + self._open_function(fn.addr, fn.name) + elif len(self._func_index): + # 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() + + def _land_without_functions(self) -> None: + """Analysis found nothing. Show the bytes and say so. + + Falling through to the symbol picker here left two empty panes and + "functions still loading…" — which is a lie, loading had finished. There + is always something to look at: the segments exist even when IDA + recognised no code in them, so open the listing at the start of the image. + + Zero functions is also the signal that a blob was described wrongly. It's + exactly what a good image loaded as the wrong processor looks like, so + the status says so rather than leaving you to guess. + """ + start = None + try: + regions = self.program.file_regions() + if regions: + start = regions[0][0] + except Exception: # noqa: BLE001 + pass + self._no_functions = self._can_reload() + if start is None: + self._status("no functions and no segments \u2014 nothing to show") + return + self._open_at(start, self.program.section_of(start) or "image", + cursor=0, push=True, is_region=True) + + def _can_reload(self) -> bool: + """Whether we're able to re-open this binary with different options.""" + if self._project is not None and self._binary is not None: + return True + return bool(self._open_path) + + def _entry_func(self) -> Func | None: + """The best startup landing function (exact-name match against + ENTRY_NAMES, in priority order), or None.""" + idx = self._func_index + if idx is None: + return None + by_name = {f.name: f for f in idx.all_loaded()} + for nm in self.ENTRY_NAMES: + if nm in by_name: + return by_name[nm] + return None def _module(self) -> str: try: @@ -2022,7 +5223,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 @@ -2057,8 +5258,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 DisasmView).focus() + self._focus_code_view() else: self.query_one("#func-table", DataTable).focus() @@ -2070,32 +5270,539 @@ class IdaTui(App): self.push_screen(StructEditor(self.program)) def action_symbols(self) -> None: - """Ctrl+N: fuzzy-find a symbol in a command-palette overlay.""" + """Ctrl+N: fuzzy-find a symbol (Ctrl+A widens it to the whole project).""" idx = self._func_index funcs = idx.all_loaded() if idx is not None else [] if not funcs: self._status("functions still loading…") return - self.push_screen(SymbolPalette(funcs), self._on_symbol_chosen) + self.push_screen(SymbolPalette(funcs, index=self._index, + binary=self._binary), + self._on_symbol_chosen) - def _on_symbol_chosen(self, addr: int | None) -> None: - if addr is None: + def _on_symbol_chosen(self, choice) -> None: # type: ignore[no-untyped-def] + if choice is None: + return + binary, addr = choice + if binary and binary != self._binary: + self._switch_then_goto(binary, addr) return f = self._func_index.by_addr(addr) if self._func_index else None self._open_function(addr, f.name if f else hex(addr)) + def _switch_then_goto(self, binary: str, addr: int) -> None: + """A project-wide hit in another binary: switch to it, then jump there + once its index has loaded. + + Records the hop so Esc can come back. Nav history is per-binary, so + without this a cross-binary jump — a project search hit, or following an + import into the library that implements it — is a one-way door: you + arrive in a binary whose history is empty and nothing takes you back. + """ + if self._binary is not None and binary != self._binary: + self._hops.append(self._binary) + self._goto_after_switch = addr + self._switch_binary(binary) + + # -- projects: switching between the binaries of one target ------------ # + # -- exit ---------------------------------------------------------------- # + def _dirty_labels(self) -> list[str]: + """Databases with edits that aren't on disk yet. + + A binary the pool has evicted was saved on the way out, so only the + active one and other still-resident binaries can be dirty. + """ + if self._pool is None: + return [os.path.basename(self._open_path or "database")] if self._dirty else [] + out = [self._binary] if (self._dirty and self._binary) else [] + out += [label for label, st in self._states.items() + if st.dirty and label != self._binary + and self._pool.is_resident(label)] + return out + + async def action_quit(self) -> None: + """Never drop edits on the floor: ask before exiting when a database has + unsaved changes (IDA/Ghidra behaviour).""" + dirty = self._dirty_labels() + if not dirty: + self._save_on_exit = False # nothing to write + self.exit() + return + self.push_screen(QuitScreen(dirty), self._on_quit_choice) + + def _on_quit_choice(self, choice: str | None) -> None: + if choice == "discard": + self._save_on_exit = False + self.exit() + elif choice == "save": + # Save with the overlay up: writing a big .i64 takes seconds, and + # doing it during teardown would look like a hang with no UI left. + self._loading_screen = LoadingScreen("saving", note="writing databases\u2026") + self.push_screen(self._loading_screen) + self._save_then_exit() + # None: cancel, stay put + + @work(thread=True, exclusive=True, group="save-exit") + def _save_then_exit(self) -> None: + try: + 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) + 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) + + def _finish_exit(self) -> None: + self._save_on_exit = False # already written above + self._dirty = False + self.exit() + + def action_help(self) -> None: + """F1: the keyboard cheatsheet (there's no permanent footer any more).""" + if self._prompt_active(): + return + if not isinstance(self.screen, HelpScreen): + self.push_screen(HelpScreen()) + + def action_switch_binary(self) -> None: + """Ctrl+O: pick another binary from the project (Ghidra-style).""" + if self._pool is None: + self._status("not a project — open one with --project") + return + if self._prompt_active(): + return + self.push_screen(ProjectPalette(self._pool.status()), + self._on_binary_chosen) + + def _on_binary_chosen(self, label: str | None) -> None: + if label and label != self._binary: + self._switch_binary(label) + + def _switch_binary(self, label: str) -> None: + # A blob nobody has described yet has to be described before its worker + # opens it — same as at boot, just reached by switching instead. + ref = self._pending_load_ref(label) + if ref is not None and self._load_for_label is None: + self._pending_switch = label + 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). + if self._binary is not None: + self._states[self._binary] = BinaryState( + label=self._binary, program=self.program, + func_index=self._func_index, nav=list(self._nav), cur=self._cur, + active=self._active, split=self._split, + filter_term=self._filter_term, dirty=self._dirty) + self._loading_screen = LoadingScreen(label, note="switching\u2026") + self.push_screen(self._loading_screen) + self._do_switch(label) + + @work(thread=True, exclusive=True, group="switch-binary") + def _do_switch(self, label: str) -> None: + assert self._pool is not None + try: + client = self._pool.get(label, progress=lambda m: + self.app.call_from_thread(self._status, m)) + except Exception as e: # noqa: BLE001 + 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 + # 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) + program = st.program if reuse else Program(client) + self.app.call_from_thread(self._after_switch, label, client, program, + st, reuse) + + def _after_switch(self, label, client, program, st, reuse) -> None: # type: ignore[no-untyped-def] + self.client = client + self.program = program + 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) + self._active = st.active if st else "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 + self._nav = list(st.nav) if st else [] + self._decomp_return = None + self._split_eamap, self._split_ea2line, self._split_range = [], {}, None + self.query_one(DecompView).loaded_ea = None # belongs to the old binary + if reuse and st.func_index is not None: + self._cur = st.cur + self._func_index = st.func_index + self._apply_filter(self._filter_term) # repopulate the names table + self._dismiss_loading() + if self._goto_after_switch is not None: + addr, self._goto_after_switch = self._goto_after_switch, None + self._goto_ea(addr, push=True) + elif self._cur is not None: + self._open_entry(self._cur, push=False) + else: + self._did_auto_land = False + self._auto_land() + return + # Cold (first visit, or the worker was evicted): rebuild the index, then + # land back where we were via _pending_restore. + self._cur = None + self._func_index = None + self._pending_restore = st.cur if st else None + # Landing is per-binary: without this the app-wide "already landed" flag + # from the first binary would stop the new one landing at all (and leave + # the switch overlay up forever). + self._did_auto_land = False + self._status("loading functions\u2026") + self._load_functions() + + def _switch_failed(self, label: str, why: str) -> None: + self._dismiss_loading() + self._status(f"could not open {label}: {why}") + + def action_strings(self) -> None: + """'\"' / Shift+F12: browse every string in the binary (filterable, F2 + widens to the whole project); Enter jumps to it in the unified listing.""" + if self.program is None: + self._status("not connected yet") + return + if self._prompt_active(): + return + self._status("collecting strings\u2026") + self._load_strings() + + @work(thread=True, exclusive=True, group="strings") + def _load_strings(self) -> None: + assert self.program is not None + try: + items, err = self.program.strings(), None + except Exception as e: # noqa: BLE001 -- report, never kill the app + items, err = [], str(e) + self.app.call_from_thread(self._present_strings, items, err) + + def _present_strings(self, items: list, err: str | None) -> None: + if err: + self._status(f"strings failed: {err}") + return + if not items: + self._status("no strings found (needs the list_strings tool)") + return + self._status(f"strings: {len(items)}") + self.push_screen(StringsPalette(items, index=self._index, + binary=self._binary), + self._on_string_chosen) + + def _on_string_chosen(self, choice) -> None: # type: ignore[no-untyped-def] + if choice is None: + return + binary, addr = choice + if binary and binary != self._binary: + self._switch_then_goto(binary, addr) + 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).""" + # Tab is a PRIORITY app binding, so it fires even while a modal is up and + # nothing inside a dialog could ever be tabbed to. Hand it back to the + # dialog: this is the only reason the load dialog's address field was + # unreachable, and it was broken the same way in every other modal. + if self.screen is not self.screen_stack[0]: + try: + self.screen.focus_next() + except Exception: # noqa: BLE001 -- screen with nothing focusable + pass + return if self._cur is None: 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" + 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": - self._active = self._pref + self._active = self._code_mode() + self._show_active() + return + if self._active == "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 = "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._active == "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() + if ea is None: + self._status("no address here to decompile") + return + # Raise the pseudocode pane + 'decompiling…' overlay NOW, before the + # background decompile runs — otherwise the (cold) decompile happens in + # _decomp_from_listing first and _show_active's re-decompile is cached, + # so the overlay only flashes for an instant. + dec = self.query_one(DecompView) + self.query_one(ListingView).display = False + self.query_one(HexView).display = False + dec.display = True + dec.loading = True + dec.focus() + self._status("decompiling\u2026") + self._decomp_from_listing(ea) + return + # active == "decomp": F5/Tab returns to the linear listing. + ret = self._decomp_return + self._decomp_return = None + if ret is not None: + self._cur = ret + self._active = "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 + # in the listing at the current pseudocode line's address. Reuse the + # entry (it is _nav[-1]) rather than spawning a detached one, so cursor + # moves keep updating it and a later edit reloads at the right spot. + dec = self.query_one(DecompView) + 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._show_active() + + @work(thread=True, group="nav") + def _toggle_to_listing(self, ea: int) -> None: + assert self.program is not None + lm = self.program.listing(ea) + idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 + self.app.call_from_thread(self._apply_toggle_listing, idx) + + def _apply_toggle_listing(self, idx: int) -> None: + cur = self._cur + if cur is None: + self._active = "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._open_entry(cur, push=False) + + @work(thread=True, group="nav") + def _decomp_from_listing(self, ea: int) -> None: + assert self.program is not None + fn = self.program.function_of(ea) + if fn is None: + self.app.call_from_thread( + self._decomp_from_listing_failed, + "F5 — cursor is not inside a defined function ('p' to make one)") + return + dec_idx = self._decomp_line_for(fn.addr, ea) + self.app.call_from_thread(self._enter_decomp, fn.addr, fn.name, dec_idx) + + def _decomp_from_listing_failed(self, msg: str) -> None: + # 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._show_active() + self._status(msg) + + def _enter_decomp(self, fn_addr: int, fn_name: str, dec_idx: int) -> None: + # Snapshot the listing position so F5/Tab returns exactly here. + lst = self.query_one(ListingView) + ret = self._cur + if ret is not None: + ret.cursor = lst.cursor + ret.cursor_x = lst.cursor_x + ret.scroll_y = round(lst.scroll_offset.y) + self._decomp_return = ret + entry = NavEntry(ea=fn_addr, name=fn_name, is_region=False, + dec_cursor=max(dec_idx, 0)) + self._cur = entry + self._active = "decomp" + self._show_active() + + def action_toggle_split(self) -> None: + """Toggle the side-by-side listing ⇄ pseudocode view (Ghidra-style).""" + if self._prompt_active(): + return # a search/rename/… prompt owns the keyboard + if self._cur is None or self.program is None: + self._status("open a function first") + return + if not self._split and self.size.width < _SPLIT_MIN_WIDTH: + self._status(f"terminal too narrow for split — need ≈{_SPLIT_MIN_WIDTH} " + f"cols (have {self.size.width})") + return + self._split = not self._split + if self._active not in ("listing", "decomp"): + self._active = "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._active == "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 = "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._active in ("listing", "disasm"): + return self.query_one(ListingView)._cursor_ea() + if self._active == "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 self._active != "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 - self._active = "decomp" if self._active == "disasm" else "disasm" - self._pref = self._active # an explicit toggle sets the preference + 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 = "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. + assert self.program is not None + lm = self.program.listing(ea) + idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 + self.app.call_from_thread(self._apply_enter_split, lm, ea, name, idx) + + def _apply_enter_split(self, lm, ea: int, name: str, idx: int) -> None: # type: ignore[no-untyped-def] + if not self._split: + return # toggled back off before the listing finished loading + lst = self.query_one(ListingView) + if lm is not None: + lst.load(lm, name, cursor=idx, scroll_y=max(idx - _JUMP_CONTEXT, 0)) + self._show_active() # split branch shows both + loads the decomp + self._sync_split(self._active) # crude link now + self._load_split_map(ea) # region map (async) if decomp is loaded def action_hex(self) -> None: """Backslash: show the raw bytes of the loaded image, synced to the code @@ -2103,11 +5810,13 @@ class IdaTui(App): if self._cur is None or self.program is None: return if self._active == "hex": - self._active = self._pref + self._active = self._code_mode() self._show_active() return - if self._active == "disasm": - ea = self.query_one(DisasmView)._cursor_ea() + if self._active in ("listing", "disasm"): + ea = self.query_one(ListingView)._cursor_ea() + elif self._active == "graph": + ea = self.query_one(GraphView)._cursor_ea() else: dec = self.query_one(DecompView) ea = dec._line_ea(dec.cursor) @@ -2145,13 +5854,23 @@ class IdaTui(App): self.clear_filter() # Esc on the list clears an active filter return if len(self._nav) > 1: + self._nav_seq += 1 # invalidate any navigation still in flight self._nav.pop() - self._open_entry(self._nav[-1], push=False) + # Say something immediately: going back can need a decompile, and + # until it lands the panes still show where you were — with no + # feedback Esc reads as a no-op. + dest = self._nav[-1] + self._status(f"\u25c2 back to {dest.name}\u2026") + self._open_entry(dest, push=False) + elif self._hops: + # Local history is spent, but we got here from another binary. + label = self._hops.pop() + self._status(f"\u25c2 back to {label}\u2026") + self._switch_binary(label) # _states restores its nav and position elif self.query_one("#left", FunctionsPanel).display: table.focus() else: - self.query_one(DecompView if self._active == "decomp" - else DisasmView).focus() + self._focus_code_view() # -- input submit (filter / goto) ------------------------------------- # def on_search_requested(self, msg: SearchRequested) -> None: @@ -2169,28 +5888,70 @@ 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() + 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() - # 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) + # _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: + 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. - nxt = view.model.cached_line(view.cursor + 1) if view.model else None - self._xrefs_disasm(ea, word, ea, nxt.ea if nxt else None) + self._push_busy("finding xrefs\u2026") + self._xrefs_disasm(ea, word, ea, view._next_ea()) elif isinstance(view, DecompView) and view._texts: here = view._line_ea(view.cursor) end = None @@ -2200,8 +5961,28 @@ class IdaTui(App): if e is not None and e > here: end = e break + self._push_busy("finding xrefs\u2026") self._xrefs_decomp(view._texts[view.cursor], word, here, end) + # -- blocking "busy" overlay for short async steps -------------------- # + def _push_busy(self, message: str) -> None: + self._xref_active = True + self._busy_screen = BusyScreen(message) + self.push_screen(self._busy_screen, self._on_busy_closed) + + def _on_busy_closed(self, result: object = None) -> None: + self._busy_screen = None + if result == "cancel": # user hit Esc while we were still gathering + self._xref_active = False + + def _dismiss_busy(self) -> None: + bs = self._busy_screen + if bs is not None: + try: + bs.dismiss() # no result -> _on_busy_closed leaves the flag alone + except Exception: # noqa: BLE001 -- already popped + pass + @staticmethod def _looks_like_symbol(word: str | None) -> bool: # A symbol has a letter but is not a bare hex number (the address column @@ -2217,7 +5998,8 @@ class IdaTui(App): # Prefer the symbol under the cursor (handles multiple refs on a line). if self._looks_like_symbol(word): try: - self._do_navigate(self.program.resolve(word), push=True) + self._do_navigate(self.program.resolve(word), push=True, + focus_name=word) return except Exception: # noqa: BLE001 -- not a resolvable name; fall back pass @@ -2231,8 +6013,55 @@ class IdaTui(App): if tgt is None or tgt.to is None: self.app.call_from_thread(self._status, "nothing to follow here") return + if self._follow_import(tgt.to): + return self._do_navigate(tgt.to, push=True) + def _import_stub(self, ea: int) -> str | None: + """The import name at ``ea``, if ``ea`` is one of this binary's import + stubs. That's the dead end phase 3 exists to open up: a call to strcmp + reaches the PLT/extern entry and Hex-Rays has nothing to decompile, + because the code lives in a library this binary only references.""" + try: + imps, _ = self.program.linkage() + except Exception: # noqa: BLE001 + return None + for i in imps: + if i.addr == ea: + return i.name + return None + + 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 + evicted — the whole reason the index exists. + """ + if self._index is None or self._project is None or not name: + return None + try: + hits = self._index.providers(name, exclude=self._binary) + except Exception: # noqa: BLE001 + return None + return (hits[0].binary, hits[0].addr) if hits else None + + def _follow_import(self, ea: int) -> bool: + """Follow an import stub into the binary that implements it. True when + it was handled (caller must not also navigate locally).""" + name = self._import_stub(ea) + if not name: + return False + found = self._cross_binary_impl(name) + if found is None: + # Leave the local navigation alone: landing on the stub is still the + # honest answer when nothing in the project provides the symbol. + return False + label, addr = found + self.app.call_from_thread( + self._status, f"{name} \u2192 {label} (import resolved)") + self.app.call_from_thread(self._switch_then_goto, label, addr) + return True + @work(thread=True, group="nav") def _follow_decomp(self, line: str, word: str | None, line_ea: int | None = None) -> None: @@ -2260,7 +6089,11 @@ class IdaTui(App): if addr is None: self.app.call_from_thread(self._status, "nothing to follow on this line") return - self._do_navigate(addr, push=True) + if self._follow_import(addr): + return + # Following FROM pseudocode: keep the reader in the decompiler when the + # target is decompilable, instead of dropping to the linear listing. + self._do_navigate(addr, push=True, prefer_decomp=True) def _ref_on_line(self, line: str) -> int | None: """Address of the first decompiler ref whose name appears on ``line``.""" @@ -2307,6 +6140,8 @@ class IdaTui(App): subj = self._ref_on_line(line) or (self._cur.ea if self._cur else None) if subj is not None: self._xrefs_present(subj, word, here_ea, here_end) + else: # nothing to xref -> still clear the busy overlay + self.app.call_from_thread(self._present_xrefs, "xrefs", [], None, 0) @staticmethod def _xref_preselect(xr, here_ea, here_end) -> int: # type: ignore[no-untyped-def] @@ -2326,6 +6161,14 @@ class IdaTui(App): here_ea: int | None = None, here_end: int | None = None) -> None: # worker assert self.program is not None + try: + return self._xrefs_present_inner(subj, subj_name, here_ea, here_end) + except Exception as e: # noqa: BLE001 -- never leave the busy overlay stuck + self.app.call_from_thread(self._status, f"xrefs failed: {e}") + self.app.call_from_thread(self._present_xrefs, "xrefs", [], None, 0) + + def _xrefs_present_inner(self, subj, subj_name, here_ea, here_end): # type: ignore[no-untyped-def] + assert self.program is not None xr = self.program.xrefs_to(subj) fn = self.program.function_of(subj) if fn is not None: @@ -2350,12 +6193,53 @@ class IdaTui(App): # not inside a function: name the section it lives in (a GOT/reloc # data slot or a loose thunk) instead of a bare '?'. loc = self.program.section_of(x.frm) or "<no seg>" - items.append((x.frm, f"{x.frm:08X} {loc:<26} [{x.type}]")) + kind = x.kind or x.type or "?" + items.append((x.frm, f"{x.frm:08X} {kind:<6} {loc}")) preselect = self._xref_preselect(xr, here_ea, here_end) + # Callers in OTHER project binaries. xrefs_to only ever sees this + # database, so an exported function looks unused from the inside even + # when half the project calls it — the import side of the linkage index + # is the only place that knowledge exists. + for label_b, addr_b, text_b in self._foreign_importers(subj, subj_name, fn): + items.append(((label_b, addr_b), text_b)) self.app.call_from_thread(self._present_xrefs, label, items, focus, preselect) - def _present_xrefs(self, label: str, items: list[tuple[int, str]], + 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. + + 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. + """ + if self._index is None or self._project is None or self._binary is None: + return [] + name = subj_name if self._looks_like_symbol(subj_name) else None + if name is None and fn is not None and fn.addr == subj: + name = fn.name + if not name: + return [] + from .domain import link_name + name = link_name(name) + try: + _, exports = self.program.linkage() + except Exception: # noqa: BLE001 + return [] + if not any(e.name == name for e in exports): + return [] # we don't export it; nobody imports it FROM US + try: + hits = self._index.importers(name, exclude=self._binary) + except Exception: # noqa: BLE001 + return [] + return [(h.binary, h.addr, f"{h.addr:08X} import [{h.binary}] {name}") + for h in hits] + + def _present_xrefs(self, label: str, items: list[tuple[object, str]], focus_name: str | None = None, preselect: int = 0) -> None: + if not self._xref_active: + return # cancelled (Esc) while we were still gathering + self._xref_active = False + self._dismiss_busy() if not items: self._status(f"{label}: none") return @@ -2363,9 +6247,17 @@ class IdaTui(App): self._status(f"{label}: {len(items)}") self.push_screen(XrefsScreen(label, items, preselect), self._on_xref_chosen) - def _on_xref_chosen(self, addr: int | None) -> None: - if addr is not None: - self._goto_ea(addr, push=True, focus_name=self._xref_focus_name) + def _on_xref_chosen(self, addr) -> None: # type: ignore[no-untyped-def] + if addr is None: + return + if isinstance(addr, tuple): # a caller in another project binary + binary, ea = addr + self._switch_then_goto(binary, ea) # records a hop, so Esc returns + return + # 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")) # -- rename ------------------------------------------------------------ # @staticmethod @@ -2383,6 +6275,41 @@ class IdaTui(App): return re.search(rf"\bgoto\s+{re.escape(name)}\b", body) is not None 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, GraphView)): + 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 @@ -2392,6 +6319,7 @@ class IdaTui(App): "(Hex-Rays goto labels aren't renamable via the API)") 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" @@ -2404,6 +6332,7 @@ class IdaTui(App): 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: @@ -2412,7 +6341,7 @@ class IdaTui(App): # -- 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): + if isinstance(view, ListingView): return view._cursor_ea() if isinstance(view, DecompView): return view._line_ea(view.cursor) @@ -2479,16 +6408,53 @@ class IdaTui(App): return self.app.call_from_thread(self._after_comment, ea, text) - def _after_comment(self, ea: int, text: str) -> None: + 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) + + 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 block caches are cleared). + # force_recompiled lazily; disasm/listing caches are cleared). self.program.bump_names() - if cur is not None: - self._save_current_pos() - self.query_one(DecompView).loaded_ea = None # force pseudocode reload - self._open_entry(cur, push=False) + 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)") @@ -2513,7 +6479,10 @@ class IdaTui(App): 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 function under the cursor (self or a referenced one) + # 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) @@ -2523,6 +6492,12 @@ class IdaTui(App): 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 @@ -2531,6 +6506,13 @@ class IdaTui(App): return self.app.call_from_thread(self._open_retype, view, kind, subject, word or "", prefill) + @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 *") + 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) @@ -2557,6 +6539,8 @@ class IdaTui(App): 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: @@ -2568,15 +6552,66 @@ class IdaTui(App): # 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() - cur = self._cur - if cur is not None: - self._save_current_pos() - self.query_one(DecompView).loaded_ea = None - self._open_entry(cur, push=False) + self._reload_active_code() self._dirty = True what = "prototype" if kind == "func" else f"'{word}'" self._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 {1: "unsigned __int8", 2: "unsigned __int16", + 4: "unsigned __int32", 8: "unsigned __int64"}.get( + sz, f"char[{sz}]" if sz > 0 else "unsigned __int8") + + 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 _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() + + @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="rename") def _do_rename(self, view, old: str, new: str) -> None: # type: ignore[no-untyped-def] assert self.program is not None @@ -2634,10 +6669,7 @@ class IdaTui(App): # 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() - if cur is not None: - self._save_current_pos() - self.query_one(DecompView).loaded_ea = None # force pseudocode reload - self._open_entry(cur, push=False) + 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: @@ -2654,6 +6686,654 @@ class IdaTui(App): self._dirty = True self._status(f"renamed {old} → {new} (Ctrl+S to save)") + @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() + # 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. + fn = None + try: + fn = self.program.function_of(addr) + except Exception: # noqa: BLE001 + fn = None + is_func_start = fn is not None and fn.addr == addr + lm = self.program.listing(addr) + label = name if is_func_start else 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, + is_func_start) + + def _open_at_named(self, label: str, addr: int, idx: int, name: str, + is_func_start: bool = False) -> None: + if is_func_start: + self.program.bump_names() + if self._func_index is not None: + self._func_index.update_name(addr, name) + for e in self._nav: + if e.ea == addr: + e.name = name + try: + table = self.query_one("#func-table", DataTable) + name_col = list(table.columns.keys())[1] + table.update_cell(str(addr), name_col, name) + except Exception: # noqa: BLE001 -- row filtered out / not streamed + pass + self._open_at(addr, label, idx, False, -1, 0, True) + self._dirty = True + self._status(f"named {addr:#x} → {name} (Ctrl+S to save)") + + # -- literal display formats (IDA 'o') -------------------------------- # + def on_op_format_requested(self, msg: OpFormatRequested) -> None: + if self.program is None or self._cur is None: + return + view = msg.view + if isinstance(view, ListingView): + ea = view._cursor_ea() + if ea is None: + self._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. + self._status("no literal on this line to reformat", priority=True) + return + self._do_op_format(msg.mode, "listing", ea, view.op_col()) + return + if isinstance(view, 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 self._cur.ea + self._do_op_format(msg.mode, "decomp", fn, view.cursor_x, view.cursor) + + @work(thread=True, exclusive=True, group="opformat") + def _do_op_format(self, mode: str, where: str, ea: int, col: int, + line: int = -1) -> None: # worker context + assert self.program is not None + try: + if where == "listing": + r = self.program.op_format(ea, mode=mode, col=col) + what = f"op{r.get('n', 0)} " + else: + r = self.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: + self.app.call_from_thread(self._status, f"format: {e.message}", + True) + return + except Exception as e: # noqa: BLE001 -- surface transport failures too + self.app.call_from_thread(self._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. + self.app.call_from_thread( + self._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))) + self.app.call_from_thread(self._after_op_format, desc, keep) + + def _after_op_format(self, desc: str, keep: tuple | None = None) -> None: + # 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. + self.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": + self.query_one(ListingView)._pending_op = keep[1] + else: + self.query_one(DecompView)._keep_lit = (keep[1], keep[2]) + self._reload_active_code() + self._dirty = True + self._status(f"{desc} (Ctrl+S to save)", priority=True) + + # -- 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, 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()) + + @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", + "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 = 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 == "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 = self.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 = self.program.set_thumb(ea) + run = self.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 = self.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 = self.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 + 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 _edit_done(self, anchor: ViewAnchor) -> None: + """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. + """ + self._dirty = True + if anchor.flash: + self._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. + self._reindex_functions() + + def _load_trace(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). + """ + from .trace import Trace + path = self._trace_path + try: + def note(n): + self.app.call_from_thread( + self._status, f"trace: {n:,} instructions\u2026") + trace = Trace.load(path, progress=note) + except OSError as e: + self.app.call_from_thread(self._status, f"trace: {e}") + return + if not trace.length: + self.app.call_from_thread( + self._status, f"trace: {os.path.basename(path)} is empty") + return + idx = self._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)) + self.app.call_from_thread(self._trace_ready, trace, slide, hit) + + def _trace_ready(self, trace, slide: int, hit: int) -> None: + self._trace = trace + self._t = 0 + dock = self.query_one(TraceDock) + dock.display = True + dock.show(trace, 0) + where = (f"rebased {slide:+#x}" if slide else "no rebase needed") + self._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.""" + 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. + self._nav_seq += 1 + self._t = max(0, min(int(idx), t.length - 1)) + self.query_one(TraceDock).show(t, self._t) + self._paint_trail() + if not follow: + return + pc = t.ip(self._t) + if self._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. + self._goto_ea(pc, push=False, + prefer_decomp=(self._active == "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. + """ + lst = self.query_one(ListingView) + dec = self.query_one(DecompView) + 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_trace_line = pc + self._resync_decomp_async(pc) + return True + self._place_decomp_at(pc) + self._sync_split(self._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. + """ + dec = self.query_one(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 = self._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) + + @work(thread=True, exclusive=True, group="split-resync") + def _resync_decomp_async(self, ea: int) -> None: + self._resync_decomp(ea) + + 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. + """ + t = self._trace + if t is None: + return + try: + hx = self.query_one(HexView) + hx.trace, hx.trace_idx = t, self._t + if hx.display: + hx.refresh() + except Exception: # noqa: BLE001 -- not mounted yet + pass + trail = t.trail(self._t) + try: + lst = self.query_one(ListingView) + lst.trail = trail + lst.refresh() + except Exception: # noqa: BLE001 -- view not mounted yet + pass + 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. + """ + try: + dec = self.query_one(DecompView) + except Exception: # noqa: BLE001 + return + ea = dec.loaded_ea + if not dec.display or ea is None or self.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: + self._apply_split_map(ea, self.program.decomp_map(ea)) + except Exception: # noqa: BLE001 + 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_trace_line = self._pending_trace_line, None + if pend is not None and self._split: + # The function was still decompiling when the step happened; land + # now that its line map exists. + self._place_decomp_at(pend) + self._sync_split(self._active) + + def _step(self, delta: int) -> None: + if self._trace is None: + self._status("no trace loaded (--trace FILE)") + return + self._seek(self._t + delta) + + 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?"). + """ + t = self._trace + if t is None: + self._status("no trace loaded (--trace FILE)") + return + if self._active == "hex": + hx = self._try_view(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 = self._active_code_view() + if isinstance(view, 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: + self._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: + self._status("no address on this line", priority=True) + return + stamps = list(t.executions(ea)) + what = f"execution of {ea:#x}" + if not stamps: + self._status(f"no {what} in this trace", priority=True) + return + import bisect as _b + if direction > 0: + i = _b.bisect_right(stamps, self._t) + else: + i = _b.bisect_left(stamps, self._t) - 1 + if not (0 <= i < len(stamps)): + edge = "last" if direction > 0 else "first" + self._status(f"already at the {edge} {what} " + f"({len(stamps)} in the trace)", priority=True) + return + self._seek(stamps[i]) + self._status(f"{what}: {i + 1} of {len(stamps)} @ t={stamps[i]:,}", + priority=True) + + def action_seek_next_hit(self) -> None: + self._seek_hit(1) + + def action_seek_prev_hit(self) -> None: + self._seek_hit(-1) + + def action_seek_reg_write(self) -> None: + """W: which instruction set each register to its current value.""" + t = self._trace + if t is None: + self._status("no trace loaded (--trace FILE)") + return + 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: + self.push_screen(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)) + + def action_step_fwd(self) -> None: + self._step(1) + + def action_step_back(self) -> None: + self._step(-1) + + 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. + """ + t = self._trace + if t is None: + self._status("no trace loaded (--trace FILE)") + return + 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 action_step_over_fwd(self) -> None: + self._step_over(1) + + def action_step_over_back(self) -> None: + self._step_over(-1) + + @work(thread=True, exclusive=True, group="load-funcs") + def _reindex_functions(self) -> None: + """Rebuild the function index in place after an edit changed it. + + 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. + """ + 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 @@ -2674,32 +7354,81 @@ class IdaTui(App): # -- navigation to an arbitrary address ------------------------------- # @work(thread=True, group="nav") def _goto_ea(self, ea: int, push: bool = True, - focus_name: str | None = None) -> None: - self._do_navigate(ea, push, focus_name) + focus_name: str | None = None, prefer_decomp: bool = False) -> None: + self._do_navigate(ea, push, focus_name, prefer_decomp) - def _do_navigate(self, ea: int, push: bool, - focus_name: str | None = None) -> None: # worker context + def _do_navigate(self, ea: int, push: bool, focus_name: str | None = None, + prefer_decomp: bool = False) -> None: # worker context assert self.program is not None + # Which navigation this is. Decompiling below can take a while, and if + # you press Esc (or jump again) in the meantime this result is stale — + # applying it anyway silently undoes what you just did. + seq = self._nav_seq fn = self.program.function_of(ea) - if fn is None: - self.app.call_from_thread(self._status, f"no function contains {ea:#x}") - return - model = self.program.disasm(fn.addr, fn.name) - idx = 0 if ea == fn.addr else model.index_of_ea(ea) - # The disasm cursor alone doesn't position the pseudocode pane. For a - # mid-function target with the decompiler active, resolve the address to - # its pseudocode line (via the per-line /*0xEA*/ markers) so the jump - # lands on the reference there too (e.g. selecting an xref). A plain - # function-entry jump is line 0 in both views -> skip the decompile. - # With a focus_name (the symbol the xref was on), also land the column - # on that token where it appears in the line. - dec_idx, dec_col = -1, 0 - if self._active == "decomp" and ea != fn.addr: - dec_idx = self._decomp_line_for(fn.addr, ea) - if dec_idx >= 0 and focus_name: - dec_col = self._decomp_col_for(fn.addr, dec_idx, focus_name) + # Jumping FROM the decompiler: stay in pseudocode when the target is a + # decompilable function, landing on the line that matches ``ea``. + if prefer_decomp and fn is not None: + dec = self.program.decompile(fn.addr) + if not dec.failed and dec.code: + if ea == fn.addr: + # Jumping to the function itself: land on its name in the + # prototype (line 0), not column 0. + dec_idx = 0 + col = self._decomp_col_for(fn.addr, 0, fn.name) + else: + # A mid-function site (e.g. an xref jump to a call site): + # anchor on the address but snap to the nearest line that + # actually holds the referenced symbol (the marker line and + # the symbol's line can differ), landing on the token. + dec_idx, col = self._decomp_locate(fn.addr, ea, + focus_name or fn.name) + self.app.call_from_thread( + self._open_decomp_entry, fn.addr, fn.name, dec_idx, col, + push, seq) + return + # Otherwise: everything opens the one continuous listing at ``ea``. A + # function name is used for the status label; a region gets a segment + # label. F5/Tab decompiles the function under the cursor from here. + lm = self.program.listing(ea) + 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, fn.addr, fn.name, idx, push, dec_idx, dec_col) + 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, + seq: int | None = None) -> None: + """Open ``fn_addr`` in the decompiler as a real navigation (nav history + aware), landing on pseudocode line ``dec_idx`` column ``dec_cursor_x``. + + ``seq`` is the navigation this result belongs to; if the user has + navigated since (Esc, another jump), the result is stale and dropped — + otherwise a slow decompile lands afterwards and undoes their Esc. + """ + if seq is not None and seq != self._nav_seq: + return + if push: + # 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: + dv = self.query_one(DecompView) + src = self._cur + src.view = "decomp" + src.dec_cursor = dv.cursor + src.dec_cursor_x = dv.cursor_x + src.dec_scroll_y = round(dv.scroll_offset.y) + if not self._nav or self._nav[-1] is not src: + self._push_nav(src) # never stack a second copy of a spot + else: + self._save_current_pos() + self._decomp_return = None # a real navigation abandons the F5 return + entry = NavEntry(ea=fn_addr, name=fn_name, view="decomp", + dec_cursor=dec_idx, dec_cursor_x=dec_cursor_x) + if push: + self._push_nav(entry) + self._open_entry(entry, push=False) def _decomp_col_for(self, fn_addr: int, line_idx: int, name: str) -> int: """Column of ``name`` (whole-word) on pseudocode line ``line_idx``, so an @@ -2717,6 +7446,40 @@ class IdaTui(App): m = re.search(rf"\b{re.escape(name)}\b", clean) return m.start() if m else 0 + def _decomp_locate(self, fn_addr: int, ea: int, + token: str | None) -> tuple[int, int]: + """Best (line, column) for address ``ea`` in ``fn_addr``'s pseudocode. + + Anchors on the /*0xEA*/ marker line for ``ea``, but Hex-Rays can attribute + an address to a line a step off from where the referenced symbol actually + appears. So when ``token`` is given, snap to the whole-word occurrence of + it NEAREST the anchor line (preferring the anchor line itself, then the + line just below), which disambiguates repeated symbols by address.""" + assert self.program is not None + try: + dec = self.program.decompile(fn_addr) + except Exception: # noqa: BLE001 + return (0, 0) + lines = (dec.code or "").splitlines() + if not lines: + return (0, 0) + anchor = self._decomp_line_for(fn_addr, ea) + anchor = anchor if anchor >= 0 else 0 + if not token: + return (anchor, 0) + pat = re.compile(rf"\b{re.escape(token)}\b") + best: tuple[int, int] | None = None + best_key: tuple[int, int] | None = None + for i, ln in enumerate(lines): + m = pat.search(_ADDR_MARK_STRIP_RE.sub("", ln)) + if not m: + continue + # rank: nearest to the anchor; tie -> the line at/after the anchor. + key = (abs(i - anchor), 0 if i >= anchor else 1) + if best_key is None or key < best_key: + best_key, best = key, (i, m.start()) + return best if best is not None else (anchor, 0) + def _decomp_line_for(self, fn_addr: int, ea: int) -> int: """Pseudocode line index best matching address ``ea``: the line whose /*0xEA*/ marker is the largest address <= ``ea``. -1 if unavailable @@ -2736,16 +7499,104 @@ class IdaTui(App): best_ea, best_idx = e, i return best_idx + @staticmethod + def _same_spot(a: NavEntry, b: NavEntry) -> bool: + """Same function, same view, same line — i.e. going from a to b is not a + move, so it has no business being a step in the history.""" + if a.ea != b.ea or a.view != b.view: + return False + return (a.dec_cursor == b.dec_cursor if a.view == "decomp" + else a.cursor == b.cursor) + + def _push_nav(self, entry: NavEntry) -> None: + """Append to the nav stack unless that would duplicate where we already are. + + Opening the place you're standing on — the app auto-lands on main, then + you pick main out of Ctrl+N — used to append an identical entry. The Esc + it created popped the stack without changing anything on screen: a dead + keypress, which is exactly what "back doesn't work" feels like. Replace + the top instead, so the newer entry's metadata still wins. + """ + top = self._nav[-1] if self._nav else None + if top is not None and self._same_spot(top, entry): + self._nav[-1] = entry + return + self._nav.append(entry) + + # -- view state across a model rebuild --------------------------------- # + def _anchor(self, flash: str | None = None) -> ViewAnchor: + """Capture where we're looking, BEFORE an edit rebuilds the model. + + Must run on the UI thread: it reads live widget state. + """ + a = ViewAnchor(view=self._active, flash=flash) + view = self._active_code_view() + model = getattr(view, "model", None) + if view is None or model is None: + return a + a.cursor_x = getattr(view, "cursor_x", 0) + try: + a.ea = view._cursor_ea() + except Exception: # noqa: BLE001 + a.ea = None + top = round(view.scroll_offset.y) + h = model.cached_line(top) or model.get(top) + a.top_ea = getattr(h, "ea", None) + return a + + @staticmethod + def _anchor_rows(a: ViewAnchor, model, fallback_ea: int | None = None): + """(cursor_row, top_row) for ``a`` in a freshly built ``model``. + + -1 means "no opinion" — the caller's own default wins. + """ + if model is None: + return (-1, -1) + + def row_of(ea): + if ea is None: + return -1 + try: + i = model.index_of_ea(ea) + except Exception: # noqa: BLE001 + return -1 + return i if i >= 0 else -1 + + cur = row_of(a.ea if a.ea is not None else fallback_ea) + if cur < 0 and fallback_ea is not None: + 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) -> None: + dec_cursor: int = -1, dec_cursor_x: int = 0, + is_region: bool = False, focus_name: str | None = None, + scroll_y: int = -1) -> None: if push: self._save_current_pos() - entry = NavEntry(ea=ea, name=name, cursor=cursor) + self._decomp_return = None # a real navigation abandons the F5 return + # Land the cursor on this token (e.g. the ref an xref jump targets) once + # the row is loaded, instead of column 0. + self._pending_focus_name = focus_name + entry = NavEntry(ea=ea, name=name, cursor=cursor, is_region=is_region) + if scroll_y >= 0: + entry.scroll_y = scroll_y if dec_cursor >= 0: entry.dec_cursor = dec_cursor entry.dec_cursor_x = dec_cursor_x if push: - self._nav.append(entry) + self._push_nav(entry) self._open_entry(entry, push=False) def on_input_changed(self, event: Input.Changed) -> None: @@ -2772,6 +7623,8 @@ 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: @@ -2794,12 +7647,17 @@ class IdaTui(App): event.prevent_default() self._end_retype() return + if self.query_one("#makedata", Input).display: + 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() + else (self._code_view() or self.query_one(ListingView))).focus() return fi = self.query_one("#func-filter", Input) if fi.display: @@ -2826,7 +7684,12 @@ class IdaTui(App): 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 @@ -2842,10 +7705,16 @@ class IdaTui(App): 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()) + return if inp.id == "goto": self._end_goto() (self.query_one(HexView) if self._active == "hex" - else self._code_view()).focus() + else (self._code_view() or self.query_one(ListingView))).focus() if value: self._goto(value) return @@ -2854,7 +7723,72 @@ class IdaTui(App): self.query_one("#func-table", DataTable).focus() def _code_view(self): # type: ignore[no-untyped-def] - return self.query_one(DecompView if self._pref == "decomp" else DisasmView) + """The code pane the user is in — where focus belongs after a prompt closes. + + This used to pick on _pref, which was always "listing", so closing the + goto prompt while reading pseudocode threw focus into the listing. + """ + return self._active_code_view() + + def _active_code_view(self): # type: ignore[no-untyped-def] + """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._active in ("listing", "disasm"): + return self.query_one(ListingView) + if self._active == "decomp": + return self.query_one(DecompView) + if self._active == "graph": + return self.query_one(GraphView) + return 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() + if view is None: + self._status("open a function first") + return + view.focus() + fn = getattr(view, f"action_{name}", None) + if fn is None: + self._status(f"'{name}' isn't available in this view") + return + fn(*args) + + def action_continuous_here(self) -> None: + """'L': open the continuous segment listing at the cursor — one long + flat view where functions, data and undefined bytes are interleaved, + instead of the function-bounded disassembly.""" + if self.program is None: + return + view = self._active_code_view() + ea = self._line_ea_for(view) if view is not None else None + if ea is None and self._cur is not None: + ea = self._cur.ea + if ea is None: + self._status("no address here to open the continuous listing") + return + if self._cur is not None and self._cur.is_region: + self._status("already in the continuous listing") + return + self._goto_continuous(ea) + + @work(thread=True, group="nav") + def _goto_continuous(self, ea: int, push: bool = True) -> None: + assert self.program is not None + 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 + self.app.call_from_thread(self._open_at, ea, name, idx, push, -1, 0, True) @work(thread=True, exclusive=True, group="goto") def _goto(self, target: str) -> None: @@ -2867,6 +7801,13 @@ class IdaTui(App): if self._active == "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._active == "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) @@ -2891,61 +7832,140 @@ class IdaTui(App): self._open_function(ea, name) def _save_current_pos(self) -> None: - """Snapshot the active views' cursor + scroll into the top history entry + """Snapshot the listing cursor + scroll into the top history entry (called just before we navigate away).""" if not self._nav: return e = self._nav[-1] - dis = self.query_one(DisasmView) - e.cursor, e.cursor_x = dis.cursor, dis.cursor_x - e.scroll_y = round(dis.scroll_offset.y) - dec = self.query_one(DecompView) - if dec.loaded_ea == e.ea: - e.dec_cursor, e.dec_cursor_x = dec.cursor, dec.cursor_x - e.dec_scroll_y = round(dec.scroll_offset.y) - e.dec_scroll_x = round(dec.scroll_offset.x) + lst = self.query_one(ListingView) + if lst.model is not None: + e.cursor, e.cursor_x = lst.cursor, lst.cursor_x + e.scroll_y = round(lst.scroll_offset.y) - def _open_function(self, ea: int, name: str, push: bool = True) -> None: - if push: - self._save_current_pos() - entry = NavEntry(ea=ea, name=name, cursor=0) - if push: - self._nav.append(entry) - self._open_entry(entry, push=False) + @work(thread=True, group="nav") + def _open_function(self, ea: int, name: str | None = None, + push: bool = True) -> None: + # Unified: opening a function is just navigating the one linear listing + # to its entry address. + self._do_navigate(ea, push) + + + def _code_mode(self) -> str: + """The code view to return to from hex — always the unified listing.""" + return "listing" def _open_entry(self, entry: NavEntry, push: bool) -> None: if self.program is None: return self._cur = entry - # Each open honours the preferred view; a decomp failure last time fell - # back to disasm without changing the preference, so retry decomp here. - self._active = self._pref - model = self.program.disasm(entry.ea, entry.name) + focus = self._pending_focus_name # one-shot: consume it here + self._pending_focus_name = None + 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" + dec = self.query_one(DecompView) + if dec.loaded_ea == entry.ea: + # already decompiled: reposition without a recompile + dec.goto(entry.dec_cursor, entry.dec_cursor_x, + entry.dec_scroll_y if entry.dec_scroll_y >= 0 else -1, + entry.dec_scroll_x) + self._show_active() # loads the pseudocode if loaded_ea != entry.ea + return + # Unified model: the code view is always the continuous listing, + # positioned at this entry. The decompiler is a per-function toggle + # (F5/Tab -> _decomp_from_listing), never opened implicitly. + lst = self.query_one(ListingView) + lm = self.program.listing(entry.ea) sy = entry.scroll_y if entry.scroll_y >= 0 else None - self.query_one(DisasmView).load( - model, entry.name, cursor=entry.cursor, cursor_x=entry.cursor_x, scroll_y=sy) - dec = self.query_one(DecompView) - # If the decompiler already shows this function, _show_active won't reload - # it (and thus won't reposition), so move its cursor to the target line - # here — e.g. an xref/goto whose target is inside the current function. - reposition_dec = dec.loaded_ea == entry.ea + if lm is not None: + if sy is None: + # Fresh jump (not a back/forward restore, which carries its own + # scroll). If the target is already on-screen in the SAME listing, + # leave the viewport alone and just move the cursor; otherwise + # scroll so the target sits a few lines below the top for context. + top = round(lst.scroll_offset.y) + if lst.model is lm and top <= entry.cursor < top + lst._visible_height(): + sy = top + else: + sy = max(entry.cursor - _JUMP_CONTEXT, 0) + lst.load( + lm, entry.name, cursor=entry.cursor, + cursor_x=entry.cursor_x, scroll_y=sy, focus=focus) + self._active = "listing" self._show_active() - if reposition_dec and self._active == "decomp": - dsy = entry.dec_scroll_y if entry.dec_scroll_y >= 0 else -1 - dec.goto(entry.dec_cursor, entry.dec_cursor_x, dsy, entry.dec_scroll_x) + # 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") + + def _prompt_active(self) -> bool: + for iid in self._PROMPT_IDS: + try: + if self.query_one(f"#{iid}", Input).display: + return True + except Exception: # noqa: BLE001 — widget not mounted yet + pass + return False def _show_active(self) -> None: - dis = self.query_one(DisasmView) dec = self.query_one(DecompView) + lst = self.query_one(ListingView) hx = self.query_one(HexView) - dis.display = dec.display = hx.display = False - if self._active == "disasm": - dis.display = True - dis.focus() - self._status_for_cur("disasm") + 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"): + # Side-by-side: listing (left) + pseudocode (right), one focused. + 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 + if busy: + dec.loading = True + self._load_decomp(self._cur.ea, self._cur.name) + else: + dec.loading = False + if grab: + (dec if self._active == "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 + # a decompile ran and Esc looked like a no-op. + self._status(f"{self._cur.name} \u2014 decompiling\u2026") + else: + self._status_for_cur("split") + return + self._split = False # a single-view target (hex, etc.) leaves split + self.query_one("#panes").set_class(False, "split") + lst.set_link(set()) + dec.set_link(None) + self._split_eamap = [] + self._split_ea2line = {} + self._split_range = None + dec.display = lst.display = hx.display = gv.display = False + if self._active == "graph": + gv.display = True + if grab: + gv.focus() + self._graph_status() + return + if self._active in ("listing", "disasm"): + lst.display = True + if grab: + lst.focus() + self._status_for_cur("listing") elif self._active == "hex": hx.display = True - hx.focus() + if grab: + hx.focus() ea = getattr(self, "_hex_pending_ea", None) self._hex_pending_ea = None if hx.model is None: @@ -2957,12 +7977,17 @@ class IdaTui(App): self._hex_status(hx.cursor_va()) else: dec.display = True - dec.focus() + if grab: + dec.focus() if self._cur is not None and dec.loaded_ea != self._cur.ea: self._status(f"{self._cur.name} — decompiling…") dec.loading = True # gray out + 'decompiling…' overlay self._load_decomp(self._cur.ea, self._cur.name) else: + # Already showing this function: clear any overlay raised by the + # F5-from-listing path (no re-decompile happens here, so nothing + # else would). + dec.loading = False self._status_for_cur("pseudocode") # -- hex view ---------------------------------------------------------- # @@ -2992,11 +8017,11 @@ class IdaTui(App): self._hex_status(msg.va) def on_hex_view_leave(self, msg: HexView.Leave) -> None: - self._active = self._pref + self._active = self._code_mode() self._show_active() def on_hex_view_to_code(self, msg: HexView.ToCode) -> None: - self._active = self._pref + self._active = self._code_mode() self._goto_ea(msg.va, push=True) def _status_for_cur(self, mode: str) -> None: @@ -3007,21 +8032,49 @@ 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: - # No pseudocode for this function: fall back to the disassembly view - # rather than showing an error panel. Keep the pseudocode preference - # so the next (decompilable) function still opens as pseudocode. + 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. 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" + # 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) + return self._active = "disasm" + 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": view.focus() # loading cover had blurred it; restore focus @@ -3034,26 +8087,222 @@ 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) - self._status(f"{name} @ {ea:#x} [pseudocode {len(dec.code or '')} chars]{note}") + 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 + self._split_status() + else: + self._status( + f"{name} @ {ea:#x} [pseudocode {len(dec.code or '')} chars]{note}") + + def _sync_split(self, source: str) -> 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.""" + if not self._split or self.program is None: + return + lst = self.query_one(ListingView) + dec = self.query_one(DecompView) + if source == "decomp": + dec.set_link(None) # the driver shows its own cursor, no band + line, screen = self._split_anchor(dec) + eas = (self._split_eamap[line] + if 0 <= line < len(self._split_eamap) else []) + if not eas: # fallback: the single /*ea*/ marker for the line + one = dec._line_ea(line) + eas = [one] if one is not None else [] + rows: set[int] = set() + if lst.model is not None: + for e in eas: + r = lst.model.ensure_ea(e) + if r is not None and r >= 0: + rows.add(r) + lst.set_link(rows) + if rows: + # keep the linked region level with the driver's anchor row so + # the eye tracks straight across the two panes + lst.align(min(rows), screen) + else: # the listing drives + lst.set_link(set()) + row, screen = self._split_anchor(lst) + head = lst._head(row) + ea = head.ea if head is not None else None + if ea is None: + dec.set_link(None) + return + rng = self._split_range + if 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) + return + line = self._split_ea2line.get(ea) + if line is None: # fallback: nearest marker + line = dec.line_for_ea(ea) + if line is not None: + dec.set_link(line) + dec.align(line, screen) + else: + dec.set_link(None) + + @staticmethod + def _split_anchor(view) -> tuple[int, int]: # type: ignore[no-untyped-def] + """(row, screen_row) the split sync anchors on: the driver's cursor while + it's visible, else the top visible row. A wheel-scroll/scrollbar drag + never moves the cursor, so anchoring on the viewport keeps the companion + following once the cursor scrolls out of sight.""" + top = round(view.scroll_offset.y) + if top <= view.cursor < top + view._visible_height(): + return view.cursor, view.cursor - top + return top, 0 + + def on_listing_view_scrolled(self, msg: "ListingView.Scrolled") -> None: + if self._split and self._active == "listing": + self._sync_split("listing") + + def on_decomp_view_scrolled(self, msg: "DecompView.Scrolled") -> None: + if self._split and self._active == "decomp": + self._sync_split("decomp") + + @work(thread=True, group="split-resync", exclusive=True) + def _resync_decomp(self, ea: int) -> None: + # The listing cursor crossed out of the decompiled function; find the + # function it's in now (off the UI thread) and re-point the decomp pane. + assert self.program is not None + fn = self.program.function_of(ea) + self.app.call_from_thread(self._apply_resync, fn) + + def _apply_resync(self, fn) -> None: # type: ignore[no-untyped-def] + if not self._split or self._cur is None: + return + dec = self.query_one(DecompView) + if fn is None: + dec.set_link(None) # over data/undefined: keep the decomp, drop the band + 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") + return + dec.loading = True + self._load_decomp(fn.addr, fn.name) # -> _apply_decomp -> map -> re-sync + + def _split_status(self) -> None: + """A split-aware status line reflecting the focused pane + the link.""" + if self._cur is None: + return + if self._active == "decomp": + dec = self.query_one(DecompView) + ea = dec._line_ea(dec.cursor) + n = (len(self._split_eamap[dec.cursor]) + if 0 <= dec.cursor < len(self._split_eamap) else 0) + at = f" @ {ea:#x}" if ea is not None else "" + rel = f" \u2194 {n} insn" if n else "" + self._status(f"{self._cur.name}{at} " + f"[split \u00b7 pseudocode line {dec.cursor + 1}{rel}]" + f" (Tab/click: drive listing)") + else: + ea = self.query_one(ListingView)._cursor_ea() + at = f" @ {ea:#x}" if ea is not None else "" + self._status(f"{self._cur.name}{at} [split \u00b7 listing]" + f" (Tab/click: drive pseudocode)") + + def on_descendant_focus(self, event) -> None: # type: ignore[no-untyped-def] + """In split, focusing a pane (Tab or a mouse click) makes it the leading/ + driver pane, so the sync direction follows where you're actually working.""" + if not self._split: + return + w = event.control + new = ("decomp" if isinstance(w, DecompView) + else "listing" if isinstance(w, ListingView) else None) + if new is not None and new != self._active: + self._active = new + self._sync_split(new) + self._split_status() + + @work(thread=True, group="split-map") + def _load_split_map(self, ea: int) -> None: + # Fetch the rich per-line instruction map off the UI thread; sync stays + # on the crude single-ea fallback until it lands. + assert self.program is not None + m = self.program.decomp_map(ea) + self.app.call_from_thread(self._apply_split_map, ea, m) + + def _apply_split_map(self, ea: int, m: list) -> None: + """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 = [] + for line, eas in enumerate(m): + for e in eas: + self._split_ea2line.setdefault(e, line) + alleas.append(e) + # 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 + # 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._trail_map, self._trail_map_ea = m, ea + self._trail_line_of = dict(self._split_ea2line) + self._trail_eas = sorted(self._trail_line_of) + self._trail_span = 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: - if self._nav: + dv = self._try_view(DecompView) + if self._nav and dv is not None: self._nav[-1].dec_cursor = msg.index - self._nav[-1].dec_cursor_x = self.query_one(DecompView).cursor_x + self._nav[-1].dec_cursor_x = dv.cursor_x + if self._split: + if self._active == "decomp": + self._sync_split("decomp") + self._split_status() + return if self._cur is not None: loc = f" @ {msg.ea:#x}" if msg.ea is not None else "" self._status(f"{self._cur.name}{loc} [pseudocode line {msg.index}]") - def on_disasm_view_cursor_moved(self, msg: DisasmView.CursorMoved) -> None: - if self._nav: - # Keep the top-of-history position current (line + column) so that - # returning here later lands exactly where we left. + def _try_view(self, cls): # type: ignore[no-untyped-def] + """The code view, or None. App.query_one searches the TOP screen, so a + cursor-moved message that lands while any modal is up (the loading + overlay, a project switch) would otherwise raise NoMatches and kill the + app from a message handler.""" + try: + return self.query_one(cls) + except Exception: # noqa: BLE001 -- NoMatches: a modal owns the screen + return None + + def on_listing_view_cursor_moved(self, msg: ListingView.CursorMoved) -> None: + lst = self._try_view(ListingView) + if self._nav and lst is not None: self._nav[-1].cursor = msg.index - self._nav[-1].cursor_x = self.query_one(DisasmView).cursor_x + self._nav[-1].cursor_x = lst.cursor_x + if msg.index >= 0: + self._nav[-1].scroll_y = round(lst.scroll_offset.y) + if self._split: + if self._active == "listing": + self._sync_split("listing") + self._split_status() + return ea = msg.ea if ea is not None: - name = self._nav[-1].name if self._nav else "" - self._status(f"{name} @ {ea:#x} (line {msg.index})") + 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: @@ -3061,5 +8310,14 @@ class IdaTui(App): self._ka.stop() if self.program is not None: self.program.close() - if self.client is not None: + if self._pool is not None: + # _save_on_exit is False only when the user chose discard or we + # already saved; an unexpected teardown still writes defensively. + self._pool.close_all(save=self._save_on_exit is not False) + 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) + except Exception: # noqa: BLE001 + pass self.client.close() diff --git a/idatui/client.py b/idatui/client.py deleted file mode 100644 index 3e5edb4..0000000 --- a/idatui/client.py +++ /dev/null @@ -1,678 +0,0 @@ -"""Persistent, thread-safe client for the ida-pro-mcp (idalib) MCP server. - -This is the foundation the whole TUI stands on. Unlike the throwaway CLI in the -ida-mcp skill (which re-does the MCP handshake and spins a fresh socket on every -call), this client: - - * Performs the MCP handshake exactly once and keeps the session warm - (measured: ~7ms/call warm vs ~60ms cold). - * Reuses TCP connections via a small keep-alive pool (no socket churn over a - long TUI session) while still allowing concurrent calls from worker threads. - * Has a precise, *grounded* error taxonomy (see below) so callers can tell - apart transport failures, protocol errors, hard tool errors, and soft - per-item "not found" results. - * Recovers automatically from an expired server session (404 -> re-handshake - -> retry once) and from transient transport hiccups (bounded retries). - * Auto-injects the mandatory ``database=<session_id>`` argument, resolving a - single open session automatically and refusing to guess when several are - open. - -Error taxonomy (verified against the live server, 2026-07): - - IDAConnectionError transport could not be established / was lost - IDATimeoutError a request exceeded its deadline - IDAProtocolError malformed / unexpected HTTP or JSON-RPC framing - IDARPCError JSON-RPC ``error`` object in the envelope - IDAToolError tool returned ``result.isError == true`` (hard failure: - bad params, unknown tool, missing database, ...) - IDASessionError no session / multiple sessions and none pinned - -Note the deliberate *non*-error: a tool that returns ``isError == false`` with an -``error`` field inside its payload (e.g. ``decompile`` on an unknown address, or -per-item failures in a batch tool) is treated as *data*, not an exception. The -caller inspects the payload. Raising on those would break every batch tool. - -The client is stdlib-only (``http.client`` / ``urllib.parse``). -""" - -from __future__ import annotations - -import http.client -import json -import socket -import threading -import time -from collections import deque -from dataclasses import dataclass -from typing import Any -from urllib.parse import urlsplit - -DEFAULT_URL = "http://127.0.0.1:8745/mcp" -PROTOCOL_VERSION = "2025-06-18" - -# Tools that must NOT receive an injected ``database`` argument. -SESSION_AGNOSTIC_TOOLS = frozenset({"idb_list", "idb_open", "int_convert"}) - -# Substrings (lowercased) that identify a stale/invalid IDB session in a tool -# error message. Verified against the live server: after a server restart the -# old session id is gone and every call fails with "Session not found: <id>". -_STALE_SESSION_MARKERS = ("session not found", "database is required") - - -# --------------------------------------------------------------------------- # -# Exceptions -# --------------------------------------------------------------------------- # -class IDAError(Exception): - """Base class for all client errors.""" - - -class IDAConnectionError(IDAError): - """The transport could not be established or was lost.""" - - -class IDATimeoutError(IDAError): - """A request exceeded its deadline.""" - - -class IDAProtocolError(IDAError): - """Malformed or unexpected HTTP / JSON-RPC framing.""" - - -class IDARPCError(IDAError): - """The JSON-RPC envelope carried an ``error`` object.""" - - def __init__(self, code: int, message: str, data: Any = None): - super().__init__(f"JSON-RPC error {code}: {message}") - self.code = code - self.message = message - self.data = data - - -class IDAToolError(IDAError): - """A tool call returned ``result.isError == true`` (a hard failure).""" - - def __init__(self, tool: str, message: str): - super().__init__(f"tool {tool!r} failed: {message}") - self.tool = tool - self.message = message - - -class IDASessionError(IDAError): - """No IDB session is open, or several are and none was pinned.""" - - -# --------------------------------------------------------------------------- # -# Session model -# --------------------------------------------------------------------------- # -@dataclass(frozen=True) -class Session: - session_id: str - filename: str - input_path: str - is_active: bool = False - is_analyzing: bool = False - - @classmethod - def from_dict(cls, d: dict) -> "Session": - return cls( - session_id=d.get("session_id", ""), - filename=d.get("filename", ""), - input_path=d.get("input_path", ""), - is_active=bool(d.get("is_active", False)), - is_analyzing=bool(d.get("is_analyzing", False)), - ) - - -# --------------------------------------------------------------------------- # -# Transport: a tiny keep-alive connection pool over http.client -# --------------------------------------------------------------------------- # -class _Transport: - """Keep-alive HTTP/1.1 pool. Thread-safe. One request == one pooled conn. - - Connections are checked out, used for exactly one request/response, then - returned to the pool if still healthy. On a dropped/broken connection the - request is retried on a fresh connection (bounded by ``max_retries``). - """ - - def __init__(self, url: str, *, max_retries: int = 2, pool_size: int = 8): - parts = urlsplit(url) - if parts.scheme not in ("http", "https"): - raise IDAConnectionError(f"unsupported scheme: {parts.scheme!r}") - self._https = parts.scheme == "https" - self._host = parts.hostname or "127.0.0.1" - self._port = parts.port or (443 if self._https else 80) - self._path = parts.path or "/" - self._max_retries = max_retries - self._pool_size = pool_size - self._idle: deque[http.client.HTTPConnection] = deque() - self._lock = threading.Lock() - - def _new_conn(self, timeout: float) -> http.client.HTTPConnection: - cls = http.client.HTTPSConnection if self._https else http.client.HTTPConnection - return cls(self._host, self._port, timeout=timeout) - - def _checkout(self, timeout: float) -> http.client.HTTPConnection: - with self._lock: - while self._idle: - conn = self._idle.popleft() - # Reuse only if the socket still looks alive. - if conn.sock is not None: - conn.timeout = timeout - try: - conn.sock.settimeout(timeout) - except OSError: - try: - conn.close() - except OSError: - pass - continue - return conn - return self._new_conn(timeout) - - def _checkin(self, conn: http.client.HTTPConnection) -> None: - with self._lock: - if len(self._idle) < self._pool_size: - self._idle.append(conn) - return - try: - conn.close() - except OSError: - pass - - def request( - self, body: bytes, headers: dict[str, str], timeout: float - ) -> tuple[int, dict[str, str], bytes]: - """POST ``body``; return (status, response_headers, response_body).""" - last_exc: Exception | None = None - for attempt in range(self._max_retries + 1): - conn = self._checkout(timeout) - try: - conn.request("POST", self._path, body=body, headers=headers) - resp = conn.getresponse() - status = resp.status - resp_headers = {k.lower(): v for k, v in resp.getheaders()} - data = resp.read() # must fully drain before reuse - except socket.timeout as e: - self._discard(conn) - raise IDATimeoutError(f"request timed out after {timeout}s") from e - except ( - http.client.RemoteDisconnected, - http.client.BadStatusLine, - ConnectionError, - OSError, - ) as e: - # A stale pooled connection, or the server closed on us. Drop it - # and retry on a fresh connection. - self._discard(conn) - last_exc = e - continue - else: - if resp.will_close: - self._discard(conn) - else: - self._checkin(conn) - return status, resp_headers, data - raise IDAConnectionError( - f"transport failed after {self._max_retries + 1} attempts: {last_exc}" - ) from last_exc - - def _discard(self, conn: http.client.HTTPConnection) -> None: - try: - conn.close() - except OSError: - pass - - def close(self) -> None: - with self._lock: - while self._idle: - self._discard(self._idle.popleft()) - - -# --------------------------------------------------------------------------- # -# The client -# --------------------------------------------------------------------------- # -class IDAClient: - """A warm, thread-safe handle to one MCP server (and one pinned IDB). - - Typical use:: - - with IDAClient(db="4f2223f9") as ida: - ida.health() - funcs = ida.call("list_funcs", queries=[{"count": 50}]) - code = ida.call("decompile", addr="main") - - Concurrency: safe to call from multiple threads (Textual workers). Requests - run on independent pooled connections; only the id counter and handshake - state are lock-guarded, so calls do not serialize on each other. - """ - - def __init__( - self, - url: str = DEFAULT_URL, - db: str | None = None, - *, - timeout: float = 30.0, - max_retries: int = 2, - pool_size: int = 8, - client_name: str = "idatui", - client_version: str = "0.0.1", - auto_recover_session: bool = True, - ): - self.url = url - self._db = db - self.timeout = timeout - self.auto_recover_session = auto_recover_session - self._transport = _Transport(url, max_retries=max_retries, pool_size=pool_size) - self._client_info = {"name": client_name, "version": client_version} - - self._rid = 0 - self._sid: str | None = None - self._ready = False - self._state_lock = threading.Lock() # guards _rid, _sid, _ready - self._handshake_lock = threading.Lock() # serializes (re)handshake - - # -- lifecycle --------------------------------------------------------- # - def __enter__(self) -> "IDAClient": - self.connect() - return self - - def __exit__(self, *exc) -> None: - self.close() - - def connect(self) -> "IDAClient": - """Ensure the MCP handshake has completed (idempotent, thread-safe).""" - if self._ready: - return self - self._handshake() - return self - - def close(self) -> None: - self._transport.close() - - # -- low-level plumbing ------------------------------------------------ # - def _next_id(self) -> int: - with self._state_lock: - self._rid += 1 - return self._rid - - def _headers(self) -> dict[str, str]: - h = { - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - } - sid = self._sid - if sid: - h["Mcp-Session-Id"] = sid - return h - - def _handshake(self) -> None: - with self._handshake_lock: - if self._ready: - return - with self._state_lock: - self._sid = None - rid = self._next_id() - status, headers, body = self._transport.request( - self._encode( - { - "jsonrpc": "2.0", - "id": rid, - "method": "initialize", - "params": { - "protocolVersion": PROTOCOL_VERSION, - "capabilities": {}, - "clientInfo": self._client_info, - }, - } - ), - { - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - }, - self.timeout, - ) - if status // 100 != 2: - raise IDAProtocolError( - f"initialize failed: HTTP {status}: {body[:200]!r}" - ) - sid = headers.get("mcp-session-id") - envelope = self._parse_body(body, rid) - self._raise_on_rpc_error(envelope) - with self._state_lock: - self._sid = sid - # notifications/initialized is a fire-and-forget notification. - self._transport.request( - self._encode({"jsonrpc": "2.0", "method": "notifications/initialized"}), - self._headers(), - self.timeout, - ) - with self._state_lock: - self._ready = True - - @staticmethod - def _encode(obj: dict) -> bytes: - return json.dumps(obj).encode() - - @staticmethod - def _parse_body(body: bytes, want_id: int) -> dict: - """Extract the JSON-RPC envelope for ``want_id`` from a (possibly SSE) body.""" - text = body.decode("utf-8", "replace") - found: dict | None = None - for line in text.splitlines(): - line = line.strip() - if line.startswith("data:"): - line = line[5:].strip() - if not line: - continue - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - if isinstance(obj, dict) and obj.get("id") == want_id: - found = obj - if found is None: - raise IDAProtocolError( - f"no JSON-RPC response with id={want_id} in body: {text[:200]!r}" - ) - return found - - @staticmethod - def _raise_on_rpc_error(envelope: dict) -> None: - err = envelope.get("error") - if err: - raise IDARPCError( - code=err.get("code", -1), - message=err.get("message", "unknown"), - data=err.get("data"), - ) - - def _rpc(self, method: str, params: dict, *, timeout: float | None = None) -> dict: - """Send a JSON-RPC request, recovering once from an expired session.""" - if not self._ready: - self._handshake() - to = self.timeout if timeout is None else timeout - rid = self._next_id() - payload = self._encode( - {"jsonrpc": "2.0", "id": rid, "method": method, "params": params} - ) - status, headers, body = self._transport.request(payload, self._headers(), to) - - if status == 404 and self._sid is not None: - # Server session expired: re-handshake and retry exactly once. - with self._state_lock: - self._ready = False - self._handshake() - rid = self._next_id() - payload = self._encode( - {"jsonrpc": "2.0", "id": rid, "method": method, "params": params} - ) - status, headers, body = self._transport.request( - payload, self._headers(), to - ) - - if status // 100 != 2: - raise IDAProtocolError(f"HTTP {status}: {body[:200]!r}") - - envelope = self._parse_body(body, rid) - self._raise_on_rpc_error(envelope) - return envelope - - # -- tool calls -------------------------------------------------------- # - @staticmethod - def _extract_payload(tool: str, result: dict) -> Any: - """Turn an MCP ``result`` object into a Python payload, or raise. - - Precedence: - 1. ``isError == true`` -> IDAToolError (hard failure) - 2. ``structuredContent`` present -> return it (already parsed) - 3. ``content[0].text`` is JSON -> return parsed JSON - 4. otherwise -> return the raw text string - """ - if result.get("isError"): - text = _first_text(result) or "(no message)" - raise IDAToolError(tool, text) - if "structuredContent" in result: - return result["structuredContent"] - text = _first_text(result) - if text is None: - return result - try: - return json.loads(text) - except (json.JSONDecodeError, TypeError): - return text - - def _prepare_args(self, tool: str, args: dict) -> tuple[dict, bool]: - """Return (arguments, db_was_injected). Never mutates the caller's dict.""" - prepared = dict(args) - if tool in SESSION_AGNOSTIC_TOOLS or "database" in prepared: - return prepared, False - prepared["database"] = self.resolve_db() - return prepared, True - - def call_envelope(self, tool: str, *, timeout: float | None = None, **args) -> dict: - """Call ``tool`` and return the full JSON-RPC envelope (for debugging).""" - prepared, _ = self._prepare_args(tool, args) - return self._rpc( - "tools/call", {"name": tool, "arguments": prepared}, timeout=timeout - ) - - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - """Call ``tool`` and return its payload (parsed JSON when possible). - - Raises IDAToolError on a hard tool failure. Soft/per-item errors (an - ``error`` field with ``isError == false``) are returned as data. - - If ``auto_recover_session`` is set and the db was auto-injected, a stale - "Session not found" error (e.g. after a server restart) triggers exactly - one transparent recovery: drop the stale pin, re-resolve the session, - and retry. A db the caller pinned explicitly is never silently switched. - """ - prepared, injected = self._prepare_args(tool, args) - envelope = self._rpc( - "tools/call", {"name": tool, "arguments": prepared}, timeout=timeout - ) - try: - return self._extract_payload(tool, envelope.get("result", {})) - except IDAToolError as e: - if not (injected and self.auto_recover_session and _is_stale_session(e)): - raise - # The pinned IDB session vanished (server restart). Re-resolve the - # sole session and retry once. resolve_db() raises IDASessionError - # if zero/many sessions exist, so we never guess. - self.set_db(None) - prepared2, _ = self._prepare_args(tool, args) - envelope2 = self._rpc( - "tools/call", {"name": tool, "arguments": prepared2}, timeout=timeout - ) - return self._extract_payload(tool, envelope2.get("result", {})) - - # -- session management ------------------------------------------------ # - def list_sessions(self) -> list[Session]: - result = self._rpc("tools/call", {"name": "idb_list", "arguments": {}}) - payload = self._extract_payload("idb_list", result.get("result", {})) - sessions = payload.get("sessions", []) if isinstance(payload, dict) else [] - return [Session.from_dict(s) for s in sessions] - - def resolve_db(self) -> str: - """Return the pinned session id, auto-resolving a lone open session. - - Raises IDASessionError if none is open, or if several are open and none - has been pinned via ``db=`` / :meth:`set_db`. - """ - if self._db: - return self._db - sessions = self.list_sessions() - usable = [s for s in sessions if s.session_id] - if len(usable) == 1: - self._db = usable[0].session_id - return self._db - if not usable: - raise IDASessionError( - "no open IDB session with a usable id; open one with idb_open" - ) - opts = ", ".join(f"{s.session_id} ({s.filename})" for s in usable) - raise IDASessionError( - f"multiple sessions open; pin one with db=... . Options: {opts}" - ) - - def set_db(self, db: str | None) -> None: - self._db = db - - @property - def db(self) -> str | None: - return self._db - - def _session_input_path(self, db: str) -> str: - for s in self.list_sessions(): - if s.session_id == db: - if not s.input_path: - raise IDASessionError(f"session {db} has no input_path to re-open") - return s.input_path - raise IDASessionError(f"session {db} not found") - - def bump_idle_ttl(self, idle_ttl_sec: int = 1_000_000_000, - path: str | None = None) -> None: - """Raise the worker's idle self-exit TTL so an interactive session never - gets reaped while the user is just reading. - - Headless idalib workers self-exit after ``idle_ttl_sec`` (default 600s) - with no requests. ``idb_open`` on the already-open path is idempotent - (returns the same session) and re-applies the TTL, so this simply opens - the pinned session's path with a huge TTL. The default (~31 years) is - effectively 'never'. - """ - p = path or self._session_input_path(self.resolve_db()) - self.call("idb_open", input_path=p, idle_ttl_sec=int(idle_ttl_sec)) - - def keepalive(self, interval: float = 120.0) -> "KeepAlive": - """Return a (not-yet-started) heartbeat that touches the worker's idle - watchdog every ``interval`` seconds. Belt-and-suspenders next to - :meth:`bump_idle_ttl`; also covers adopted sessions we didn't open. - """ - return KeepAlive(self, interval=interval) - - # -- convenience ------------------------------------------------------- # - def health(self) -> dict: - return self.call("server_health") - - def list_tools(self) -> list[tuple[str, str]]: - envelope = self._rpc("tools/list", {}) - out = [] - for t in envelope.get("result", {}).get("tools", []): - desc = (t.get("description") or "").splitlines() - out.append((t["name"], desc[0] if desc else "")) - return out - - def schema(self, tool: str) -> dict: - envelope = self._rpc("tools/list", {}) - for t in envelope.get("result", {}).get("tools", []): - if t["name"] == tool: - return t.get("inputSchema", {}) - raise IDAError(f"tool not found: {tool}") - - -class KeepAlive: - """Background heartbeat that keeps an idalib worker from idling out. - - Any forwarded request resets the worker's idle timer, so a periodic cheap - ``server_health`` is enough. Failures are swallowed (the next real call will - auto-recover); the heartbeat just keeps a chilling TUI's worker alive. - """ - - def __init__(self, client: "IDAClient", interval: float = 120.0): - if interval <= 0: - raise ValueError("interval must be > 0") - self._client = client - self.interval = interval - self._stop = threading.Event() - self._thread: threading.Thread | None = None - self.beats = 0 - self.failures = 0 - - def start(self) -> "KeepAlive": - if self._thread is not None: - return self - self._stop.clear() - self._thread = threading.Thread( - target=self._run, daemon=True, name="idatui-keepalive" - ) - self._thread.start() - return self - - def stop(self) -> None: - self._stop.set() - t = self._thread - self._thread = None - if t is not None: - t.join(timeout=2.0) - - def __enter__(self) -> "KeepAlive": - return self.start() - - def __exit__(self, *exc) -> None: - self.stop() - - def _run(self) -> None: - while not self._stop.wait(self.interval): - try: - self._client.health() - self.beats += 1 - except IDAError: - self.failures += 1 - - -def _is_stale_session(e: IDAToolError) -> bool: - msg = e.message.lower() - return any(marker in msg for marker in _STALE_SESSION_MARKERS) - - -def _first_text(result: dict) -> str | None: - content = result.get("content") - if isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - return item.get("text") - return None - - -# --------------------------------------------------------------------------- # -# Tiny self-check CLI: python -m idatui.client [--db ID] [--url URL] [health] -# --------------------------------------------------------------------------- # -def _main(argv: list[str]) -> int: - import os - - url = DEFAULT_URL - db = os.environ.get("IDA_MCP_DB") - rest: list[str] = [] - it = iter(argv) - for a in it: - if a == "--url": - url = next(it) - elif a == "--db": - db = next(it) - else: - rest.append(a) - - ida = IDAClient(url, db=db) - try: - ida.connect() - t0 = time.time() - sessions = ida.list_sessions() - print(f"sessions ({(time.time() - t0) * 1e3:.1f}ms):") - for s in sessions: - mark = "*" if s.is_active else " " - print(f" {mark} {s.session_id or '<none>':10} {s.filename}") - try: - h = ida.health() - print("health:", json.dumps(h, indent=1)[:400]) - except IDASessionError as e: - print(f"health: skipped ({e})") - finally: - ida.close() - return 0 - - -if __name__ == "__main__": - import sys - - raise SystemExit(_main(sys.argv[1:])) diff --git a/idatui/domain.py b/idatui/domain.py index 3dd21ea..222f7bb 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -28,14 +28,17 @@ import threading import urllib.request from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace -from typing import Callable +from typing import Callable, TYPE_CHECKING -from .client import IDAClient, IDAToolError +from .errors import IDAToolError + +if TYPE_CHECKING: # type hint only + from .worker_client import WorkerClient # noqa: F401 # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 DISASM_BLOCK = 256 # instructions per cached/fetched block (<= disasm cap) -HEX_BLOCK = 4096 # bytes per cached/fetched hex block +HEX_BLOCK = 16384 # bytes per cached/fetched hex block (compact read_raw -> cheap) DECOMPILE_TIMEOUT = 15.0 # s; cap per decompile so a failing one can't hang the CLI _TRUNC_RE = re.compile(r"\[(\d+) chars total\]\s*$") @@ -58,7 +61,14 @@ class Func: @classmethod def from_raw(cls, d: dict) -> "Func": - return cls(addr=_as_int(d["addr"]), name=d["name"], size=_as_int(d["size"])) + addr = _as_int(d["addr"]) + name = d.get("name") + # An unnamed function (server returns null/empty) 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: + name = f"sub_{addr:X}" + return cls(addr=addr, name=name, size=_as_int(d.get("size", 0))) @dataclass(frozen=True) @@ -79,6 +89,81 @@ 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.""" + + ea: int + kind: str # 'code' | 'data' | 'unknown' | 'member' + size: int + 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 the worker didn't provide them (older worker, or the spans + #: disagreed with the plain text, in which case the text wins). + spans: tuple[tuple[str, str], ...] | 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. + ops: tuple[tuple[int, int, int], ...] | 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": + sp = d.get("spans") + ops = d.get("ops") + 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"), + spans=(tuple((str(k), str(t)) for k, t in sp) + if isinstance(sp, list) and sp else None), + ops=(tuple((int(a), int(b), int(n)) for a, b, n in ops) + if isinstance(ops, list) and ops else 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 @@ -90,9 +175,10 @@ class Ref: class Xref: frm: int # the referencing address to: int | None # the referenced address - type: str # "code" | "data" | ... + type: str # coarse: "code" | "data" fn_name: str | None # function containing `frm` fn_addr: int | None + kind: str | None = None # fine: call/jump/flow/read/write/offset/text/info @dataclass(frozen=True) @@ -129,6 +215,44 @@ class Struct: ) +@dataclass(frozen=True) +class StrLit: + """A string literal IDA found in the binary (the Shift+F12 list).""" + addr: int + text: str + length: int + type: str = "" + + +def link_name(raw: str) -> str: + """A linkage name reduced to what actually joins across binaries. + + ELF symbol versioning means the importer sees ``strrchr@@GLIBC_2.2.5`` while + the provider may export ``strrchr``, ``strrchr@GLIBC_2.2.5`` or the versioned + spelling — comparing raw names silently resolves almost nothing. Cut at the + first '@' so both sides meet on the bare symbol. + """ + n = (raw or "").strip() + at = n.find("@") + return n[:at] if at > 0 else n + + +@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), + ``ordinal`` for exports. + + ``name`` is the joinable name; ``raw`` keeps the spelling IDA reported, which + is what the user sees in the listing. + """ + addr: int + name: str + module: str = "" + ordinal: int = 0 + raw: str = "" + + @dataclass class Decompilation: ea: int @@ -277,8 +401,14 @@ class DisasmModel: self._lock = threading.Lock() self._inflight: set[int] = set() + def _end_kw(self) -> dict: + end = self._function_end() + return {"end": hex(end)} if end is not None else {} + def total(self) -> int: - """Total instruction count (fetched once; ~200ms on huge funcs).""" + """Instruction/row count of the function (fetched once). Uses disasm's + ``include_total`` — one fast call, no response-size truncation. For a + 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( @@ -332,15 +462,22 @@ class DisasmModel: self._max_raw = biggest return out + @staticmethod + def _line_from_head(r: dict) -> Line: + """Adapt a ``heads`` row to a disasm Line (label = the head's name).""" + return Line(ea=_as_int(r["ea"]), text=r.get("text", ""), + label=r.get("name")) + def _fetch_block(self, b: int) -> list[Line]: - # Over-fetch one instruction so the block knows where its last - # instruction ends (variable-length archs give no size field). + # 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( - "disasm", addr=hex(self.ea), offset=b * self.BLOCK, - max_instructions=self.BLOCK + 1, + "heads", addr=hex(self.ea), offset=b * self.BLOCK, + count=self.BLOCK + 1, **self._end_kw(), ) - raw = payload.get("asm", {}).get("lines", []) if isinstance(payload, dict) else [] - fetched = [Line.from_raw(r) for r in raw] + rows = payload.get("heads", []) if isinstance(payload, dict) else [] + fetched = [self._line_from_head(r) for r in rows] lines = fetched[:self.BLOCK] if len(fetched) > self.BLOCK: end_ea: int | None = fetched[self.BLOCK].ea @@ -470,6 +607,308 @@ class DisasmModel: # --------------------------------------------------------------------------- # +# Listing model: lazily-grown flat listing (code + data + undefined) per segment +# --------------------------------------------------------------------------- # +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 + 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) + + def __init__(self, program: "Program", seg_start: int, seg_end: int, + name: str | None = None): + self._prog = program + self.seg_start = seg_start + self.seg_end = seg_end + self.name = name or f"seg @ {seg_start:#x}" + self._heads: list[Head] = [] + self._by_ea: dict[int, int] = {} + # Logical rows != physical heads. A run of undefined bytes arrives as ONE + # head ("db 2044 dup(?)") because materialising millions of one-byte rows + # for a .bss would be absurd — but you must still be able to put the + # cursor on any byte in it and press `c`, exactly as in IDA. So a run of + # N bytes PRESENTS as N rows and the text for each is synthesised on + # 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 + 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 + self._done = False + self._max_raw = 0 # widest opcode length (bytes) seen, for the op column + self._lock = threading.Lock() + # Serializes page loads so a background grower and an in-view search can + # both drive loading without double-fetching the same page. + self._load_lock = threading.Lock() + + # Opcode bytes are only worth showing for code; cap the bulk read so a page + # 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) + 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 + + def max_raw_len(self) -> int: + with self._lock: + return self._max_raw + + def load_next_page(self) -> int: + """Load one more page of heads; returns how many were added.""" + return self._load_next_page() + + def _load_next_page(self) -> int: + with self._load_lock: + return self._load_next_page_locked() + + def _load_next_page_locked(self) -> int: + with self._lock: + if self._done or self._next is None: + return 0 + frm = self._next + payload = self._prog.client.call( + "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) + with self._lock: + for h in page: + # Banner/label rows (function headers, separators, code labels) + # are display-only; don't index them so navigation lands on the + # real code/data head at that address. + if h.kind not in ("sep", "funchdr", "label"): + self._by_ea.setdefault(h.ea, self._rows) + self._row_at.append(self._rows) + self._head_eas.append(h.ea) + self._heads.append(h) + self._rows += self._span(h) + nxt = cur.get("next") + if nxt is None: + self._done = True + self._next = None + else: + self._next = _as_int(nxt) + return len(rows) + + @staticmethod + def _span(h: Head) -> int: + """How many logical rows head ``h`` occupies.""" + return h.size if (h.kind == "unknown" and h.size > 1) else 1 + + def _phys(self, row: int) -> tuple[int, int]: + """(physical head index, byte offset into it) for logical ``row``.""" + import bisect + i = bisect.bisect_right(self._row_at, row) - 1 + if i < 0: + return (-1, 0) + return (i, row - self._row_at[i]) + + def _unknown_bytes(self, ea: int, n: int) -> bytes: + """Bytes behind an undefined run, read in blocks and cached. + + Undefined rows are the ones you carve, so their VALUES are the whole + point — "db ?" with no byte tells you nothing about where an instruction + stream might start. + """ + BLK = 1024 + out = bytearray() + a = ea + while len(out) < n: + b0 = (a // BLK) * BLK + blk = self._ubytes.get(b0) + if blk is None: + try: + blk = self._prog.read_bytes(b0, BLK) + except Exception: # noqa: BLE001 + blk = b"" + self._ubytes[b0] = blk + off = a - b0 + take = min(BLK - off, n - len(out)) + chunk = blk[off:off + take] if blk else b"" + if not chunk: + break + out += chunk + a += len(chunk) + return bytes(out) + + def _row_head(self, i: int, off: int) -> Head: + """The Head for one logical row: the physical head, or a synthesised + single-byte row inside an undefined run. + + The run's FIRST row is synthesised too. Leaving "db 2044 dup(?)" there + would say the row covers 2044 bytes when it now covers one, and the + column of byte values would start an address late. + """ + h = self._heads[i] + if self._span(h) == 1: + return h + ea = h.ea + off + b = self._unknown_bytes(ea, 1) + text = f"db {b[0]:02X}h" if b else "db ?" + return Head(ea=ea, kind="unknown", size=1, text=text, + name=h.name if off == 0 else None) + + def ensure(self, n: int) -> None: + """Ensure at least ``n`` logical rows are loaded (or all, if fewer).""" + while not self._done and self._rows < n: + if self._load_next_page() == 0: + break + + def ensure_ea(self, ea: int) -> int: + """Walk forward until the head containing ``ea`` is loaded; return its + line index (or the nearest head at/after it), or -1 if past the end.""" + while True: + idx = self.index_of_ea(ea) + if idx >= 0: + return idx + with self._lock: + have = self._rows + last_ea = (self._heads[-1].ea + max(self._heads[-1].size, 1) - 1 + if self._heads else -1) + done = self._done + if done or (have and last_ea >= ea): + # Loaded past ea without an exact head hit: return the first head + # at/after ea (a mid-item address lands on its containing head). + return self._first_at_or_after(ea) + if self._load_next_page() == 0: + return self._first_at_or_after(ea) + + def _first_at_or_after(self, ea: int) -> int: + with self._lock: + j = self._head_index_at(ea) + if j >= 0: + h = self._heads[j] + if h.ea <= ea < h.ea + max(h.size, 1): + off = (ea - h.ea) if self._span(h) > 1 else 0 + return self._row_at[j] + off + for i, h in enumerate(self._heads): + if h.ea <= ea < h.ea + max(h.size, 1): + # Inside an undefined run, land on the exact BYTE. + off = (ea - h.ea) if self._span(h) > 1 else 0 + return self._row_at[i] + off + if h.ea > ea: + return self._row_at[i] + return -1 + + def load_all(self, progress: Callable[[int], None] | None = None) -> None: + while not self._done: + if self._load_next_page() == 0: + break + if progress: + progress(self._rows) + + @property + def complete(self) -> bool: + with self._lock: + return self._done + + def loaded(self) -> int: + with self._lock: + return self._rows + + def __len__(self) -> int: + return self.loaded() + + def get(self, i: int) -> Head | None: + 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 + + def window(self, start: int, count: int) -> list[Head]: + """``count`` logical rows from ``start`` (synthesising undefined ones).""" + self.ensure(start + count) + 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 + plain = [(j, off, heads[j]) for j, off in spans if j >= 0] + return [self._row_head(j, off) if self._span(h) > 1 else h + for j, off, h in plain] + + def index_of_ea(self, ea: int) -> int: + with self._lock: + hit = self._by_ea.get(ea) + if hit is not None: + return hit + # An address INSIDE an undefined run is a real row now, not a + # mid-item address: that is what makes `g <addr>` + `c` work + # anywhere in a blob. Heads are address-ordered, so bisect rather + # than scan — a big listing has hundreds of thousands of them and + # this is on the navigation path. + j = self._head_index_at(ea) + if j >= 0: + h = self._heads[j] + if self._span(h) > 1 and h.ea <= ea < h.ea + h.size: + return self._row_at[j] + (ea - h.ea) + return -1 + + 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 + + # -- DisasmModel-compatible accessors (unified model) ------------------ # + def cached_line(self, idx: int) -> Head | None: + """Alias of get() for the disasm view's Line interface.""" + return self.get(idx) + + def lines(self, start: int, count: int, prefetch: bool = True) -> list[Head]: + return self.window(start, count) + + def is_cached(self, start: int, count: int) -> bool: + with self._lock: + return start + count <= self._rows + + def ensure_async(self, start: int, count: int) -> None: + pass # the background grower streams the rest in; nothing to prefetch + + +# --------------------------------------------------------------------------- # # Hex model: block-cached byte view over the loaded image (VA-addressed) # --------------------------------------------------------------------------- # class HexModel: @@ -558,18 +997,31 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: IDAClient, prefetch_workers: int = 2): + def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" ) self._indices: dict[str | None, FunctionIndex] = {} 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 + self._segments_cache: list[tuple[int, int, int, str]] | None = None 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._lock = threading.Lock() # -- prefetch plumbing ------------------------------------------------- # @@ -592,23 +1044,48 @@ class Program: return idx # -- sections / segments ---------------------------------------------- # + def _segments(self) -> list[tuple[int, int, int, str]]: + """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. + """ + 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") + 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"]), + int(d.get("file_off", -1)), d.get("name", "") or "")) + except IDAToolError: + segs = [] + if not segs: # older server without file_regions -> survey_binary (slow) + try: + sb = self.client.call("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, + s.get("name", "") or "")) + except (KeyError, ValueError, TypeError): + continue + except Exception: # noqa: BLE001 -- best-effort; callers handle empty + segs = [] + segs.sort() + with self._lock: + self._segments_cache = segs + return segs + def sections(self) -> list[tuple[int, int, str]]: - """Sorted, non-overlapping [(start, end, name)] segment map (cached). - Fetched once from ``survey_binary`` (~0.7s) on first use.""" + """Sorted, non-overlapping [(start, end, name)] segment map (cached).""" if self._sections is not None: return self._sections - secs: list[tuple[int, int, str]] = [] - try: - sb = self.client.call("survey_binary") - for s in (sb.get("segments", []) if isinstance(sb, dict) else []): - try: - secs.append((_as_int(s["start"]), _as_int(s["end"]), - s.get("name", ""))) - except (KeyError, ValueError, TypeError): - continue - except Exception: # noqa: BLE001 -- best-effort; callers handle None - secs = [] - secs.sort() + secs = [(s, e, nm) for s, e, _fo, nm in self._segments()] with self._lock: self._sections = secs return secs @@ -636,16 +1113,7 @@ class Program: the injected ``file_regions`` server tool.""" if self._fileregions is not None: return self._fileregions - regions: list[tuple[int, int, int]] = [] - try: - r = self.client.call("file_regions") - for d in (r.get("regions", []) if isinstance(r, dict) else []): - if isinstance(d, dict) and "start" in d: - regions.append((_as_int(d["start"]), _as_int(d["end"]), - int(d.get("file_off", -1)))) - except IDAToolError: - regions = [] - regions.sort() + regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] with self._lock: self._fileregions = regions return regions @@ -658,9 +1126,29 @@ class Program: return None def read_bytes(self, ea: int, n: int) -> bytes: - """Raw bytes [ea, ea+n) from IDA (gaps read as zero).""" + """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. + """ if n <= 0: return b"" + if not self._no_read_raw: + try: + r = self.client.call("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) + return out[:n] if len(out) >= n else out + b"\x00" * (n - len(out)) + except IDAToolError as e: + # Tool missing on this server: stop trying it, use get_bytes. + if "read_raw" in str(e) or "Unknown tool" in str(e) or "not found" in str(e): + self._no_read_raw = True + else: + return b"\x00" * n + 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)}]) except IDAToolError: @@ -680,14 +1168,33 @@ class Program: def section_of(self, ea: int) -> str | None: """Name of the segment/section containing ``ea`` (e.g. '.got', '.text', '.data.rel.ro', 'LOAD'), or None if unmapped.""" + b = self.segment_bounds(ea) + return b[2] if b else None + + def segment_bounds(self, ea: int) -> tuple[int, int, str] | None: + """(start, end, name) of the segment containing ``ea``, or None.""" secs = self.sections() if not secs: return None i = bisect.bisect_right([s[0] for s in secs], ea) - 1 if 0 <= i < len(secs) and secs[i][0] <= ea < secs[i][1]: - return secs[i][2] + return secs[i] return None + def listing(self, ea: int) -> ListingModel | None: + """Flat listing (code+data+undefined) for the segment containing ``ea``, + cached per segment. None if ``ea`` is unmapped.""" + seg = self.segment_bounds(ea) + if seg is None: + return None + start, end, name = seg + with self._lock: + m = self._listings.get(start) + if m is None: + m = ListingModel(self, start, end, name) + self._listings[start] = m + return m + # -- structs / local types -------------------------------------------- # def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted @@ -759,6 +1266,27 @@ class Program: return None return row.get("error") or "failed to set the prototype" + 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.""" + try: + r = self.client.call("data_type", addr=hex(ea)) + except IDAToolError: + return None + if not isinstance(r, dict) or r.get("error"): + return None + return r + + 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( + "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 {} + if row.get("ok"): + return None + 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). None on success, else an error string.""" @@ -844,15 +1372,233 @@ class Program: return dec def bump_names(self) -> None: - """Signal that symbol names changed (a rename). Disasm names are live in - the IDB, so clearing the block caches is enough for those; decompilation - is generation-checked and force-recompiled lazily on next access.""" + """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.""" with self._lock: self._name_gen += 1 models = list(self._disasm.values()) + self._listings.clear() # listing head rows cache names -> refetch + self._pc_nums.clear() # a reformat moves every literal on its line for m in models: m.invalidate() + def bump_items(self) -> 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).""" + 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() + for m in models: + m.invalidate() + + # -- item / function structure edits (IDA c/d/u/p) --------------------- # + @staticmethod + def _first_result(payload) -> dict: + """Unwrap the first row of a batch tool response ({result:[...]} or a + bare list); soft per-item ``error`` fields ride along in the dict.""" + data = payload.get("result", payload) if isinstance(payload, dict) else payload + if isinstance(data, list): + return data[0] if data and isinstance(data[0], dict) else {} + return data if isinstance(data, dict) else {} + + def define_code(self, ea: int) -> None: + """Convert the bytes at ``ea`` into a code instruction (IDA's 'c'). + 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)}]) + 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)}])) + 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.call("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.call("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.call("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. + """ + try: + r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit)) + except IDAToolError: + self.define_code(ea) + return {"count": 1, "stopped": "single", "end": hex(ea)} + if not isinstance(r, dict) or r.get("error"): + raise IDAToolError("define_code_run", + f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + return r + + def define_func(self, ea: int) -> dict: + """Create a function starting at ``ea`` (IDA's 'p'). + + Prefers the injected tool, which works out the end when IDA can't; + falls back to the plain one for an older worker. + """ + try: + r = self.client.call("define_func_run", addr=hex(ea)) + except IDAToolError: + 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']}") + 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])) + if res.get("error"): + raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") + + def make_data(self, ea: int, type_decl: str, name: str | None = None) -> None: + """Create a typed data item at ``ea`` (IDA's 'd', but typed). ``type_decl`` + is a C type, e.g. 'int', 'unsigned __int32', 'char[5]', 'my_struct'.""" + item: dict = {"addr": hex(ea), "type": type_decl} + if name: + item["name"] = name + res = self._first_result(self.client.call("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'}") + + 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) + 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.call("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.call("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.call("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: + sec = self.section_of(ea) + except Exception: # noqa: BLE001 + 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).""" @@ -862,6 +1608,180 @@ class Program: 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.""" + if not refresh: + with self._lock: + hit = self._strings + if hit is not None: + return hit + out: list[StrLit] = [] + offset, page = 0, 2000 + while True: + try: + payload = self.client.call( + "list_strings", offset=offset, count=page, min_len=min_len, + refresh=(refresh and offset == 0)) + except IDAToolError: + return [] + rows = payload.get("strings", []) if isinstance(payload, dict) else [] + for r in rows: + if not isinstance(r, dict): + continue + out.append(StrLit( + addr=_as_int(r.get("addr", 0)), + text=r.get("text", ""), + length=int(r.get("len", 0) or 0), + type=r.get("type", "") or "", + )) + total = int(payload.get("total", 0) or 0) if isinstance(payload, dict) else 0 + if len(rows) < page or len(out) >= total: + break + offset += len(rows) + with self._lock: + self._strings = out + return out + + 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.""" + with self._lock: + hit = self._linkage + if hit is not None: + return hit + try: + payload = self.client.call("list_linkage", kind="both") + except IDAToolError: + return ([], []) + if not isinstance(payload, dict): + return ([], []) + imps = [Linkage(addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + module=r.get("module", "") or "", + raw=r.get("name", "") or "") + for r in payload.get("imports", []) if isinstance(r, dict)] + exps = [Linkage(addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + ordinal=int(r.get("ordinal", 0) or 0), + raw=r.get("name", "") or "") + for r in payload.get("exports", []) if isinstance(r, dict)] + out = ([i for i in imps if i.name], [e for e in exps if e.name]) + with self._lock: + self._linkage = out + return out + + 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 + decompiler attributes to that line (may be empty). Cached per function + + name generation; ``[]`` if the tool is unavailable.""" + with self._lock: + hit = self._decomp_maps.get(ea) + gen = self._name_gen + if hit is not None and hit[1] == gen: + return hit[0] + try: + payload = self.client.call("decomp_map", addr=hex(ea)) + except IDAToolError: + return [] + lines = payload.get("lines", []) if isinstance(payload, dict) else [] + out = [[_as_int(e) for e in (ln.get("eas") or [])] + for ln in lines if isinstance(ln, dict)] + with self._lock: + 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.call("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) + hi = max(b.end for b in blocks) + rows = self._heads_between(lo, hi) + for b in blocks: + b.rows = [h for h in rows if b.start <= h.ea < 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 + + 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.call("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).""" @@ -878,11 +1798,14 @@ class Program: return _parse_xrefs(payload) def xrefs_to(self, ea: int, limit: int = 2000) -> list[Xref]: - payload = self.client.call( - "xref_query", - queries=[{"addr": hex(ea), "direction": "to", "include_fn": True, - "dedup": True, "count": limit}], - ) + q = [{"addr": hex(ea), "direction": "to", "include_fn": True, + "dedup": True, "count": limit}] + 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) + except IDAToolError: + payload = self.client.call("xref_query", queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -895,8 +1818,20 @@ class Program: return int(s, 16) if re.fullmatch(r"[0-9a-fA-F]+", s): return int(s, 16) - # Symbol name -> ask the server. lookup_funcs takes an array of strings - # and returns [{"query":..., "fn": {addr,name,size} | null, "error":...}]. + # Symbol name -> resolve to the address the NAME denotes (get_name_ea via + # resolve_names). This handles functions, data AND mid-function labels + # (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]) + 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 + # 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]) except IDAToolError as e: @@ -904,9 +1839,33 @@ class Program: res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None if not fn: - raise KeyError(f"cannot resolve {target!r}") + raise KeyError(f"cannot resolve {target!r}{self._name_suggestion(s)}") return _as_int(fn["addr"]) + def _name_suggestion(self, s: str) -> str: + """Best-effort ' — did you mean …?' hint for a failed name resolve. + + ``lookup_funcs`` matches exact function names only, so a demangled or + partial name (``QuaziLies`` for ``_Z9QuaziLiesPcS_ii``) or a data symbol + (``checkKey``) resolves to nothing. Surface the substring matches from + the function index so the caller can retype the exact name instead of + getting a bare ``cannot resolve``. Never raises — suggestions are a + nicety, not a contract. + """ + try: + idx = self.functions(filter=s) + idx.ensure(6) + cands = idx.window(0, 6) + except Exception: # noqa: BLE001 -- suggestions are strictly optional + return "" + if not cands: + return (" (no function name contains it; it may be a data symbol or " + "not a function — pass an address like 0x1234)") + shown = cands[:5] + names = ", ".join(f"{c.name} @ {c.addr:#x}" for c in shown) + more = " …" if len(cands) > len(shown) else "" + return f" — did you mean: {names}{more}?" + # -- comments ---------------------------------------------------------- # def set_comment(self, ea: int, text: str): """Set (empty text clears) the comment at ``ea``; affects both the disasm @@ -947,6 +1906,7 @@ def _parse_xrefs(payload) -> list[Xref]: type=d.get("type", "?"), fn_name=fn.get("name"), fn_addr=_as_int(fn["addr"]) if fn.get("addr") else None, + kind=d.get("kind"), )) return out diff --git a/idatui/drive.py b/idatui/drive.py index 6c96335..b83a8ba 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -165,6 +165,31 @@ def cmd_names(c, args): or "(no match)" +def cmd_binaries(c, args): + """Project inventory: which binaries, which is active, which have a live + worker, how much of each is indexed.""" + r = c.call("binaries") + out = [] + for b in r["binaries"]: + mark = "*" if b["active"] else ("~" if b["resident"] else " ") + out.append(f" {mark} {b['label']:<28} indexed={b['indexed']:<7} {b['source']}") + if r.get("hops"): + out.append(f" (Esc returns to: {' <- '.join(r['hops'])})") + return "\n".join(out) + "\n * active ~ worker resident" + + +def cmd_switch(c, args): + if not args: + raise SystemExit("usage: switch <binary> [addr|name]") + p = {"binary": args[0]} + if len(args) > 1: + a = args[1] + p["addr"] = a if a.lower().startswith("0x") else c.call("resolve", name=a)["ea"] + r = c.call("switch", **p) + fn = (r.get("function") or {}).get("name") + return f" now on {r.get('binary') or args[0]}" + (f" @ {fn}" if fn else "") + + def _rename_one(c, old, new): c.call("goto", target=old, delay_ms=0) st = c.call("rename", name=new, word=old, delay_ms=0) @@ -194,10 +219,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): @@ -208,6 +238,58 @@ 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" @@ -231,7 +313,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, + "binaries": cmd_binaries, "switch": cmd_switch, } diff --git a/idatui/errors.py b/idatui/errors.py new file mode 100644 index 0000000..29b09ae --- /dev/null +++ b/idatui/errors.py @@ -0,0 +1,76 @@ +"""Transport-agnostic error hierarchy and the 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. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +# --------------------------------------------------------------------------- # +# Exceptions +# --------------------------------------------------------------------------- # +class IDAError(Exception): + """Base class for all client errors.""" + + +class IDAConnectionError(IDAError): + """The transport could not be established or was lost.""" + + +class IDATimeoutError(IDAError): + """A request exceeded its deadline.""" + + +class IDAProtocolError(IDAError): + """Malformed or unexpected HTTP / JSON-RPC framing.""" + + +class IDARPCError(IDAError): + """The JSON-RPC envelope carried an ``error`` object.""" + + def __init__(self, code: int, message: str, data: Any = None): + super().__init__(f"JSON-RPC error {code}: {message}") + self.code = code + self.message = message + self.data = data + + +class IDAToolError(IDAError): + """A tool call returned ``result.isError == true`` (a hard failure).""" + + def __init__(self, tool: str, message: str): + super().__init__(f"tool {tool!r} failed: {message}") + self.tool = tool + self.message = message + + +class IDASessionError(IDAError): + """No IDB session is open, or several are and none was pinned.""" + + +# --------------------------------------------------------------------------- # +# Session model +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Session: + session_id: str + filename: str + input_path: str + is_active: bool = False + is_analyzing: bool = False + + @classmethod + def from_dict(cls, d: dict) -> "Session": + return cls( + session_id=d.get("session_id", ""), + filename=d.get("filename", ""), + input_path=d.get("input_path", ""), + is_active=bool(d.get("is_active", False)), + is_analyzing=bool(d.get("is_analyzing", False)), + ) diff --git a/idatui/formats.py b/idatui/formats.py new file mode 100644 index 0000000..f09bb99 --- /dev/null +++ b/idatui/formats.py @@ -0,0 +1,143 @@ +"""Recognising what a file is, and what to ask when we can't. + +IDA picks a loader by looking at the file. When one matches, it knows the +processor, the load address and often the entry point, and there is nothing to +ask. When none matches — a raw firmware dump, a flash image, a decompressed blob +— IDA falls back to a plain binary load: **x86 at address 0**. That is a silent +wrong answer, not an error; the database opens, analysis runs, and finds nothing. + +Interactive IDA handles this by asking. This module is the "should we ask, and +what are the sensible answers" half of doing the same. + +The sniff deliberately only recognises formats IDA definitely has a loader for. +Being wrong in the cautious direction costs one dismissible dialog; being wrong +the other way is the silent-nothing case we're trying to kill. +""" + +from __future__ import annotations + +import os + +#: (magic, offset, name) for formats IDA loads without help. +_MAGIC: tuple[tuple[bytes, int, str], ...] = ( + (b"\x7fELF", 0, "ELF"), + (b"MZ", 0, "PE/MZ"), + (b"\xfe\xed\xfa\xce", 0, "Mach-O"), + (b"\xfe\xed\xfa\xcf", 0, "Mach-O 64"), + (b"\xce\xfa\xed\xfe", 0, "Mach-O (LE)"), + (b"\xcf\xfa\xed\xfe", 0, "Mach-O 64 (LE)"), + (b"\xca\xfe\xba\xbe", 0, "Mach-O fat/Java class"), + (b"dex\n", 0, "Dalvik"), + (b"\x00asm", 0, "WebAssembly"), + (b"!<arch>", 0, "ar archive"), + (b"\x4c\x01", 0, "COFF i386"), + (b"\x64\x86", 0, "COFF x64"), + (b"\x00\x00\x03\xf3", 0, "Amiga hunk"), + (b"\x7fCGC", 0, "CGC"), +) + +#: Text-ish container formats: recognised by their first line, not a magic word. +_TEXT_PREFIX: tuple[tuple[bytes, str], ...] = ( + (b":", "Intel HEX"), + (b"S0", "Motorola S-record"), + (b"S1", "Motorola S-record"), + (b"S2", "Motorola S-record"), + (b"S3", "Motorola S-record"), +) + + +def sniff(path: str) -> str | None: + """The container format of ``path``, or None when nothing recognises it. + + None is the interesting answer: it means IDA will guess, and its guess is + x86-at-zero regardless of what the bytes actually are. + """ + try: + with open(path, "rb") as f: + head = f.read(64) + except OSError: + return None + if not head: + return None + for magic, off, name in _MAGIC: + if head[off:off + len(magic)] == magic: + return name + # Only treat a text prefix as a format if the whole head is printable — + # a raw blob starting with 0x3a (':') is far more likely than Intel HEX. + if all(32 <= b < 127 or b in (9, 10, 13) for b in head): + for prefix, name in _TEXT_PREFIX: + if head.startswith(prefix): + return name + return None + + +def needs_load_options(path: str) -> bool: + """True when we should ask how to load ``path`` rather than let IDA guess.""" + return os.path.isfile(path) and sniff(path) is None + + +#: Processors worth offering, as (IDA -p name, human label). +#: +#: IDA ships 73 processor modules, most of which are museum pieces; a list that +#: long is a worse answer than a short one plus free text. These are the targets +#: that actually turn up in firmware work, endianness spelled out because +#: getting it wrong is the most common way to end up with zero functions. +#: +#: EVERY NAME HERE IS VERIFIED against a real IDA by opening a scratch blob with +#: ``-p<name>`` and reading back ``inf_get_procname()`` — see +#: ``tools/verify_procs.py``. That is not ceremony: a wrong name is REJECTED by +#: IDA (rc=4) with no useful message, which is the same silent-failure class this +#: dialog exists to prevent. Two of the first twenty were wrong ("h8" and +#: "sparc" are module filenames, not processor names), and the aliases people +#: reach for first — arm64, aarch64, mips, m68k — are all invalid. Re-run the +#: script before adding to this list. +PROCESSORS: tuple[tuple[str, str], ...] = ( + # 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"), + ("mipsb", "MIPS — big-endian"), + ("ppc", "PowerPC — big-endian"), + ("ppcl", "PowerPC — little-endian"), + ("sh4", "SuperH SH-4"), + ("68k", "Motorola 68000 / m68k"), + ("riscv", "RISC-V"), + ("tricore", "Infineon TriCore"), + ("xtensa", "Tensilica Xtensa (ESP32 etc.)"), + ("avr", "Atmel AVR"), + ("z80", "Zilog Z80"), + ("tms320c6", "TI TMS320C6x DSP"), + ("m32r", "Renesas M32R"), + ("arc", "Synopsys ARC"), + ("h8300", "Renesas H8/300"), + ("sparcb", "SPARC — big-endian"), + ("sparcl", "SPARC — little-endian"), + ("s390", "IBM S/390"), +) + + +def load_args(processor: str = "", base: int = 0, extra: str = "") -> str: + """``processor``/``base`` as IDA command-line switches. + + ``-b`` is in PARAGRAPHS, not bytes: ``-b1000`` loads at 0x10000. Every caller + goes through here so that conversion is done in exactly one place. + """ + parts = [] + if processor: + parts.append(f"-p{processor}") + if base: + parts.append(f"-b{int(base) >> 4:x}") + if extra: + parts.append(extra) + return " ".join(parts) diff --git a/idatui/graph.py b/idatui/graph.py new file mode 100644 index 0000000..e0a9422 --- /dev/null +++ b/idatui/graph.py @@ -0,0 +1,715 @@ +"""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 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: + ps = sorted(pos[j] for j in side[i]) + if not ps: + return -1.0 + m = len(ps) // 2 + return float(ps[m]) if len(ps) % 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 = _pair_cross(a, b, down, pos) + _pair_cross(a, b, up, pos) + swap = _pair_cross(b, a, down, pos) + _pair_cross(b, a, 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 9fc692f..5f504bf 100644 --- a/idatui/highlight.py +++ b/idatui/highlight.py @@ -15,19 +15,28 @@ from pygments.lexers import CLexer from pygments.token import Token # Token -> style, checked in priority order (first hierarchical match wins). -# Palette roughly follows a dark IDE theme. +# +# 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 +# blue whether you're reading disassembly or pseudocode. Contrast ratios are +# against the app background #12161c; signal sits at 7:1+ and structure recedes +# to 3-4:1 so punctuation stops competing with the code. +# +# 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="grey54", italic=True)), - (Token.Keyword.Type, Style(color="#4ec9b0")), # __int64, char, ... - (Token.Keyword, Style(color="#c586c0", bold=True)), # if/else/return/goto - (Token.Name.Builtin, Style(color="#4ec9b0")), - (Token.Literal.String, Style(color="#ce9178")), # "..." - (Token.Literal.Number, Style(color="#b5cea8")), # 0x10, 42 - (Token.Operator, Style(color="#d4d4d4")), - (Token.Punctuation, Style(color="grey70")), - (Token.Name, Style(color="#dcdcaa")), # identifiers / calls + (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 ] -_DEFAULT = Style(color="#d4d4d4") +_DEFAULT = Style(color="#c3cad3") # 11.0:1 body _lexer = CLexer(stripnl=False, ensurenl=False) diff --git a/idatui/index.py b/idatui/index.py new file mode 100644 index 0000000..751ad1f --- /dev/null +++ b/idatui/index.py @@ -0,0 +1,206 @@ +"""ProjectIndex — one searchable index over every binary in a project. + +Phase 2 of docs/PROJECTS.md: searching across binaries must work for binaries +whose worker isn't running, so the index lives on disk rather than in the +Programs' caches. + +SQLite FTS5 with the **trigram** tokenizer, which is stdlib (no dependency) and +indexes arbitrary *substrings* rather than just word prefixes — the right shape +for symbol names and string literals. Measured on a 300k-entry corpus: 1.9 ms per +substring query (vs 11.8 ms for a Python scan and 28.9 ms for plain LIKE), 0.2 ms +per incremental insert. + +Sizing, from real binaries: libcrypto.so.3 (5.9 MB, ~10k functions + ~20k +strings) contributes 0.52 MB of text, and the index runs ~5.7x the text it +covers. A 20-binary project therefore lands around 12-23 MB — against the ``.i64`` +files already in the sidecar, where libcrypto's alone is 72 MB. The index is +roughly 1% of what the project already costs on disk. + +Caveat baked into ``search``: a trigram index cannot answer queries shorter than +three characters — it silently returns nothing rather than erroring — so short +queries fall back to LIKE. Without that, typing "e" then "er" would show "no +matches" until the third keystroke. +""" +from __future__ import annotations + +import os +import sqlite3 +from dataclasses import dataclass + +#: Trigram indexes can't match fewer than 3 characters; below this we scan. +MIN_TRIGRAM = 3 + +KIND_FUNC = "func" +KIND_STRING = "string" +#: Cross-binary linkage: what a binary takes from, and offers to, other modules. +KIND_IMPORT = "import" +KIND_EXPORT = "export" + + +@dataclass(frozen=True) +class Hit: + """One index match.""" + + binary: str + kind: str + addr: int + text: str + + +def _fts_phrase(query: str) -> str: + """``query`` as an FTS5 phrase: quoted so operators are literal, with any + embedded quote doubled.""" + return '"' + query.replace('"', '""') + '"' + + +class ProjectIndex: + """Symbol/string index for a whole project, keyed by binary label.""" + + def __init__(self, path: str) -> None: + self.path = os.path.abspath(path) + parent = os.path.dirname(self.path) + if parent: + os.makedirs(parent, exist_ok=True) + # check_same_thread=False: the TUI indexes from a worker thread and + # queries from the UI thread. Writes are serialised by the caller. + self._db = sqlite3.connect(self.path, check_same_thread=False) + self._db.executescript( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS entries USING fts5( + text, + binary UNINDEXED, kind UNINDEXED, addr UNINDEXED, + tokenize='trigram'); + CREATE TABLE IF NOT EXISTS stamps( + binary TEXT PRIMARY KEY, + size INTEGER, mtime INTEGER, n INTEGER); + """ + ) + self._db.commit() + + # -- freshness --------------------------------------------------------- # + def stamp(self, label: str) -> tuple[int, int, int] | None: + """(size, mtime, entry count) recorded when ``label`` was last indexed.""" + row = self._db.execute( + "SELECT size, mtime, n FROM stamps WHERE binary = ?", (label,)).fetchone() + return tuple(row) if row else None # type: ignore[return-value] + + def is_stale(self, label: str, source: str) -> bool: + """True when ``label`` has never been indexed, or its source changed.""" + st = self.stamp(label) + if st is None: + return True + try: + s = os.stat(source) + except OSError: + return False # source gone: keep what we have rather than wipe it + return (st[0], st[1]) != (s.st_size, int(s.st_mtime)) + + # -- population -------------------------------------------------------- # + def reindex(self, label: str, entries, source: str | None = None) -> int: + """Replace ``label``'s entries with ``entries`` — (kind, addr, text) + triples. Per-binary, so re-indexing one never touches the others.""" + rows = [(text, label, kind, int(addr)) + for kind, addr, text in entries if text] + self._db.execute("DELETE FROM entries WHERE binary = ?", (label,)) + self._db.executemany( + "INSERT INTO entries(text, binary, kind, addr) VALUES(?,?,?,?)", rows) + size = mtime = 0 + if source: + try: + s = os.stat(source) + size, mtime = s.st_size, int(s.st_mtime) + except OSError: + pass + self._db.execute( + "INSERT INTO stamps(binary, size, mtime, n) VALUES(?,?,?,?) " + "ON CONFLICT(binary) DO UPDATE SET size=?, mtime=?, n=?", + (label, size, mtime, len(rows), size, mtime, len(rows))) + self._db.commit() + return len(rows) + + def forget(self, label: str) -> None: + """Drop a binary from the index (removed from the project).""" + self._db.execute("DELETE FROM entries WHERE binary = ?", (label,)) + self._db.execute("DELETE FROM stamps WHERE binary = ?", (label,)) + self._db.commit() + + # -- query -------------------------------------------------------------- # + def search(self, query: str, kind: str | None = None, + limit: int = 500) -> list[Hit]: + """Substring search across every indexed binary, newest-agnostic. + + Uses the trigram index at >= 3 characters and falls back to a LIKE scan + below that (the index can't answer shorter queries and would silently + return nothing). + """ + q = (query or "").strip() + if not q: + return [] + sql = ["SELECT binary, kind, addr, text FROM entries WHERE "] + args: list = [] + if len(q) >= MIN_TRIGRAM: + sql.append("text MATCH ?") + args.append(_fts_phrase(q)) + else: + sql.append("text LIKE ?") + args.append(f"%{q}%") + if kind: + sql.append(" AND kind = ?") + args.append(kind) + sql.append(" LIMIT ?") + args.append(int(limit)) + try: + rows = self._db.execute("".join(sql), args).fetchall() + except sqlite3.OperationalError: + return [] # malformed FTS expression: treat as no matches + return [Hit(binary=b, kind=k, addr=int(a), text=t) for b, k, a, t in rows] + + def exact(self, name: str, kind: str, exclude: str | None = None) -> list[Hit]: + """Every entry whose text is EXACTLY ``name``, for the linkage join. + + Deliberately not ``search()``: an import must resolve to the export of + that name, not to everything containing it (``read`` would otherwise + match ``pread``, ``read_line``, ``thread_start``). Exact match also + works below the trigram floor, which matters — plenty of real exports + are one or two characters. + """ + n = (name or "").strip() + if not n: + return [] + sql = "SELECT binary, kind, addr, text FROM entries WHERE text = ? AND kind = ?" + args: list = [n, kind] + if exclude: + sql += " AND binary <> ?" + args.append(exclude) + rows = self._db.execute(sql + " ORDER BY binary, addr", args).fetchall() + return [Hit(binary=b, kind=k, addr=int(a), text=t) for b, k, a, t in rows] + + def providers(self, name: str, exclude: str | None = None) -> list[Hit]: + """Binaries in the project that EXPORT ``name`` (skip ``exclude``, the + binary asking). This is the 'follow an import to its implementation' + half of the join.""" + return self.exact(name, KIND_EXPORT, exclude) + + def importers(self, name: str, exclude: str | None = None) -> list[Hit]: + """Binaries in the project that IMPORT ``name`` — 'who in the project + calls this export'.""" + return self.exact(name, KIND_IMPORT, exclude) + + # -- introspection ------------------------------------------------------ # + def counts(self) -> dict[str, int]: + """Indexed entry count per binary.""" + return {b: n for b, n in + self._db.execute("SELECT binary, n FROM stamps").fetchall()} + + def total(self) -> int: + return int(self._db.execute( + "SELECT count(*) FROM entries").fetchone()[0]) + + def close(self) -> None: + try: + self._db.close() + except Exception: # noqa: BLE001 + pass + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"<ProjectIndex {self.total()} entries {self.path}>" 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 new file mode 100644 index 0000000..1f9c51c --- /dev/null +++ b/idatui/launch.py @@ -0,0 +1,181 @@ +"""One-shot launcher: ``ida-tui foo.elf`` and you're in the 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. + +Usage: + + ida-tui /path/to/binary # open a binary and drive it + +Extras: --ttl, --no-keepalive, --rpc (all forwarded to the TUI). +""" +from __future__ import annotations + +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).""" + from .formats import load_args + return load_args(load.get("processor", ""), int(load.get("base", 0) or 0), + str(load.get("ida_args", "") or "")) + + +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 main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser( + prog="ida-tui", + description="Open a binary in the IDA TUI (private idalib worker).") + p.add_argument("binary", nargs="*", + help="binary to open and analyze (several with --project " + "creates/extends that project)") + p.add_argument("--project", metavar="FILE", + 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)") + p.add_argument("--no-keepalive", action="store_true", + help="do not run the keepalive 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 " + "falls back to x86 at address 0 — which analyses to nothing. These say " + "how to read it, and are recorded per binary in a project.") + g.add_argument("--processor", metavar="NAME", + help="IDA processor: arm, armb (big-endian), mipsb, metapc, …") + 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") + args = p.parse_args(argv) + + load: dict = {} + if args.processor: + load["processor"] = args.processor + if args.base: + try: + base = int(args.base, 0) + except ValueError: + _log(f"--base must be a number (got {args.base!r})") + return 2 + if base % 16: + # -b is in paragraphs, so an unaligned base cannot be expressed and + # would silently load somewhere else. + _log(f"--base must be 16-byte aligned (got {base:#x})") + return 2 + load["base"] = base + if args.ida_args: + load["ida_args"] = args.ida_args + + project = None + binary = None + if args.project: + from .project import Project, ProjectError + ppath = os.path.abspath(os.path.expanduser(args.project)) + try: + if os.path.isfile(ppath): + project = Project.load(ppath) + if args.binary: # extend an existing project, skipping repeats + before = len(project.refs) + for b in args.binary: + project.add(b, load=load or None) + added = len(project.refs) - before + dupes = len(args.binary) - added + if added: + project.save() + _log(f"added {added} binary(ies) to {ppath}") + if dupes: + _log(f"{dupes} already in the project (matched by path) " + f"— left alone") + elif args.binary: + project = Project.create(ppath, args.binary, load=load or None) + _log(f"created project {ppath} with {len(project.refs)} binaries") + else: + _log(f"no such project: {ppath} (pass binaries to create it)") + return 2 + except ProjectError as e: + _log(str(e)) + return 2 + # Everything IDA writes lives in the project's sidecar, so the source + # tree is never touched and never needs to be writable. + try: + project.stage_all(progress=lambda m: _log(m)) + except ProjectError as e: + _log(str(e)) + return 2 + else: + if len(args.binary) != 1: + _log("give exactly 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}") + 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)}") + 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. + try: + from .app import IdaTui + except ImportError as e: + _log(f"the TUI needs textual; run with ~/ida-venv/bin/python ({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), + trace_path=(os.path.abspath(os.path.expanduser(args.trace)) + if args.trace else "")).run() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/idatui/pane.py b/idatui/pane.py index 6ee45ca..2592581 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -1,16 +1,17 @@ -"""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 # -> {"sock": "/run/user/1000/idatui-3f2a.sock", "pane": "%7", "ready": true, ...} - # or attach to an existing session id - python -m idatui.pane spawn --db 80d83396 - # drive it (see docs/RPC.md / the idatui-rpc skill) python -m idatui.rpcclient --sock <sock> pseudocode target=main @@ -18,9 +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, and the ida-pro-mcp supervisor already up -(./spawn.sh). Uses ~/ida-venv/bin/python for the TUI (needs textual) unless ---python / IDATUI_PYTHON says otherwise. + # 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 spawns its own private idalib +worker (no shared supervisor). Uses ~/ida-venv/bin/python for the TUI (needs +textual) unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -28,14 +33,12 @@ import argparse import json import os import secrets -import socket +import signal import subprocess import sys import time from typing import Any -from urllib.parse import urlparse -from .client import DEFAULT_URL from .rpcclient import RpcClient, RpcError REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -66,10 +69,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: @@ -77,121 +117,268 @@ def _tmux(*args: str) -> str: check=True).stdout.strip() -# --------------------------------------------------------------------------- # -# supervisor (ida-pro-mcp server) — auto-start if down -# --------------------------------------------------------------------------- # -def _server_addr(url: str) -> tuple[str, int]: - u = urlparse(url) - return (u.hostname or "127.0.0.1", u.port or 8745) +def _zellij(*args: str) -> str: + return subprocess.run([*_zellij_argv(), *args], capture_output=True, + text=True, check=True).stdout.strip() -def _server_up(host: str, port: int, timeout: float = 0.75) -> bool: - """Is something listening on host:port? (Cheap TCP probe; the readiness poll - that follows catches a half-up server.)""" +def _zellij_panes() -> list[dict[str, Any]]: try: - with socket.create_connection((host, port), timeout=timeout): - return True - except OSError: - return False + out = subprocess.run([*_zellij_argv(), "action", "list-panes", + "--state", "--json"], + capture_output=True, text=True) + rows = json.loads(out.stdout or "[]") + except (OSError, ValueError): + return [] + return rows if isinstance(rows, list) else [] -def _ensure_server(host: str, port: int, timeout: float, - detached: bool = True) -> dict[str, Any]: - """Make sure the supervisor is up; start ./spawn.sh in a tmux pane if not. +def _pane_alive(pane: str, mux: str | None = None) -> bool: + """True if the pane exists *and* its command is still running. - Only auto-starts a *local* server (can't launch a remote one). spawn.sh binds - the port, so starting a duplicate is impossible — the probe guards that. - ``IDATUI_SERVER_CMD`` overrides the launch command (used by the tests). + 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 _server_up(host, port): - return {"server_started": False, "server_up": True} - if host not in ("127.0.0.1", "localhost", "::1"): - return {"server_started": False, "server_up": False, - "error": f"server at {host}:{port} is down and not local; " - "cannot auto-start"} - cmd_str = os.environ.get("IDATUI_SERVER_CMD", "./spawn.sh") - cmd = f"cd {_q(REPO)} && exec {cmd_str}" - split = ["split-window", "-v", "-P", "-F", "#{pane_id}"] + if 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) - pane = _tmux(*split) - deadline = time.time() + timeout - while time.time() < deadline: - if not _pane_alive(pane): - return {"server_started": True, "server_up": False, "server_pane": pane, - "error": "supervisor pane exited during startup (check it)"} - if _server_up(host, port): - return {"server_started": True, "server_up": True, "server_pane": pane} - time.sleep(0.5) - return {"server_started": True, "server_up": False, "server_pane": pane, - "error": "supervisor did not come up in time"} + 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) + + +# --------------------------------------------------------------------------- # +# 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 _worker_pids() -> list[int]: + """PIDs of our private per-pane idalib worker processes (idatui/worker.py), + never our own PID.""" + try: + out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], + capture_output=True, text=True) + except OSError: + 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 + + +def _count_live_panes() -> int: + 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 # --------------------------------------------------------------------------- # # 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 args.db: - print("error: pass --open <binary> or --db <session>", file=sys.stderr) + if not args.open and not getattr(args, "project", None): + print("error: pass --open <binary> or --project <file>", file=sys.stderr) return 2 sock = args.sock or os.path.join(_sockdir(), f"idatui-{secrets.token_hex(3)}.sock") - target = os.path.abspath(os.path.expanduser(args.open)) if args.open else args.db - if args.open and not os.path.exists(target): + project = (os.path.abspath(os.path.expanduser(args.project)) + if getattr(args, "project", None) else None) + target = os.path.abspath(os.path.expanduser(args.open)) if args.open else None + if target is not None and not os.path.exists(target): print(f"error: no such binary: {target}", file=sys.stderr) return 2 + if project is not None and target is None and not os.path.exists(project): + print(f"error: no such project: {project}", file=sys.stderr) + return 2 - # make sure the ida-pro-mcp supervisor is up (auto-start it if not) - srv: dict[str, Any] = {"server_started": False, "server_up": True} - if not args.no_ensure_server: - host, port = _server_addr(args.url or DEFAULT_URL) - srv = _ensure_server(host, port, args.server_timeout) - if srv.get("server_started"): - print(f"supervisor was down — started it ({srv.get('server_pane')})", - file=sys.stderr) - if not srv.get("server_up"): - print(json.dumps({"ready": False, **srv}), file=sys.stderr) - return 3 + # 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: become the TUI so kill-pane kills it cleanly - inner = [args.python, "-m", "idatui.tui", "--rpc", sock] - if args.open: - inner += ["--open", target] + # 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. + 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). + inner = [args.python, "-m", "idatui.launch", "--project", project] + if target is not None: + inner.append(target) + inner += ["--rpc", sock] else: - inner += ["--db", target] - if args.url: - inner += ["--url", args.url] - cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) + inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] + # 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": target, - "kind": "open" if args.open else "db", "started": time.time(), - "server_started": srv.get("server_started", False)} - if srv.get("server_pane"): - row["server_pane"] = srv["server_pane"] + 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 @@ -202,12 +389,20 @@ def _q(s: str) -> str: return shlex.quote(s) -def _wait_ready(sock: str, timeout: float, pane: str) -> dict[str, Any]: - """Poll the socket + ping until the TUI reports ready (or timeout).""" - deadline = time.time() + timeout +def _wait_ready(sock: str, timeout: float, pane: str, + 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. + """ + 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: @@ -217,6 +412,13 @@ def _wait_ready(sock: str, timeout: float, pane: str) -> dict[str, Any]: return last except (OSError, RpcError, ConnectionError): pass + if not warned and (time.time() - start) > stuck_after: + 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) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -237,25 +439,50 @@ 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) except OSError: pass _save_registry([r for r in reg if r not in rows]) - print(json.dumps({"stopped": [r.get("sock") or r.get("pane") for r in rows]})) + # Reap the workers those panes leaked (safe: only fires once no pane is live). + reaped = _reap_orphan_workers() + 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 @@ -264,7 +491,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"]): @@ -272,32 +500,114 @@ 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: _save_registry(alive) + reaped = _reap_orphan_workers() + if reaped: + print(f"reaped {reaped} orphaned idalib worker(s)", file=sys.stderr) print(json.dumps(alive, indent=2)) return 0 +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) + 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 (dir must be writable)") - sp.add_argument("--db", metavar="SESSION", help="attach to an existing session id") + 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("--url", help="MCP server URL (default: idatui's default)") - sp.add_argument("--no-ensure-server", action="store_true", - help="don't auto-start ./spawn.sh if the supervisor is down") - sp.add_argument("--server-timeout", type=float, default=90.0, - help="seconds to wait for an auto-started supervisor") sp.add_argument("--vertical", action="store_true", help="split vertically (stacked)") - sp.add_argument("--size", help="new pane size (tmux -l value, e.g. 60%% or 120)") + sp.add_argument("--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) @@ -305,12 +615,34 @@ 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.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 new file mode 100644 index 0000000..ae37c25 --- /dev/null +++ b/idatui/pool.py @@ -0,0 +1,235 @@ +"""WorkerPool — keeps a live idalib worker per project binary, within a budget. + +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). + +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. +""" +from __future__ import annotations + +import os + +from .project import BinaryRef, Project + +#: Fallback budget if /proc/meminfo can't be read (MB). +_FALLBACK_BUDGET_MB = 2048 + + +def _total_ram_mb() -> int: + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + return int(line.split()[1]) // 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + +def _pss_mb(pid: int | None) -> int: + """Proportional set size of a worker, 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. + """ + if not pid: + return 0 + try: + with open(f"/proc/{pid}/smaps_rollup") as f: + for line in f: + if line.startswith("Pss:"): + return int(line.split()[1]) // 1024 + except (OSError, ValueError, IndexError): + pass + 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) + + +class WorkerPool: + """Live workers for a project's binaries, keyed by label.""" + + def __init__(self, project: Project, *, budget_mb: int | None = None, + ttl: int = 1800, spawn=None, mem_fn=None) -> None: + self.project = project + self._ttl = ttl + self._spawn = spawn or _default_spawn + self._mem = mem_fn or (lambda c: _pss_mb(getattr(c, "pid", None))) + self._clients: dict[str, object] = {} + self._lru: list[str] = [] # least-recently-used first + self._pinned: set[str] = set() + self.active: str | None = None # never evicted + if budget_mb is None: + ram = _total_ram_mb() + budget_mb = (ram * project.memory_pct // 100) if ram else _FALLBACK_BUDGET_MB + self.budget_mb = max(budget_mb, 256) + self.evicted: list[str] = [] # labels evicted, most recent last + + # -- residency --------------------------------------------------------- # + def resident(self) -> list[str]: + """Labels with a live worker, least-recently-used first.""" + return list(self._lru) + + def is_resident(self, label: str) -> bool: + return label in self._clients + + def memory_mb(self) -> int: + return sum(self._mem(c) for c in self._clients.values()) + + def pin(self, label: str, on: bool = True) -> None: + """Keep ``label`` resident regardless of the budget.""" + self._pinned.add(label) if on else self._pinned.discard(label) + + def is_pinned(self, label: str) -> bool: + return label in self._pinned + + # -- acquire ----------------------------------------------------------- # + def get(self, label: str, progress=None): + """A live client for ``label``, spawning it (and making room) if 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. + """ + client = self._clients.get(label) + if client is not None: + self._touch(label) + return client + ref = self.project.by_label(label) + if ref is None: + raise KeyError(f"no such binary in the project: {label}") + + def note(msg: str) -> None: + if progress is not None: + progress(msg) + + 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) + 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._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. + + 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 + no evidence at all, so we allow one — that is the case where the budget + is certainly free. + """ + if label in self._clients: + return False + if self.project.by_label(label) is None: + return False + used = self.memory_mb() + est = max((self._mem(c) for c in self._clients.values()), default=0) + if used + est > self.budget_mb: + 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 + # 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 _touch(self, label: str) -> None: + if label in self._lru: + self._lru.remove(label) + self._lru.append(label) + + def set_active(self, label: str | None) -> None: + self.active = label + if label: + self._touch(label) + + # -- release ----------------------------------------------------------- # + def evict(self, label: str, save: bool = True) -> bool: + """Drop a resident worker, persisting its database first.""" + client = self._clients.pop(label, None) + if client is None: + return False + if label in self._lru: + self._lru.remove(label) + if save: + try: # persist analysis + edits so the next open is a load + client.call("idb_save") + except Exception: # noqa: BLE001 -- evict regardless + pass + try: + client.close() + except Exception: # noqa: BLE001 + pass + self.evicted.append(label) + return True + + def _evictable(self, protect: str | None) -> str | None: + for label in self._lru: # least-recently-used first + if label == protect or label == self.active or label in self._pinned: + continue + return label + return None + + def _enforce_budget(self, protect: str | None = None) -> int: + """Evict LRU workers until the pool fits its budget. Returns how many.""" + n = 0 + while self.memory_mb() > self.budget_mb: + victim = self._evictable(protect) + if victim is None: # everything left is active/pinned/protected + break + self.evict(victim) + n += 1 + return n + + def close_all(self, save: bool = True) -> None: + for label in list(self._clients): + self.evict(label, save=save) + self.active = None + + # -- introspection ------------------------------------------------------ # + def status(self) -> list[dict]: + """Per-binary residency for the switcher UI.""" + out = [] + for ref in self.project.refs: + client = self._clients.get(ref.label) + out.append({ + "label": ref.label, + "source": ref.source, + "resident": client is not None, + "pinned": ref.label in self._pinned, + "active": ref.label == self.active, + "analysed": self.project.has_db(ref), + "memory_mb": self._mem(client) if client is not None else 0, + }) + return out + + def __repr__(self) -> str: # pragma: no cover - debug aid + return (f"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") diff --git a/idatui/project.py b/idatui/project.py new file mode 100644 index 0000000..e2fc542 --- /dev/null +++ b/idatui/project.py @@ -0,0 +1,353 @@ +"""Multi-binary projects: group the binaries of one target under a project file. + +A project is an explicit JSON file plus a sidecar directory beside it:: + + router-fw.json # the project file + router-fw.idatui.d/ + bin/httpd # hardlink (or copy) of the source binary + bin/httpd.i64 # IDA's DB + scratch land here automatically + idx/httpd.json # cached index (phase 2) + +Binaries are **staged** into ``bin/`` and IDA opens the staged file, so the +``.i64`` and every ``.id0/.id1/.id2/.nam/.til`` scratch file are created inside +the sidecar instead of next to the original. Source trees stay pristine. + +Staging **copies** rather than hardlinks. A hardlink would be free, but source +and staged would be one inode: an in-place rebuild (``cp newbuild /path/bin`` +truncates instead of replacing) would silently swap the bytes under an already +analysed database, with nothing to detect it. A copy costs a fraction of the +``.i64`` it will grow, decouples the two completely, and leaves the sidecar +self-contained — the project still opens after the sources go away (an unmounted +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. +""" +from __future__ import annotations + +import json +import os +import shutil +from dataclasses import dataclass + +SIDECAR_SUFFIX = ".idatui.d" +DEFAULT_MEMORY_PCT = 25 + +#: The database plus the working files IDA unpacks beside it while it's open. +DB_SUFFIXES = (".i64", ".idb", ".id0", ".id1", ".id2", ".nam", ".til") +#: Just the unpacked scratch — safe to delete when no worker holds the DB. +SCRATCH_SUFFIXES = (".id0", ".id1", ".id2", ".nam", ".til") + + +class ProjectError(Exception): + """A malformed project file, or a binary that can't be staged.""" + + +@dataclass(frozen=True) +class BinaryRef: + """One binary in a project: where it came from, and where IDA works on it.""" + + label: str # unique within the project; names the staged file + source: str # absolute path to the original binary + staged: str # absolute path IDA actually opens (inside the sidecar) + #: How to LOAD it. Only meaningful for a headerless blob: an ELF/PE says what + #: 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 + + @property + def db(self) -> str: + """The database IDA creates for the staged file.""" + return self.staged + ".i64" + + @property + def load_args(self) -> str: + """``processor``/``base`` as IDA command-line switches. + + ``-b`` is in PARAGRAPHS, not bytes — ``-b1000`` loads at 0x10000. That + trap is worth hiding: projects say ``"base": "0x8000000"`` and the one + conversion lives in ``formats.load_args``. + """ + from .formats import load_args + return load_args(self.processor, self.base, self.ida_args) + + +def _as_addr(v) -> int: + """A load address from JSON: int, or a string in any base ("0x8000000"). + + Addresses are written by hand in a project file, so accept how people write + them rather than demanding decimal. + """ + if v is None or v == "": + return 0 + if isinstance(v, int): + return v + try: + return int(str(v), 0) + except ValueError: + return 0 + + +def _stat_key(path: str) -> tuple[int, int] | None: + """(size, mtime) identity used to spot a source that changed under us.""" + try: + st = os.stat(path) + except OSError: + return None + return (st.st_size, int(st.st_mtime)) + + +def _unlink(path: str) -> bool: + try: + os.remove(path) + return True + except OSError: + return False + + +class Project: + """A set of binaries analysed together, with all IDA artifacts corralled.""" + + def __init__(self, path: str, name: str, entries: list[dict], + memory_pct: int = DEFAULT_MEMORY_PCT) -> None: + self.path = os.path.abspath(os.path.expanduser(path)) + self.name = name + self.memory_pct = memory_pct + self._entries = entries # raw, as written to the file + self._refs = self._build_refs() + + # -- construction ------------------------------------------------------ # + @classmethod + def load(cls, path: str) -> "Project": + path = os.path.abspath(os.path.expanduser(path)) + try: + with open(path) as f: + raw = json.load(f) + except OSError as e: + raise ProjectError(f"cannot read project {path}: {e}") from e + except ValueError as e: + raise ProjectError(f"malformed project {path}: {e}") from e + if not isinstance(raw, dict): + raise ProjectError(f"malformed project {path}: expected an object") + entries = raw.get("binaries") + if not isinstance(entries, list) or not entries: + raise ProjectError(f"project {path} lists no binaries") + norm: list[dict] = [] + for e in entries: + if isinstance(e, str): + e = {"path": e} + if not isinstance(e, dict) or not e.get("path"): + raise ProjectError(f"project {path}: bad binary entry {e!r}") + # Keep every recognised key: a whitelist of path/label silently + # dropped the load options on the first save, so a blob's processor + # and base vanished the moment the project was reopened. + norm.append({k: e[k] for k in + ("path", "label", "processor", "base", "ida_args") + if e.get(k) not in (None, "")}) + name = raw.get("name") or os.path.splitext(os.path.basename(path))[0] + try: + pct = int(raw.get("memory_pct", DEFAULT_MEMORY_PCT)) + except (TypeError, ValueError): + pct = DEFAULT_MEMORY_PCT + return cls(path, str(name), norm, max(1, min(pct, 90))) + + @classmethod + def create(cls, path: str, binaries: list[str], name: str | None = None, + memory_pct: int = DEFAULT_MEMORY_PCT, load: dict | None = None) -> "Project": + """Write a new project file listing ``binaries`` (an ad-hoc project). + + ``load`` carries per-binary load options (processor/base/ida_args) that + apply to every binary given here — a headerless blob needs them, and one + command line normally adds blobs of the same kind. + """ + if not binaries: + raise ProjectError("a project needs at least one binary") + entries, seen = [], set() + for b in binaries: # the same file twice on one command line is a typo + p = os.path.abspath(os.path.expanduser(b)) + key = os.path.realpath(p) + if key in seen: + continue + seen.add(key) + e = {"path": p} + e.update({k: v for k, v in (load or {}).items() if v}) + entries.append(e) + path = os.path.abspath(os.path.expanduser(path)) + proj = cls(path, name or os.path.splitext(os.path.basename(path))[0], + entries, memory_pct) + proj.save() + return proj + + def save(self) -> None: + data = {"name": self.name, "memory_pct": self.memory_pct, + "binaries": self._entries} + tmp = self.path + ".tmp" + os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) + with open(tmp, "w") as f: + json.dump(data, f, indent=2) + f.write("\n") + os.replace(tmp, self.path) + + # -- layout ------------------------------------------------------------ # + @property + def sidecar(self) -> str: + """The directory holding staged binaries, databases and indexes.""" + return os.path.splitext(self.path)[0] + SIDECAR_SUFFIX + + @property + def bin_dir(self) -> str: + return os.path.join(self.sidecar, "bin") + + @property + def index_dir(self) -> str: + return os.path.join(self.sidecar, "idx") + + def index_path(self, ref: BinaryRef) -> str: + """Where the cached per-binary index lives (phase 2).""" + return os.path.join(self.index_dir, ref.label + ".json") + + # -- binaries ---------------------------------------------------------- # + def _build_refs(self) -> tuple[BinaryRef, ...]: + base = os.path.dirname(self.path) + refs: list[BinaryRef] = [] + used: set[str] = set() + for e in self._entries: + src = os.path.expanduser(str(e["path"])) + if not os.path.isabs(src): # relative to the project file + src = os.path.join(base, src) + src = os.path.abspath(src) + label = str(e.get("label") or os.path.basename(src)) or "binary" + label = label.replace("/", "_").replace(os.sep, "_") + if label in used: # labels name files; keep them unique + stable + n = 2 + while f"{label}_{n}" in used: + n += 1 + label = f"{label}_{n}" + used.add(label) + refs.append(BinaryRef( + label=label, source=src, + staged=os.path.join(self.bin_dir, label), + processor=str(e.get("processor") or ""), + base=_as_addr(e.get("base")), + ida_args=str(e.get("ida_args") or ""))) + return tuple(refs) + + @property + def refs(self) -> tuple[BinaryRef, ...]: + return self._refs + + def set_load(self, label: str, processor: str = "", base: int = 0, + ida_args: str = "") -> BinaryRef | None: + """Record how ``label`` should be loaded, and persist it. + + Answered once: the dialog that asks writes the answer here, so reopening + the project doesn't ask again — and neither does adding the same blob to + another project, since it travels with the entry. + """ + ref = self.by_label(label) + if ref is None: + return None + i = self._refs.index(ref) + e = self._entries[i] + if processor: + e["processor"] = processor + if base: + e["base"] = int(base) + if ida_args: + e["ida_args"] = ida_args + self._refs = self._build_refs() + self.save() + return self._refs[i] + + def by_label(self, label: str) -> BinaryRef | None: + return next((r for r in self._refs if r.label == label), None) + + def by_source(self, binary: str) -> BinaryRef | None: + """The entry for ``binary``, matched by resolved path. + + Identity is the real path, not the file name: a project can legitimately + hold two different ``foo.elf`` from different directories (the labels + disambiguate them), but the same file must not be listed twice — and + ``./a.elf``, ``/abs/a.elf`` and a symlink to it are all the same file. + """ + key = os.path.realpath(os.path.abspath(os.path.expanduser(binary))) + return next((r for r in self._refs + if os.path.realpath(r.source) == key), None) + + def add(self, binary: str, label: str | None = None, + load: dict | None = None) -> BinaryRef: + """Add a binary, or return the existing entry if it's already here.""" + existing = self.by_source(binary) + if existing is not None: + return existing + entry = {"path": os.path.abspath(os.path.expanduser(binary))} + if label: + entry["label"] = label + entry.update({k: v for k, v in (load or {}).items() if v}) + self._entries.append(entry) + self._refs = self._build_refs() + return self._refs[-1] + + def remove(self, label: str) -> bool: + ref = self.by_label(label) + if ref is None: + return False + i = self._refs.index(ref) + del self._entries[i] + self._refs = self._build_refs() + return True + + # -- staging ----------------------------------------------------------- # + def is_stale(self, ref: BinaryRef) -> bool: + """True when the staged file is missing or no longer matches its source.""" + staged = _stat_key(ref.staged) + return staged is None or staged != _stat_key(ref.source) + + def stage(self, ref: BinaryRef) -> str: + """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). + """ + if not os.path.isfile(ref.source): + raise ProjectError(f"no such binary: {ref.source}") + if not self.is_stale(ref): + return ref.staged + os.makedirs(self.bin_dir, exist_ok=True) + tmp = ref.staged + ".staging" + _unlink(tmp) + # copy2 (not link): keeps size+mtime so freshness compares cleanly, while + # leaving the staged bytes immune to an in-place rewrite of the source. + shutil.copy2(ref.source, tmp) + os.replace(tmp, ref.staged) + for suf in DB_SUFFIXES: # the old DB describes the old bytes + _unlink(ref.staged + suf) + return ref.staged + + def stage_all(self, progress=None) -> list[BinaryRef]: + out = [] + for ref in self._refs: + if progress is not None: + progress(f"staging {ref.label}\u2026") + self.stage(ref) + out.append(ref) + return out + + def sweep_scratch(self, ref: BinaryRef) -> int: + """Delete IDA's unpacked working files (never the ``.i64``) for ``ref``. + + A hard-killed worker leaves them behind and the database then refuses to + reopen. Only safe when no worker holds it. + """ + return sum(1 for suf in SCRATCH_SUFFIXES if _unlink(ref.staged + suf)) + + def has_db(self, ref: BinaryRef) -> bool: + """True once the binary has been analysed and saved at least once.""" + return os.path.isfile(ref.db) + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"<Project {self.name!r} {len(self._refs)} binaries {self.path}>" diff --git a/idatui/rpc.py b/idatui/rpc.py index f41c734..43a34b7 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -28,7 +28,7 @@ from typing import Any from rich.console import Console from ._sync import drain, settle -from .app import DecompView, DisasmView, HexView +from .app import DecompView, GraphView, HexView, ListingView PROTO_VERSION = 1 TYPE_DELAY_MS = 35 # default per-char delay for high-level typed ops (aesthetic) @@ -38,6 +38,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", } # Self-documenting method table (returned by the 'methods' verb). @@ -59,22 +60,57 @@ 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", "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,13 +130,79 @@ def _active_widget(app): """The currently *shown* code widget (mirrors app._active).""" if app._active == "hex": return app.query_one(HexView) - if app._active == "disasm": - return app.query_one(DisasmView) + if app._active == "graph": + return app.query_one(GraphView) + if app._active in ("listing", "disasm"): + 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._active == "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._active == "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.""" @@ -129,6 +231,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: @@ -145,6 +253,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.)""" @@ -168,12 +284,14 @@ def snapshot(app) -> dict[str, Any]: pass return { "active": app._active, - "pref": app._pref, + "pref": app._code_mode(), # kept for wire compat; a constant now "function": ({"ea": cur.ea, "name": cur.name} if cur else None), "cursor": _cursor_info(app, w), "status": st, "filter": app._filter_term, + "binary": app._binary, # None outside project mode "nav_depth": len(app._nav), + "hops": list(getattr(app, "_hops", [])), "dirty": bool(app._dirty), "modal": _modal_snapshot(app), **_readiness(app), @@ -187,6 +305,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)) @@ -220,21 +343,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)): @@ -246,9 +406,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 @@ -310,7 +468,7 @@ def disassembly(app, target=None, max_lines: int = 2000) -> dict[str, Any]: def _xref_dicts(xs, limit: int) -> list[dict[str, Any]]: - return [{"frm": x.frm, "to": x.to, "type": x.type, + return [{"frm": x.frm, "to": x.to, "type": x.type, "kind": x.kind, "fn_addr": x.fn_addr, "fn_name": x.fn_name} for x in xs[:limit]] @@ -438,14 +596,94 @@ class RpcServer: result = await self._dispatch(method, params) return {"id": rid, "result": result} except Exception as e: # noqa: BLE001 — report, never kill the connection - return {"id": rid, "error": {"message": f"{type(e).__name__}: {e}"}} + # ``str(KeyError("msg"))`` returns ``repr("msg")`` (adds quotes), which + # mangles our friendly resolve messages; unwrap the single arg instead. + if isinstance(e, KeyError) and len(e.args) == 1 and isinstance(e.args[0], str): + msg = e.args[0] + else: + msg = str(e) + 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._active == "graph": + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + if action == "close" and app._active != "graph": + return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} + want = "graph" if action in ("open", "toggle") and \ + app._active != "graph" else None + res = await self._press( + ["space"], + (lambda: app._active == "graph") if want else + (lambda: app._active != "graph"), + timeout, f"graph {action}") + return {**res, "graph": graph_info(app, blocks=want_blocks)} + + if app._active != "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.""" @@ -455,12 +693,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.call, "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.call, "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 @@ -474,8 +801,33 @@ class RpcServer: 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: @@ -487,9 +839,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") @@ -520,6 +888,74 @@ class RpcServer: if method == "functions": return functions(app, params.get("filter"), int(params.get("limit", 50))) + # -- projects ------------------------------------------------------ # + if method == "trace": + if app._trace is None: + raise ValueError("no trace loaded (launch with --trace FILE)") + t = app._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) + app._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") + app._seek(first) + elif "step" in params: + n = int(params.get("step") or 1) + over = bool(params.get("over")) + for _ in range(abs(n)): + (app._step_over if over else app._step)(1 if n > 0 else -1) + await settle(app, timeout=float(params.get("timeout", 20.0))) + snap = snapshot(app) + snap["trace"] = {"idx": app._t, "length": t.length, + "pc": hex(t.ip(app._t)), + "changed": sorted(t.changed(app._t))} + return snap + + if method == "binaries": + if app._project is None: + raise ValueError("not a project session (launch with --project)") + counts = app._index.counts() if app._index is not None else {} + resident = set(app._pool.resident()) if app._pool is not None else set() + return {"active": app._binary, "hops": list(app._hops), + "binaries": [{"label": r.label, "source": r.source, + "active": r.label == app._binary, + "resident": r.label in resident, + "indexed": int(counts.get(r.label, 0))} + for r in app._project.refs]} + + if method == "switch": + if app._project is None: + raise ValueError("not a project session (launch with --project)") + label = str(params.get("binary") or params.get("label") or "") + if app._project.by_label(label) is None: + have = ", ".join(r.label for r in app._project.refs) + raise ValueError(f"no such binary {label!r} (have: {have})") + addr = params.get("addr") + if label == app._binary and addr is None: + return snapshot(app) + if addr is None: + app._switch_binary(label) + else: + # Same path a project search hit takes, so it records a hop and + # Esc comes back here. + app._switch_then_goto(label, int(str(addr), 0) + if isinstance(addr, str) else int(addr)) + await settle(app, lambda: app._binary == label + and app._func_index is not None + and app._func_index.complete, + timeout=float(params.get("timeout", 300.0))) + return snapshot(app) + # -- structured introspection (heavy: run off the UI loop) -------- # loop = asyncio.get_running_loop() if method == "pseudocode": @@ -563,15 +999,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._active == "hex": + # backslash leaves hex for the code view (which may be decomp). + await self._press(["backslash"], + lambda: app._active != "hex", timeout, + "leave the hex view") + if app._active == "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._active == "listing", + timeout, "switch to the listing") + if app._active != "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._active == "hex": + await self._press(["backslash"], lambda: app._active != "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) @@ -582,7 +1109,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": @@ -606,12 +1134,16 @@ class RpcServer: 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) + return await self._press(["backslash"], lambda: app._active == "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) @@ -622,7 +1154,8 @@ 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 == "close": return await self._press(["escape"], timeout=timeout) if method == "save": @@ -667,13 +1200,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/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/tui.py b/idatui/tui.py deleted file mode 100644 index b399ffe..0000000 --- a/idatui/tui.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Launcher for the idatui TUI. - - # attach to the single open session on a running server - python -m idatui.tui - - # attach to a specific session - python -m idatui.tui --db 80d83396 - - # open (or reopen) an arbitrary binary, then drive it - python -m idatui.tui --open /path/to/binary - -The server (supervisor) must already be running (see spawn.sh). --open creates a -session via idb_open; the binary's directory must be writable (idalib writes a -.i64 next to it). -""" -from __future__ import annotations - -import argparse -import os - -from .app import IdaTui -from .client import DEFAULT_URL - - -def main(argv: list[str] | None = None) -> int: - p = argparse.ArgumentParser(prog="idatui", description="Minimal TUI for IDA over MCP") - p.add_argument("--url", default=os.environ.get("IDA_MCP_URL", DEFAULT_URL), - help=f"MCP server URL (default {DEFAULT_URL})") - p.add_argument("--db", default=os.environ.get("IDA_MCP_DB"), - help="attach to an existing session id") - p.add_argument("--open", metavar="PATH", - help="open (or reopen) a binary and drive it (dir must be writable)") - p.add_argument("--no-keepalive", action="store_true", - help="Do not bump idle-TTL / run the keepalive heartbeat") - p.add_argument("--rpc", metavar="PATH", - help="listen for RPC on this unix socket path (puppeteer the TUI)") - args = p.parse_args(argv) - rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None - IdaTui(url=args.url, db=args.db, open_path=args.open, - keepalive=not args.no_keepalive, rpc_path=rpc_path).run() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/idatui/worker.py b/idatui/worker.py new file mode 100644 index 0000000..a4e3509 --- /dev/null +++ b/idatui/worker.py @@ -0,0 +1,233 @@ +"""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 new file mode 100644 index 0000000..79a6db9 --- /dev/null +++ b/idatui/worker_client.py @@ -0,0 +1,234 @@ +"""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() |