diff options
Diffstat (limited to 'server')
| -rw-r--r-- | server/patch_server.py | 406 |
1 files changed, 405 insertions, 1 deletions
diff --git a/server/patch_server.py b/server/patch_server.py index f42cc3a..6862785 100644 --- a/server/patch_server.py +++ b/server/patch_server.py @@ -1,6 +1,10 @@ #!/usr/bin/env python3 """Inject idatui's extra ida-pro-mcp tools into the installed server package. +DEPRECATED along with the ida-pro-mcp transport: the default backend is now the +idalib worker (idatui/worker.py), which registers these same tools in-process and +needs no patching. Kept only for `--backend mcp`; slated for removal. + ida-pro-mcp lacks a few tools idatui needs. Rather than vendor/fork the server, we keep the tool source here and inject it (idempotently) into the installed ``api_types.py``. That module is imported by every worker @@ -42,6 +46,25 @@ def _idatui_lv_get(x): @tool @idasync +def resolve_names( + queries: Annotated[list, "Symbol name(s) to resolve to their OWN address"], +) -> list: + """Resolve named locations (functions, labels like loc_/locret_, data) to the + exact address the NAME denotes, via get_name_ea. Unlike lookup_funcs, a + mid-function label resolves to the label's address, not the containing + function's entry.""" + import idaapi + qs = queries if isinstance(queries, list) else [queries] + out = [] + for q in qs: + q = str(q).strip() + ea = idaapi.get_name_ea(idaapi.BADADDR, q) + out.append({"query": q, "ea": (hex(ea) if ea != idaapi.BADADDR else None)}) + return out + + +@tool +@idasync def del_type( name: Annotated[str, "Local type name to delete (struct/union/enum/typedef)"], ) -> dict: @@ -175,9 +198,390 @@ def file_regions() -> dict: fo = -1 if fo < 0 or fo >= (1 << 48): fo = -1 - out.append({"start": hex(seg.start_ea), "end": hex(seg.end_ea), "file_off": fo}) + try: + nm = ida_segment.get_segm_name(seg) or "" + except Exception: + nm = "" + out.append({"start": hex(seg.start_ea), "end": hex(seg.end_ea), + "file_off": fo, "name": nm}) seg = ida_segment.get_next_seg(seg.start_ea) return {"regions": out} + + +@tool +@idasync +def make_string( + addr: Annotated[str, "Address of the string start"], + length: Annotated[int, "Length in bytes (0 = auto-detect to the terminator)"] = 0, + kind: Annotated[str, "String kind: c | c16 | c32 | pascal"] = "c", +) -> dict: + """Create a string literal at ``addr`` (IDA's 'A'). ``length`` 0 auto-detects + to the terminator. Undefines any items in the way first, like the UI does. + Returns the created byte size and the decoded contents.""" + import ida_bytes + import ida_nalt + + ea = parse_address(addr) + strtype = { + "c": ida_nalt.STRTYPE_C, + "c16": ida_nalt.STRTYPE_C_16, + "c32": ida_nalt.STRTYPE_C_32, + "pascal": ida_nalt.STRTYPE_PASCAL, + }.get(str(kind).lower(), ida_nalt.STRTYPE_C) + n = max(int(length), 0) + # Free any existing item(s) so create_strlit can carve the literal. + ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, n if n > 0 else 1) + ok = bool(ida_bytes.create_strlit(ea, n, strtype)) + if not ok: + return {"addr": addr, "ok": False, "error": "create_strlit failed"} + size = int(ida_bytes.get_item_size(ea)) + try: + raw = ida_bytes.get_strlit_contents(ea, -1, strtype) + text = raw.decode("utf-8", "replace") if raw else "" + except Exception: + text = "" + return {"addr": addr, "ok": True, "size": size, "text": text} + + +@tool +@idasync +def read_raw( + addr: Annotated[str, "Start address (hex or name)"], + size: Annotated[int, "Number of bytes to read"], +) -> dict: + """Read ``size`` bytes at ``addr`` as ONE contiguous lowercase hex string + (no per-byte '0x'/spaces). The hot path for the hex view and disasm opcode + bytes. + + Fast: does a single bulk ``ida_bytes.get_bytes`` (C-speed) instead of the + per-byte read_bytes_bss_safe loop (2 IDA calls/byte). Unloaded bytes come + back from IDA as the 0xFF sentinel, so we only re-check is_loaded for the + (usually sparse) 0xFF bytes and zero the genuinely-unloaded ones — matching + get_bytes' bss semantics without paying per-byte for the whole range. + + Encoding is compact hex (~2.5x smaller than get_bytes' '0x..'-with-spaces) + and, unlike get_bytes, does not truncate on large reads.""" + import ida_bytes + + ea = parse_address(addr) + n = max(int(size), 0) + if n == 0: + return {"addr": addr, "hex": "", "n": 0} + raw = ida_bytes.get_bytes(ea, n) + if raw is None or len(raw) < n: # nothing (or not all) mapped + base = bytearray(raw or b"") + base.extend(b"\\xff" * (n - len(base))) + raw = bytes(base) + ba = bytearray(raw) + # Only unloaded bytes read as 0xFF; correct just those to 0 (bss => zero). + i = ba.find(0xFF) + while i != -1: + if not ida_bytes.is_loaded(ea + i): + ba[i] = 0 + i = ba.find(0xFF, i + 1) + return {"addr": addr, "hex": bytes(ba).hex(), "n": len(ba)} + + +def _idatui_head_row(ea): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name.""" + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) + if ida_bytes.is_code(f): + kind = "code" + elif ida_bytes.is_data(f): + kind = "data" + else: + kind = "unknown" + line = ida_lines.generate_disasm_line(ea, 0) + text = ida_lines.tag_remove(line) if line else "" + text = " ".join(text.split()) # collapse IDA's column padding + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +def _idatui_unknown_row(ea, size): + """One collapsed row for a run of ``size`` undefined bytes starting at + ``ea``. A single byte is rendered normally (shows its value); a longer run + collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions + of one-byte rows.""" + import ida_name + + if size <= 1: + return _idatui_head_row(ea) + row = {"ea": hex(ea), "kind": "unknown", "size": int(size), + "text": f"db {size} dup(?)"} + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +def _idatui_struct_member_rows(ea): + """Indented member rows for a struct-typed data item at ``ea`` (expansion), + or [] if it isn't a struct. Top-level fields only.""" + import ida_nalt + import ida_typeinf + import idaapi + + tif = ida_typeinf.tinfo_t() + if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): + return [] + udt = ida_typeinf.udt_type_data_t() + if not tif.get_udt_details(udt): + return [] + rows = [] + for m in udt: + off = m.begin() // 8 + try: + mtype = m.type._print() or "" + except Exception: + mtype = "" + try: + sz = int(m.type.get_size()) + if sz == idaapi.BADSIZE: + sz = 0 + except Exception: + sz = 0 + name = m.name or "" + text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") + rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, + "text": text}) + return rows + + +def _idatui_func_header_rows(ea): + """IDA-style subroutine banner rows shown just before a function's entry.""" + import ida_funcs + + name = ida_funcs.get_func_name(ea) or "sub_%X" % ea + bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 + return [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + + +def _idatui_func_footer_rows(ea, func): + """End-of-function marker shown just after a function's last item.""" + import ida_funcs + + name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea + return [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + + +@tool +@idasync +def heads( + addr: Annotated[str, "Start address or name to walk from"], + count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, + offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, + end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", + back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, + annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, +) -> dict: + """Walk item heads from ``addr`` as a flat listing: every head is rendered + (code OR data OR undefined) via generate_disasm_line and stepped with + next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data + byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's + real disassembly view. Address-paged: page forward by re-calling with + ``addr`` = the returned cursor.next; page up with ``back=true``.""" + import ida_bytes + import ida_segment + import idaapi + + count = 2000 if count > 2000 else (1 if count < 1 else count) + offset = max(int(offset), 0) + try: + start = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} + seg = ida_segment.getseg(start) + if not seg: + return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} + lo, hi = seg.start_ea, seg.end_ea + if end: + try: + hi = min(hi, parse_address(end)) + except Exception: + pass + + rows = [] + if back: + # Collect up to (count+offset) heads strictly before `start`, then take + # the window closest to `start`, returned in forward order. + walk = [] + cur = ida_bytes.prev_head(start, lo) + while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: + walk.append(cur) + cur = ida_bytes.prev_head(cur, lo) + walk.reverse() + chosen = walk[: len(walk) - offset] if offset else walk + chosen = chosen[-count:] + rows = [_idatui_head_row(e) for e in chosen] + first = chosen[0] if chosen else start + pea = ida_bytes.prev_head(first, lo) + cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} + return {"addr": str(addr), "heads": rows, "cursor": cursor} + + # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a + # flat listing must show them (IDA renders undefined as `db ?` lines, and + # navigating to an unmarked address must land ON it). Defined items advance + # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its + # end found in O(1) via next_head, which skips undefined) so a large .bss or + # gap doesn't explode into millions of one-byte rows. + def _is_unknown(e): + f = ida_bytes.get_flags(e) + return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) + + def _run_end(e): + """End (exclusive) of the undefined run starting at ``e``.""" + nh = ida_bytes.next_head(e, hi) + return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi + + def _advance(e): + if _is_unknown(e): + return _run_end(e) + nxt = ida_bytes.get_item_end(e) + return nxt if nxt > e else e + 1 + + def _rows_for(e): + if _is_unknown(e): + return [_idatui_unknown_row(e, _run_end(e) - e)] + func = idaapi.get_func(e) if annotate else None + at_start = func is not None and func.start_ea == e + out = [] + if at_start: + out.extend(_idatui_func_header_rows(e)) + row = _idatui_head_row(e) + if at_start: + row = dict(row) + row["name"] = None # the name is shown on the proc header line + elif annotate and row.get("kind") == "code" and row.get("name"): + # A code label (loc_XXX/jump target) gets its OWN line at depth 0, + # like IDA; strip it from the instruction row below. + nm = row["name"] + out.append({"ea": hex(e), "kind": "label", "size": 0, + "text": nm + ":", "name": nm}) + row = dict(row) + row["name"] = None + out.append(row) + if row.get("kind") == "data": + out.extend(_idatui_struct_member_rows(e)) # expand struct fields + if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: + out.extend(_idatui_func_footer_rows(e, func)) + return out + + ea = ida_bytes.get_item_head(start) + for _ in range(offset): + if ea >= hi or ea == idaapi.BADADDR: + break + ea = _advance(ea) + more = False + while ea != idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True + break + rows.extend(_rows_for(ea)) # a struct head expands into member rows + ea = _advance(ea) + cursor = {"next": hex(ea)} if more else {"done": True} + return {"addr": str(addr), "heads": rows, "cursor": cursor} + + +@tool +@idasync +def xref_types( + queries: Annotated[list, "[{addr, direction:'to'|'from'|'both', include_fn, dedup, count}]"], +) -> dict: + """Like xref_query, but every row carries a fine-grained ``kind`` derived from + the IDA xref type \u2014 call/jump/flow for code, read/write/offset/text/info for + data \u2014 alongside the coarse ``type`` (code/data). Feeds the xref dialog's + r/w/call badges. Same query/envelope shape as xref_query.""" + import idaapi, idautils, ida_funcs, ida_bytes, ida_xref + code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", + ida_xref.fl_JF: "jump", ida_xref.fl_JN: "jump", + ida_xref.fl_F: "flow"} + data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", + ida_xref.dr_R: "read", ida_xref.dr_T: "text", ida_xref.dr_I: "info"} + + def _kind(xr): + table = code_kind if xr.iscode else data_kind + return table.get(xr.type, "code" if xr.iscode else "data") + + def _fn(ea): + f = ida_funcs.get_func(ea) + if not f: + return None + return {"addr": hex(f.start_ea), "name": ida_funcs.get_func_name(f.start_ea)} + + def _resolve(raw): + raw = str(raw).strip() + try: + return int(raw, 16) # handles '0x2490' and '2490' + except ValueError: + return idaapi.get_name_ea(idaapi.BADADDR, raw) + + qs = queries if isinstance(queries, list) else [queries] + result = [] + for q in qs: + q = q if isinstance(q, dict) else {"addr": q} + raw = str(q.get("addr", "")).strip() + direction = str(q.get("direction", "to") or "to").lower() + include_fn = bool(q.get("include_fn", True)) + dedup = bool(q.get("dedup", True)) + try: + count = int(q.get("count", 2000) or 2000) + except (TypeError, ValueError): + count = 2000 + target = _resolve(raw) + rows = [] + if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + row = {"direction": "to", "addr": hex(xr.frm), "from": hex(xr.frm), + "to": hex(target), "type": "code" if xr.iscode else "data", + "kind": _kind(xr)} + if include_fn: + row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + row = {"direction": "from", "addr": hex(xr.to), "from": hex(target), + "to": hex(xr.to), "type": "code" if xr.iscode else "data", + "kind": _kind(xr)} + if include_fn: + row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen = set() + deduped = [] + for r in rows: + k = (r["direction"], r["from"], r["to"], r["kind"]) + if k in seen: + continue + seen.add(k) + deduped.append(r) + rows = deduped + rows = rows[:count] + result.append({"query": raw, "data": rows, "next_offset": None}) + return {"result": result} ''' SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n" |