"""Domain / paging layer: address-centric models over the raw MCP client. This is where the "millions of lines" problem is solved, so the TUI widgets only ever see a viewport-sized slice. Every hard-won constraint from ``docs/PAGING_FINDINGS.md`` is encoded here: * Per-call caps are silent (over the cap the server returns 10, not a clamp), so we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. * ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. * ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly is **block-cached** (revisits are free) and **prefetches** the next block on a background thread (the client is concurrency-safe). * ``include_total`` scans the whole function (~200ms on monsters); totals are fetched once and cached. * ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is null); that is surfaced as data, not an exception. Everything here is synchronous and thread-safe. The TUI runs these calls from Textual worker threads; the internal prefetch pool is separate and small. """ from __future__ import annotations import bisect import json import re import threading import urllib.request from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace from typing import Callable, TYPE_CHECKING 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 = 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*$") # --------------------------------------------------------------------------- # # Value models # --------------------------------------------------------------------------- # def _as_int(v) -> int: if isinstance(v, int): return v return int(v, 16) if isinstance(v, str) and v.startswith("0x") else int(v, 16) @dataclass(frozen=True) class Func: addr: int name: str size: int @classmethod def from_raw(cls, d: dict) -> "Func": addr = _as_int(d["addr"]) name = d.get("name") # An unnamed function (server returns null/empty) must still have a # 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) class Line: """One rendered disassembly line. ``ea`` is the instruction address.""" ea: int text: str label: str | None = None raw: bytes | None = None # opcode bytes; filled in by DisasmModel post-fetch @classmethod def from_raw(cls, d: dict) -> "Line": return cls( ea=_as_int(d["addr"]), text=d.get("instruction", ""), label=d.get("label"), ) @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) @property def label(self) -> str | None: # Line-compatible alias return self.name @classmethod def from_raw(cls, d: dict) -> "Head": 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"), ) @dataclass class Ref: addr: int name: str string: str | None = None @dataclass class Xref: frm: int # the referencing address to: int | None # the referenced address 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) class LVar: name: str type: str is_arg: bool @dataclass class FuncTypes: addr: int name: str prototype: str # e.g. 'int __fastcall foo(int a, char *b)' lvars: list[LVar] @dataclass(frozen=True) class Struct: name: str size: int is_union: bool members: int # field count ordinal: int @classmethod def from_raw(cls, d: dict) -> "Struct": return cls( name=d.get("name", ""), size=int(d.get("size", 0) or 0), is_union=bool(d.get("is_union", False)), members=int(d.get("cardinality", 0) or 0), ordinal=int(d.get("ordinal", 0) or 0), ) @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 code: str | None failed: bool error: str | None truncated: bool total_chars: int | None refs: list[Ref] = field(default_factory=list) # --------------------------------------------------------------------------- # # Query-payload unwrapping (shape: {"result":[{"data":[...],"next_offset":N}]}) # --------------------------------------------------------------------------- # def _query_data(payload) -> list: res = payload.get("result", payload) if isinstance(payload, dict) else payload if isinstance(res, list): res = res[0] if res else {} return res.get("data", []) if isinstance(res, dict) else [] # --------------------------------------------------------------------------- # # Function index: lazy, clamped, cached, filterable # --------------------------------------------------------------------------- # class FunctionIndex: """A lazily-paginated, cached view of the function list. Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by ``next_offset``). A single index instance corresponds to one server-side ``filter`` glob (``None`` = all functions). """ def __init__(self, program: "Program", filter: str | None = None): self._prog = program self.filter = filter self._funcs: list[Func] = [] self._by_addr: dict[int, Func] = {} self._done = False self._lock = threading.Lock() def _load_next_page(self) -> int: with self._lock: if self._done: return 0 offset = len(self._funcs) query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter data = _query_data(self._prog.client.call("list_funcs", queries=[query])) added = 0 with self._lock: for d in data: f = Func.from_raw(d) if f.addr not in self._by_addr: self._by_addr[f.addr] = f self._funcs.append(f) added += 1 if len(data) < LIST_PAGE: self._done = True return len(data) def load_next_page(self) -> int: """Load one more page; return the number of rows fetched (0 at end).""" return self._load_next_page() def ensure(self, n: int) -> None: """Ensure at least ``n`` functions are loaded (or all, if fewer exist).""" while not self._done and len(self._funcs) < n: if self._load_next_page() == 0: break def load_all(self, progress: Callable[[int], None] | None = None) -> None: """Fully enumerate (background-friendly). ~3s for 10k funcs.""" while not self._done: self._load_next_page() if progress: progress(len(self._funcs)) @property def complete(self) -> bool: return self._done def __len__(self) -> int: return len(self._funcs) def loaded(self) -> int: return len(self._funcs) def get(self, i: int) -> Func | None: self.ensure(i + 1) with self._lock: return self._funcs[i] if 0 <= i < len(self._funcs) else None def window(self, start: int, count: int) -> list[Func]: """Viewport slice ``[start, start+count)`` (loads as needed).""" self.ensure(start + count) with self._lock: return list(self._funcs[start : start + count]) def by_addr(self, ea: int) -> Func | None: with self._lock: return self._by_addr.get(ea) def all_loaded(self) -> list[Func]: with self._lock: return list(self._funcs) def update_name(self, addr: int, new_name: str) -> None: """Reflect a rename in the cached index (Func is frozen -> replace).""" with self._lock: old = self._by_addr.get(addr) if old is None: return nf = Func(addr=old.addr, name=new_name, size=old.size) self._by_addr[addr] = nf try: self._funcs[self._funcs.index(old)] = nf except ValueError: pass # --------------------------------------------------------------------------- # # Disassembly model: block-cached windowed listing for ONE function # --------------------------------------------------------------------------- # class DisasmModel: """Windowed, block-cached disassembly of a single function. Because ``disasm offset=N`` is O(N) and uncacheable server-side, we fetch and cache fixed ``DISASM_BLOCK``-sized blocks; a viewport read slices from cached blocks and returns instantly on revisit. The block just past the viewport is prefetched on a background thread. """ BLOCK = DISASM_BLOCK def __init__(self, program: "Program", ea: int, name: str | None = None): self._prog = program self.ea = ea self.name = name self._blocks: dict[int, list[Line]] = {} self._total: int | None = None self._ea_list: list[int] | None = None self._max_raw = 0 # widest opcode length seen (bytes) self._func_end: int | None = None self._func_end_done = False 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: """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( "disasm", addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") if total is None: total = payload.get("instruction_count", 0) with self._lock: self._total = int(total) return self._total def _function_end(self) -> int | None: """End address (exclusive) of this function; used to size the final instruction's opcode bytes. Cached (one lookup per model).""" if self._func_end_done: return self._func_end end: int | None = None try: fn = self._prog.function_of(self.ea) if fn is not None: end = fn.addr + fn.size except Exception: # noqa: BLE001 -- best-effort; falls back to a width guess end = None with self._lock: self._func_end = end self._func_end_done = True return end def _attach_bytes(self, lines: list[Line], end_ea: int | None) -> list[Line]: """Read the opcode bytes for ``lines`` in one request and slice them per instruction using consecutive addresses (variable-length safe).""" if not lines: return lines last_end = end_ea if last_end is None or last_end <= lines[-1].ea: last_end = lines[-1].ea + 15 # x86 max insn len; only the final line start = lines[0].ea data = self._prog.read_bytes(start, last_end - start) out: list[Line] = [] biggest = 0 for i, ln in enumerate(lines): nxt = lines[i + 1].ea if i + 1 < len(lines) else last_end length = max(nxt - ln.ea, 0) off = ln.ea - start b = bytes(data[off:off + length]) biggest = max(biggest, len(b)) out.append(replace(ln, raw=b)) with self._lock: if biggest > self._max_raw: 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]: # 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( "heads", addr=hex(self.ea), offset=b * self.BLOCK, count=self.BLOCK + 1, **self._end_kw(), ) 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 else: # this block ends the function end_ea = self._function_end() lines = self._attach_bytes(lines, end_ea) with self._lock: self._blocks[b] = lines self._inflight.discard(b) return lines def max_raw_len(self) -> int: """Widest opcode length (bytes) across the blocks fetched so far.""" with self._lock: return self._max_raw def scan_bytes(self) -> int: """Fetch every block (populating opcode bytes) and return the widest instruction length across the whole function. Used to size the opcode column so its padding doesn't jump as the listing streams in.""" total = self.total() off = 0 while off < total: got = self.lines(off, self.BLOCK, prefetch=False) if not got: break off += len(got) return self.max_raw_len() def _get_block(self, b: int) -> list[Line]: with self._lock: hit = self._blocks.get(b) if hit is not None: return hit return self._fetch_block(b) def _prefetch_block(self, b: int) -> None: if b < 0: return with self._lock: if b in self._blocks or b in self._inflight: return self._inflight.add(b) self._prog.submit(self._fetch_block, b) def lines(self, start: int, count: int, prefetch: bool = True) -> list[Line]: """Return rendered lines for viewport ``[start, start+count)``.""" if count <= 0 or start < 0: return [] end = start + count b0, b1 = start // self.BLOCK, (end - 1) // self.BLOCK out: list[Line] = [] for b in range(b0, b1 + 1): block = self._get_block(b) lo = start - b * self.BLOCK if b == b0 else 0 hi = end - b * self.BLOCK if b == b1 else self.BLOCK out.extend(block[max(lo, 0):hi]) if prefetch: self._prefetch_block(b1 + 1) # forward scroll self._prefetch_block(b0 - 1) # backward scroll return out def cached_line(self, idx: int) -> Line | None: """Non-blocking single-line peek: return the cached Line or None. Never touches the network — used by the virtualized view's render path.""" if idx < 0: return None b = idx // self.BLOCK off = idx - b * self.BLOCK with self._lock: block = self._blocks.get(b) if block is None or off >= len(block): return None return block[off] def is_cached(self, start: int, count: int) -> bool: """True if every block covering [start, start+count) is already cached.""" if count <= 0: return True b0, b1 = start // self.BLOCK, (start + count - 1) // self.BLOCK with self._lock: return all(b in self._blocks for b in range(b0, b1 + 1)) def ensure_async(self, start: int, count: int) -> None: """Schedule background fetches for any missing blocks in the range (non-blocking). Safe to call every render.""" if count <= 0: return b0, b1 = start // self.BLOCK, (start + count - 1) // self.BLOCK for b in range(b0, b1 + 1): self._prefetch_block(b) def cached_blocks(self) -> int: with self._lock: return len(self._blocks) def ensure_ea_index(self) -> list[int]: """Build (once) a sorted list of every line's ea, for ea->index lookup. Fetches the whole function; cached. Only needed for mid-function jumps.""" if self._ea_list is not None: return self._ea_list total = self.total() eas: list[int] = [] off = 0 while off < total: lines = self.lines(off, self.BLOCK, prefetch=False) if not lines: break eas.extend(ln.ea for ln in lines) off += len(lines) with self._lock: self._ea_list = eas return eas def index_of_ea(self, ea: int) -> int: """Instruction index of the line at/containing ``ea`` (0 if before start).""" eas = self.ensure_ea_index() i = bisect.bisect_right(eas, ea) - 1 return i if 0 <= i < len(eas) else 0 def invalidate(self) -> None: with self._lock: self._blocks.clear() self._total = None self._ea_list = None self._max_raw = 0 # --------------------------------------------------------------------------- # # 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] = {} 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: base = len(self._heads) for i, h in enumerate(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, base + i) self._heads.append(h) nxt = cur.get("next") if nxt is None: self._done = True self._next = None else: self._next = _as_int(nxt) return len(rows) def ensure(self, n: int) -> None: """Ensure at least ``n`` heads are loaded (or all, if fewer exist).""" while not self._done and len(self._heads) < 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 = len(self._heads) last_ea = self._heads[-1].ea 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: heads = self._heads for i, h in enumerate(heads): if h.ea <= ea < h.ea + max(h.size, 1): return i if h.ea > ea: return 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(len(self._heads)) @property def complete(self) -> bool: with self._lock: return self._done def loaded(self) -> int: with self._lock: return len(self._heads) def __len__(self) -> int: return self.loaded() def get(self, i: int) -> Head | None: with self._lock: return self._heads[i] if 0 <= i < len(self._heads) else None def window(self, start: int, count: int) -> list[Head]: self.ensure(start + count) with self._lock: return list(self._heads[start:start + count]) def index_of_ea(self, ea: int) -> int: with self._lock: return self._by_ea.get(ea, -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 <= len(self._heads) 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: """Windowed, block-cached raw bytes of the loaded image, addressed by virtual address. Format-agnostic: the range/segments come from IDA, not from any file header. Gaps between segments read back as zeros.""" BLOCK = HEX_BLOCK def __init__(self, program: "Program", start: int, end: int): self._prog = program self.start = start self.end = end self.size = max(end - start, 0) self._blocks: dict[int, bytes] = {} self._lock = threading.Lock() self._inflight: set[int] = set() def total_rows(self) -> int: return (self.size + 15) // 16 def file_offset(self, va: int) -> int | None: return self._prog.file_offset(va) def _fetch_block(self, b: int) -> bytes: addr = self.start + b * self.BLOCK n = min(self.BLOCK, self.end - addr) data = self._prog.read_bytes(addr, n) if n > 0 else b"" with self._lock: self._blocks[b] = data self._inflight.discard(b) return data def _prefetch(self, b: int) -> None: if b < 0 or b * self.BLOCK >= self.size: return with self._lock: if b in self._blocks or b in self._inflight: return self._inflight.add(b) self._prog.submit(self._fetch_block, b) def row(self, r: int, prefetch: bool = True) -> tuple[int, bytes | None]: """Return (va, bytes<=16) for row ``r``, or (va, None) if not yet cached. Non-blocking; used by the virtualized view's render path.""" off = r * 16 va = self.start + off b = off // self.BLOCK with self._lock: block = self._blocks.get(b) if block is None: return (va, None) bo = off - b * self.BLOCK return (va, block[bo:bo + 16]) def ensure(self, r0: int, count: int) -> None: """Blocking: fetch the blocks covering rows [r0, r0+count) if missing.""" if count <= 0: return b0 = (r0 * 16) // self.BLOCK b1 = ((r0 + count) * 16) // self.BLOCK for b in range(b0, b1 + 1): with self._lock: have = b in self._blocks if not have: self._fetch_block(b) def is_cached(self, r0: int, count: int) -> bool: if count <= 0: return True b0 = (r0 * 16) // self.BLOCK b1 = ((r0 + count - 1) * 16) // self.BLOCK with self._lock: return all(b in self._blocks for b in range(b0, b1 + 1)) def ensure_async(self, r0: int, count: int) -> None: b0 = (r0 * 16) // self.BLOCK b1 = ((r0 + max(count, 1) - 1) * 16) // self.BLOCK for b in range(b0 - 1, b1 + 2): self._prefetch(b) # --------------------------------------------------------------------------- # # Program: top-level handle, model registry, prefetch pool # --------------------------------------------------------------------------- # class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" 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]] = {} self._decomp_maps: dict[int, tuple[list[list[int]], int]] = {} # line->ea sets 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 ------------------------------------------------- # def submit(self, fn, *args) -> None: try: self._pool.submit(fn, *args) except RuntimeError: pass # pool shut down def close(self) -> None: self._pool.shutdown(wait=False, cancel_futures=True) # -- functions --------------------------------------------------------- # def functions(self, filter: str | None = None) -> FunctionIndex: with self._lock: idx = self._indices.get(filter) if idx is None: idx = FunctionIndex(self, filter) self._indices[filter] = idx 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).""" if self._sections is not None: return self._sections secs = [(s, e, nm) for s, e, _fo, nm in self._segments()] with self._lock: self._sections = secs return secs def image_range(self) -> tuple[int, int] | None: """[start, end) spanning all loaded segments (the hex view's document).""" secs = self.sections() if not secs: return None return (min(s[0] for s in secs), max(s[1] for s in secs)) def hex_model(self) -> "HexModel | None": """Cached block-backed byte view over the whole loaded image.""" if self._hexmodel is not None: return self._hexmodel rng = self.image_range() # calls sections() -> don't hold _lock (it re-locks) with self._lock: if self._hexmodel is None and rng is not None: self._hexmodel = HexModel(self, rng[0], rng[1]) return self._hexmodel def file_regions(self) -> list[tuple[int, int, int]]: """Sorted [(start, end, file_off)] mapping loaded segments to raw file offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs the injected ``file_regions`` server tool.""" if self._fileregions is not None: return self._fileregions regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] with self._lock: self._fileregions = regions return regions def file_offset(self, va: int) -> int | None: """Raw on-disk file offset for ``va``, or None if not file-backed.""" for start, end, fo in self.file_regions(): if start <= va < end: return (fo + (va - start)) if fo >= 0 else None return None def read_bytes(self, ea: int, n: int) -> bytes: """Raw bytes [ea, ea+n) from IDA (gaps read as zero). Fast path: the injected ``read_raw`` tool returns one contiguous hex string (C-speed both ends). Falls back to the stock ``get_bytes`` (a per-byte '0x..'-with-spaces string) on an older server without it. """ 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: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] data = res[0].get("data", "") if res and isinstance(res[0], dict) else "" out = bytearray() for tok in data.split(): try: out.append(int(tok, 16) & 0xFF) except ValueError: out.append(0) if len(out) < n: # pad short reads (unmapped tail) out.extend(b"\x00" * (n - len(out))) return bytes(out[:n]) 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] 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 by name.""" payload = self.client.call("search_structs", filter=filter) res = payload.get("result", []) if isinstance(payload, dict) else [] out = [Struct.from_raw(d) for d in res if isinstance(d, dict) and d.get("name") and not str(d["name"]).startswith("$")] # skip anonymous UDTs out.sort(key=lambda s: s.name.lower()) return out def struct_source(self, name: str) -> str: """A C definition for ``name`` reconstructed from its member layout (the server exposes members, not printable source). Faithful to IDA's field names/types; array dims are moved after the field name.""" payload = self.client.call( "type_inspect", queries=[{"name": name, "include_members": True}]) res = payload.get("result", []) if isinstance(payload, dict) else [] info = res[0] if res and isinstance(res[0], dict) else {} kw = "union" if info.get("is_union") else "struct" lines = [f"{kw} {name}", "{"] for m in info.get("members", []) or []: if not isinstance(m, dict): continue t = str(m.get("type", "")).strip() fn = m.get("name", "") base, arr = t, "" am = re.match(r"^(.*?)((?:\s*\[\d+\])+)\s*$", t) if am: base, arr = am.group(1).rstrip(), am.group(2).replace(" ", "") sep = "" if base.endswith("*") else " " lines.append(f" {base}{sep}{fn}{arr};") lines.append("};") return "\n".join(lines) def declare_type(self, decl: str) -> str | None: """Create or update a C type. Returns None on success, else the parse error. (Re-declaring a name updates it in place.)""" payload = self.client.call("declare_type", decls=decl) res = payload.get("result", []) if isinstance(payload, dict) else [] if res and isinstance(res[0], dict): return res[0].get("error") return None # -- function / variable types ---------------------------------------- # def func_types(self, ea: int) -> FuncTypes | None: """Structured decompiler types for the function at ``ea`` (prototype + local variables). None if ``ea`` isn't a decompilable function. Requires the injected ``func_types`` server tool.""" try: r = self.client.call("func_types", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): return None lvars = [LVar(name=lv.get("name", ""), type=lv.get("type", ""), is_arg=bool(lv.get("is_arg"))) for lv in r.get("lvars", []) if isinstance(lv, dict)] return FuncTypes(addr=_as_int(r.get("addr", hex(ea))), name=r.get("name", ""), prototype=r.get("prototype", ""), lvars=lvars) def set_function_type(self, ea: int, signature: str) -> str | None: """Set a function's prototype. None on success, else an error string.""" r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}]) 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 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.""" r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) if isinstance(r, dict) and r.get("error"): return r["error"] if isinstance(r, dict) and not r.get("ok"): return "failed to set the variable type" return None def delete_type(self, name: str) -> str | None: """Delete a named type. Returns None on success, else an error string. Requires a server-side ``del_type`` tool; if absent, a clear message is returned instead of raising.""" try: self.client.call("del_type", name=name) return None except IDAToolError as e: msg = e.message if "not found" in msg.lower() and "del_type" in msg: return "delete needs a 'del_type' tool on the ida-pro-mcp server" return msg # -- disassembly ------------------------------------------------------- # def disasm(self, ea: int, name: str | None = None) -> DisasmModel: with self._lock: m = self._disasm.get(ea) if m is None: m = DisasmModel(self, ea, name) self._disasm[ea] = m return m # -- decompilation ----------------------------------------------------- # def decompile(self, ea: int, refresh: bool = False) -> Decompilation: """Full pseudocode for a function. The server truncates responses over 50KB (strings clipped to 1000 chars) but caches the full output and exposes it at ``_meta.ida_mcp.download_url``. We transparently fetch that so the view always gets the complete body, not a 1KB stub. """ if not refresh: with self._lock: hit = self._decomp.get(ea) gen = self._name_gen if hit is not None: dec, hit_gen = hit if hit_gen == gen: return dec # Cached before a rename: names may be stale. Drop the server's # Hex-Rays cache so the refetch reflects the new names. try: self.client.call("force_recompile", items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 pass # Bound the decompile: a function Hex-Rays can't handle tends to stall # near the client's default 30s timeout, and the transport retries a # dropped connection up to max_retries+1 times, re-running the failing # decompile each time. Cap it so the worst case stays well under the # rpcclient socket timeout, and cache the failure below so a re-request # returns instantly instead of re-grinding. try: envelope = self.client.call_envelope( "decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT ) except Exception as e: # noqa: BLE001 -- surface as a failed decompile dec = Decompilation(ea, None, True, f"decompile error: {e}", False, None) with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec result = envelope.get("result", {}) payload = result.get("structuredContent") if payload is None: # fall back to text content payload = self.client._extract_payload("decompile", result) meta = (result.get("_meta") or {}).get("ida_mcp") if isinstance(meta, dict) and meta.get("download_url"): full = self._fetch_output(meta["download_url"]) if isinstance(full, dict) and full.get("code"): payload = full dec = _parse_decompilation(ea, payload) with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec def bump_names(self) -> None: """Signal that symbol names changed (a rename). Disasm/listing names are live in the IDB, so clearing the cached rows is enough for those; decompilation is generation-checked and force-recompiled lazily.""" with self._lock: self._name_gen += 1 models = list(self._disasm.values()) self._listings.clear() # listing head rows cache names -> refetch 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() 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(f"define code @ {ea:#x}: {res['error']}") def define_func(self, ea: int) -> None: """Create a function starting at ``ea`` (IDA's 'p').""" res = self._first_result( self.client.call("define_func", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError(f"create function @ {ea:#x}: {res['error']}") 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(f"undefine @ {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( f"make data @ {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( f"make string @ {ea:#x}: {res.get('error') or 'rejected'}") return res.get("text", "") 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" @ {ea:#x}" @staticmethod def _fetch_output(url: str, timeout: float = 15.0): """GET the server's cached full-output blob (plain HTTP, not MCP).""" try: with urllib.request.urlopen(url, timeout=timeout) as r: return json.loads(r.read().decode("utf-8", "replace")) except Exception: # noqa: BLE001 -- fall back to the truncated preview return None def strings(self, min_len: int = 4, refresh: bool = False) -> list[StrLit]: """Every string literal in the binary (IDA's Shift+F12 list), paged in full and cached. ``[]`` if the tool is unavailable.""" 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 # -- cross-references & containing function --------------------------- # def function_of(self, ea: int) -> Func | None: """Return the function containing ``ea`` (resolves mid-function addrs).""" payload = self.client.call("lookup_funcs", queries=[hex(ea)]) res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None return Func.from_raw(fn) if fn else None def xrefs_from(self, ea: int) -> list[Xref]: payload = self.client.call( "xref_query", queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) return _parse_xrefs(payload) def xrefs_to(self, ea: int, limit: int = 2000) -> list[Xref]: 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 ------------------------------------------------ # def resolve(self, target: int | str) -> int: """Resolve an int/hex-string/symbol name to an address (ea).""" if isinstance(target, int): return target s = target.strip() if s.startswith("0x") or s.startswith("0X"): return int(s, 16) if re.fullmatch(r"[0-9a-fA-F]+", s): return int(s, 16) # 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: raise KeyError(f"cannot resolve {target!r}: {e}") from e 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}{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 and decompiler views. Returns the raw payload so the caller can surface a soft per-item error. The caller must invalidate/recompile to see it.""" return self.client.call("set_comments", items=[{"addr": hex(ea), "comment": text}]) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: """Drop caches for a function after a rename/comment/patch/etc.""" with self._lock: self._decomp.pop(ea, None) m = self._disasm.get(ea) if m is not None: m.invalidate() def invalidate_functions(self) -> None: """Drop the function index caches (after rename/define/undefine).""" with self._lock: self._indices.clear() def _parse_xrefs(payload) -> list[Xref]: res = payload.get("result", []) if isinstance(payload, dict) else [] if not res: return [] data = res[0].get("data", []) or [] out: list[Xref] = [] for d in data: if not isinstance(d, dict): continue fn = d.get("fn") or {} frm = d.get("from", d.get("addr")) to = d.get("to") out.append(Xref( frm=_as_int(frm) if frm is not None else 0, to=_as_int(to) if to is not None else None, 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 def _parse_decompilation(ea: int, payload) -> Decompilation: if not isinstance(payload, dict): return Decompilation(ea, None, True, "unexpected payload", False, None) code = payload.get("code") error = payload.get("error") if not code: return Decompilation(ea, None, True, error or "decompilation failed", False, None) m = _TRUNC_RE.search(code) truncated = m is not None total_chars = int(m.group(1)) if m else len(code) refs = [ Ref(addr=_as_int(r["addr"]), name=r.get("name", ""), string=r.get("string")) for r in payload.get("refs", []) if isinstance(r, dict) and "addr" in r ] return Decompilation(ea, code, False, error, truncated, total_chars, refs)