diff options
Diffstat (limited to 'idatui')
| -rw-r--r-- | idatui/app.py | 767 | ||||
| -rw-r--r-- | idatui/launch.py | 6 | ||||
| -rw-r--r-- | idatui/pane.py | 26 | ||||
| -rw-r--r-- | idatui/rpc.py | 165 | ||||
| -rw-r--r-- | idatui/trace.py | 495 |
5 files changed, 1424 insertions, 35 deletions
diff --git a/idatui/app.py b/idatui/app.py index ea68620..1edc65e 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -91,6 +91,18 @@ _ASM_KEYWORDS = frozenset({ "gs", "ss", "align", "public", "assume", "end", }) _S_CURSOR = Style(bgcolor="#2a313c") +#: Execution trails. Deliberately faint: they sit UNDER the code palette and +#: must not compete with it — the trail says "you came through here", the text +#: still has to be readable as code. Now is the loudest because there is exactly +#: one of it. +_S_TRAIL_NOW = Style(bgcolor="#3f3410") +_S_TRAIL_PAST = Style(bgcolor="#2b1c17") # warm: behind you +_S_TRAIL_FUTURE = Style(bgcolor="#152230") # cool: ahead of you +#: Hex with a trace loaded: bytes the trace SAW at this timestamp vs bytes we're +#: still showing from the file. The distinction matters more than the values — +#: one is evidence, the other is an assumption. +_S_HEX_LIVE = Style(color="#9ece6a") +_S_HEX_STALE = Style(color="#5e6875") _S_DIM = Style(color="#7c8b9e", italic=True) _S_MATCH = Style(bgcolor="#7a5c00") # all search matches _S_MATCH_CUR = Style(bgcolor="#d0a215", color="#12161c") # the current match @@ -777,6 +789,8 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._search_loading = False self._search_pending: list = [] # done-callbacks awaiting the load self._link_rows: set[int] = set() # split-view: linked instruction rows + #: {address: 'now'|'past'|'future'} painted under the code (trace mode). + self.trail: dict[int, str] = {} # -- text helpers ------------------------------------------------------ # def _head(self, idx: int) -> Head | None: @@ -1045,6 +1059,12 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru 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)) @@ -1253,6 +1273,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self._gutter = 0 # line-number gutter width (cells) self._line_eas: list[int | None] = [] # per-line address (marker stripped) self._link_line: int | None = None # split-view: linked pseudocode line + #: {line index: 'now'|'past'|'future'} — the execution trail, mapped from + #: instructions onto pseudocode via decomp_map. + self.trail: dict[int, str] = {} self._term = "" self._matches: list[int] = [] self._ranges: dict[int, list[tuple[int, int]]] = {} @@ -1386,6 +1409,11 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True base = self._strips[idx] if linked: base = base.apply_style(_S_LINK) # split-view companion band + kind = self.trail.get(idx) if self.trail else None + if kind is not None: + base = base.apply_style( + _S_TRAIL_NOW if kind == "now" else + _S_TRAIL_PAST if kind == "past" else _S_TRAIL_FUTURE) if idx in self._ranges: base = _overlay_ranges(base, self._ranges[idx], self._match_style(idx)) if self._hl_word: @@ -1538,6 +1566,10 @@ class HexView(ScrollView, can_focus=True): super().__init__(id="hex") self.model = None self.total = 0 + #: Trace to read memory from, and the timestamp to read it at. When set, + #: the dump shows what memory HELD then rather than what the file holds. + self.trace = None + self.trace_idx = 0 self._internal_top: int | None = None # scroll target we set ourselves # -- public API -------------------------------------------------------- # @@ -1724,6 +1756,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] = [ @@ -1738,15 +1779,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)) @@ -2286,6 +2341,221 @@ class HelpScreen(ModalScreen): self.dismiss(None) +class RegWriteScreen(ModalScreen): + """Registers, and the instruction that set each one. + + "Which instruction set this register to its current value?" is the question + a trace exists to answer, and seeking backwards to it is a single keypress + here rather than a manual walk. Forward is offered too, but backward is what + people actually want — you notice a bad value after it has been used. + """ + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + Binding("enter", "choose", show=False, priority=True), + Binding("f", "choose_forward", show=False), + ] + + def __init__(self, rows, idx: int) -> None: + super().__init__() + self._rows = rows # (name, value, last_write, next_write) + self._idx = idx + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box"): + 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. @@ -3016,7 +3286,7 @@ class IdaTui(App): #xref-list { height: auto; max-height: 100%; } /* every #pal-box palette centres, not just the symbol one */ SymbolPalette, StringsPalette, ProjectPalette, - LoadOptionsScreen { align: center middle; } + 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; } @@ -3026,6 +3296,11 @@ class IdaTui(App): #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 @@ -3071,6 +3346,18 @@ class IdaTui(App): 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), Binding("f1", "help", "Keys", show=False), Binding("g", "goto", "Goto"), Binding("slash", "filter", "Filter", show=False), @@ -3082,7 +3369,7 @@ class IdaTui(App): def __init__(self, open_path: str | None = None, keepalive: bool = True, rpc_path: str | None = None, ttl: int = 1800, - project=None, load_args: str = "") -> None: + project=None, load_args: str = "", trace_path: str = "") -> None: super().__init__() # Project mode is additive: with no project this is the plain # single-binary app, unchanged. @@ -3114,6 +3401,16 @@ class IdaTui(App): 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 @@ -3168,6 +3465,11 @@ class IdaTui(App): hx = HexView() hx.display = False yield hx + # 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 @@ -3438,8 +3740,12 @@ class IdaTui(App): self._flash_until = _time.monotonic() + 8.0 elif self._flash and _time.monotonic() < self._flash_until: text = self._flash - if self._binary: # project mode: always say which binary you're in - text = f"[{self._binary}] {text}" + # 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. @@ -3585,6 +3891,7 @@ class IdaTui(App): 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( @@ -3607,7 +3914,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) @@ -3615,18 +3921,20 @@ class IdaTui(App): if rows: self.app.call_from_thread(self._append_rows, rows) self.app.call_from_thread( - self._status, f"{module} — {last} functions…" + self._status, f"{last} functions…" ) # If a filter is active (typed during load), re-apply it over the full set. if self._filter_term: self.app.call_from_thread(self._apply_filter, self._filter_term) else: self.app.call_from_thread( - self._status, f"{module} — {len(idx)} functions (Ctrl+N: find symbol)") + self._status, f"{len(idx)} functions (Ctrl+N: find symbol)") # Land somewhere useful instead of an empty pane: main() if present, # otherwise pop the fuzzy symbol picker. self.app.call_from_thread(self._auto_land) self._index_binary() # project mode: keep the cross-binary index fresh + if self._trace_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: @@ -3726,7 +4034,19 @@ class IdaTui(App): if fn is not None: self._open_function(fn.addr, fn.name) elif len(self._func_index): - self.action_symbols() + # No main(): land on the first function rather than pushing the + # symbol palette. A modal as the *startup* state leaves a human + # staring at a picker over an empty pane, and silently swallows + # every keystroke an RPC driver injects while `ping` still says + # ready:true. Landing somewhere real is better for both; Ctrl+N is + # one keypress away. + first = self._func_index.get(0) + if first is not None: + self._open_function(first.addr, first.name) + self._status(f"no entry function — opened {first.name} " + "(Ctrl+N: find symbol)") + else: + self.action_symbols() else: self._land_without_functions() @@ -3837,7 +4157,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 @@ -4039,6 +4359,7 @@ class IdaTui(App): 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 "" @@ -5159,12 +5480,38 @@ class IdaTui(App): 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 = self.program.region_label(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) + 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) -> None: + 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)") @@ -5299,6 +5646,355 @@ class IdaTui(App): # 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. @@ -5381,7 +6077,8 @@ class IdaTui(App): idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 name = fn.name if fn is not None else self.program.region_label(ea) self.app.call_from_thread( - self._open_at, ea, name, idx, push, -1, 0, fn is None, focus_name) + self._open_at_if_current, seq, ea, name, idx, push, fn is None, + focus_name) def _open_decomp_entry(self, fn_addr: int, fn_name: str, dec_idx: int, dec_cursor_x: int, push: bool, @@ -5554,6 +6251,18 @@ class IdaTui(App): cur = row_of(fallback_ea) return (cur, row_of(a.top_ea)) + def _open_at_if_current(self, seq: int, ea: int, name: str, cursor: int, + push: bool, is_region: bool, + focus_name: str | None) -> None: + """Apply a navigation result only if it's still the one being awaited. + + The decompiler path has had this since 756589a; the listing path hadn't, + so a slow navigation could still land on top of a newer one. + """ + if seq != self._nav_seq: + return + self._open_at(ea, name, cursor, push, -1, 0, is_region, focus_name) + def _open_at(self, ea: int, name: str, cursor: int, push: bool, dec_cursor: int = -1, dec_cursor_x: int = 0, is_region: bool = False, focus_name: str | None = None, @@ -6167,8 +6876,18 @@ class IdaTui(App): self.app.call_from_thread(self._apply_split_map, ea, m) def _apply_split_map(self, ea: int, m: list) -> None: - if not self._split or self._cur is None or self._cur.ea != ea: - return # left split / navigated away + """Index the per-line instruction map for the decompiled function. + + Keyed to what the DECOMPILER holds, not to _cur, and not conditional on + split being on. The old guard dropped the result whenever _cur had moved + while the fetch was in flight — during trace stepping that is almost + always — leaving the split view working from the map of the function you + just left. _cur follows the cursor; this map describes the pseudocode on + screen, and those are different things. + """ + dec = self._try_view(DecompView) + if dec is not None and dec.loaded_ea is not None and ea != dec.loaded_ea: + return # a stale fetch for a function we no longer show self._split_eamap = m self._split_ea2line = {} alleas = [] @@ -6179,7 +6898,15 @@ class IdaTui(App): # ea span of the decompiled function: when the listing cursor leaves it, # _sync_split re-points the decomp to the function under the cursor. self._split_range = (min(alleas), max(alleas)) if alleas else None - self._sync_split(self._active) # re-link with the region map + # ONE index, shared with the trace path: it used to keep a parallel copy + # of exactly this, fetched separately and keyed differently, which is how + # the two ended up describing different functions. + self._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: dv = self._try_view(DecompView) diff --git a/idatui/launch.py b/idatui/launch.py index f51e5af..64e1546 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -64,6 +64,8 @@ def main(argv: list[str] | None = None) -> int: 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 " @@ -157,7 +159,9 @@ def main(argv: list[str] | None = None) -> int: rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None IdaTui(open_path=binary, keepalive=not args.no_keepalive, rpc_path=rpc_path, ttl=args.ttl, project=project, - load_args=_load_args(load)).run() + load_args=_load_args(load), + trace_path=(os.path.abspath(os.path.expanduser(args.trace)) + if args.trace else "")).run() return 0 diff --git a/idatui/pane.py b/idatui/pane.py index 53afef3..37a5f0c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -165,6 +165,8 @@ def spawn(args) -> int: inner += ["--rpc", sock] else: inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] + if getattr(args, "trace", None): + inner += ["--trace", os.path.abspath(os.path.expanduser(args.trace))] cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) split = ["split-window", "-v" if args.vertical else "-h", @@ -246,18 +248,29 @@ 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") + 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 + # 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): + time.sleep(0.25) if pane and _pane_alive(pane): subprocess.run(["tmux", "kill-pane", "-t", pane], capture_output=True) + killed.append(pane) if sock: try: os.unlink(sock) @@ -269,6 +282,12 @@ def stop(args) -> int: out = {"stopped": [r.get("sock") or r.get("pane") for r in rows]} if reaped: out["reaped_workers"] = reaped + if killed: + # Only ever reached on timeout: say so, because it means a save may have + # been cut short rather than "clean teardown". + out["force_killed"] = killed + out["warning"] = (f"pane(s) did not exit within {args.timeout}s and were " + "killed; unsaved database changes may be lost") print(json.dumps(out)) return 0 @@ -316,6 +335,8 @@ def main(argv: list[str]) -> int: sp = sub.add_parser("spawn", help="open a TUI pane and wait until ready") sp.add_argument("--open", metavar="PATH", help="binary to open (its dir must be writable)") + sp.add_argument("--trace", metavar="FILE", + help="Tenet execution trace to load alongside the binary") sp.add_argument("--project", metavar="FILE", help="project file to open instead of a single binary; " "any --open paths are added to it (created if absent)") @@ -331,6 +352,9 @@ 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") diff --git a/idatui/rpc.py b/idatui/rpc.py index 0921b89..717f4df 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -59,7 +59,7 @@ 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", @@ -72,6 +72,7 @@ METHODS = { "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)", @@ -103,6 +104,33 @@ def _active_widget(app): _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.""" @@ -147,6 +175,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.)""" @@ -451,9 +487,15 @@ class RpcServer: return {"id": rid, "error": {"message": f"{type(e).__name__}: {msg}"}} # -- composed helpers (semantic verbs) -------------------------------- # - async def _press(self, keys, pred=None, timeout=20.0): + async def _press(self, keys, pred=None, timeout=20.0, what=""): await self.app._press_keys([str(k) for k in keys]) - await settle(self.app, pred, timeout=timeout) + ok = await settle(self.app, pred, timeout=timeout) + if pred is not None and not ok: + # Never report success for an action that did not happen: the caller + # would go on to edit whatever the *previous* location was. + raise TimeoutError( + f"{what or 'action'} did not complete within {timeout}s " + f"(still at {_where(self.app)}); retry with a larger timeout=") return snapshot(self.app) async def _fill_prompt(self, open_key, input_id, value, delay_ms, clear): @@ -465,7 +507,14 @@ 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)) @@ -484,8 +533,32 @@ 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", + } + #: Modals the driver is expected to interact with (they have their own verbs). + _DRIVABLE_MODALS = {"XrefsScreen", "SymbolPalette", "StructEditor", + "ProjectPalette", "QuitScreen"} + + def _modal_kind(self) -> str | None: + scr = self.app.screen + name = type(scr).__name__ + return name if name in _MODALS or name in self._DRIVABLE_MODALS else None + async def _dispatch(self, method: str, params: dict[str, Any]) -> Any: app = self.app + params = _coerce_params(params) + if method in self._NEEDS_NO_MODAL: + modal = self._modal_kind() + if modal is not None: + raise RuntimeError( + f"modal {modal!r} is on top and will swallow this verb's " + f"keystrokes; dismiss it first (close) or use its own verb " + f"(select/symbols/xrefs). Note: a binary with no entry " + f"function can land in the symbol palette on startup.") if method in (None, "ping"): module = None try: @@ -497,9 +570,25 @@ class RpcServer: if method == "methods": return METHODS if method == "quit": + # Route through the same teardown a human gets, so a dirty database + # is written instead of dropped. `app.exit()` alone skips the dirty + # check entirely, and the caller (pane stop) then kills the pane -- + # which used to destroy a whole session's annotations. + dirty = list(app._dirty_labels()) + save = params.get("save", True) + if isinstance(save, str): + save = save.lower() not in ("0", "false", "no", "") + + def _go(): + if dirty and save: + app._on_quit_choice("save") # saves, then exits + else: + app._on_quit_choice("discard") + # answer first, then tear down (so this response still gets written) - asyncio.get_running_loop().call_later(0.2, app.exit) - return {"ok": True, "quitting": True} + asyncio.get_running_loop().call_later(0.2, _go) + return {"ok": True, "quitting": True, "saving": bool(dirty and save), + "dirty": dirty} if method in _PROGRAM_METHODS and app.program is None: raise ValueError("not ready: still connecting / loading functions") @@ -531,6 +620,38 @@ class RpcServer: 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)") @@ -609,7 +730,15 @@ 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 == "rename": @@ -617,7 +746,13 @@ class RpcServer: await settle(app, timeout=timeout) return snapshot(app) if method == "comment": - await self._fill_prompt("semicolon", "comment", str(params["text"]), delay, + # Comments can be long; skip the per-char delay so the agent isn't + # blocked for seconds watching the typing animation. Also: the + # prompt is single-line, so literal newlines (0x0a) get swallowed by + # the Input widget. The app's _do_comment converts the two-char + # sequence '\n' into a real newline for IDA, so we escape here. + ctext = str(params["text"]).replace("\n", "\\n") + await self._fill_prompt("semicolon", "comment", ctext, 0, clear=True) await settle(app, timeout=timeout) return snapshot(app) @@ -628,7 +763,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": @@ -652,12 +788,14 @@ 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 == "xrefs": return await self._press( - ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", timeout) + ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", + timeout, "xrefs") if method == "symbols": await app._press_keys(["ctrl+n"]) await settle(app, lambda: type(app.screen).__name__ == "SymbolPalette", timeout=10) @@ -668,7 +806,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": 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) |
