path: root/idatui/codemode_client.py

"""Client adapter from ida-tui's domain operations to IDA Code Mode.

``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered
GUI database, reuses a shared managed idalib worker, or starts one when needed.
The TUI never owns or terminates an IDA process.  Closing this client releases
only its lease.

The Code Mode transport intentionally exposes one broad operation,
``execute_python``.  ``CodeModeClient.invoke`` turns the small, address-centric
operations needed by the paging layer into self-contained snippets.  The
snippets prefer the public ``ida-domain`` ``db`` object.  A handful of features
that ida-domain does not currently expose (IDA-coloured listing rows, creating
instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the
IDAPython modules that Code Mode deliberately makes importable.
"""
from __future__ import annotations

import json
import os
import shlex
import threading
import time
from textwrap import dedent
from typing import Any

from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session

# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost.
#
# The paging/graph/trace layers and their offline test suites must keep importing
# `idatui` on a machine with no IDA and no Code Mode installed -- that is the
# house rule the stdlib-only worker client used to satisfy for free, and
# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level
# import here makes the whole package unimportable, so the failure is deferred to
# the first operation that genuinely needs the library.
_CODEMODE_ERROR: Exception | None = None
try:
    from ida_codemode.client import (
        ClientError,
        DatabaseHandle,
        InstanceDisconnectedError,
        RemoteError,
    )
    from ida_codemode.registry import (
        REGISTRY_DIR,
        FileLock,
        RegistryEntry,
        canonical_path,
        idb_key,
        scan_instances,
    )
    from ida_codemode.resolver import IdbBusy, expected_idb_path
except ImportError as _exc:  # library absent: usable only for offline layers
    _CODEMODE_ERROR = _exc
    # Bound to None rather than left undefined so the names stay patchable: the
    # offline contract tests inject a fake DatabaseHandle here.
    ClientError = InstanceDisconnectedError = RemoteError = None  # type: ignore[assignment,misc]
    DatabaseHandle = RegistryEntry = FileLock = None  # type: ignore[assignment,misc]
    REGISTRY_DIR = canonical_path = idb_key = scan_instances = None  # type: ignore[assignment]
    IdbBusy = expected_idb_path = None  # type: ignore[assignment]


def _require_codemode() -> None:
    """Raise an actionable error when the Code Mode library is missing.

    Gated on the binding, not on the original import result, so a test that
    injects a fake ``DatabaseHandle`` exercises the real adapter logic.
    """
    if DatabaseHandle is None:
        raise IDAConnectionError(
            "ida-codemode-mcp is not installed in this environment "
            f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or "
            "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a "
            "database.") from _CODEMODE_ERROR


def database_owner(idb_path: str, staged_path: str | None = None):
    """The registry entry that owns ``idb_path``/``staged_path``, else None.

    Returns None when the Code Mode library is absent: with no library there is
    no client in this environment that could be holding the database, and the
    IDA-free layers (project staging) must keep working. Registry errors that
    happen WITH the library installed still propagate -- those mean "we could
    not determine ownership", which is not the same as "nobody owns it".
    """
    if DatabaseHandle is None:
        return None
    expected_key = idb_key(idb_path)
    staged = canonical_path(staged_path) if staged_path else None
    for item in scan_instances(timeout=0.5):
        entry = item.entry
        if entry.idb_key == expected_key:
            return entry
        if staged and entry.exe_path and canonical_path(entry.exe_path) == staged:
            return entry
    return None


def registered_database(path: str, output_database: str | None = None) -> bool:
    """Whether a live/lock-held Code Mode instance owns this target."""
    _require_codemode()
    source = canonical_path(path)
    expected = canonical_path(output_database) if output_database else expected_idb_path(source)
    expected_key = idb_key(expected)
    for instance in scan_instances(timeout=0.5):
        entry = instance.entry
        if entry.idb_key == expected_key:
            return True
        if not output_database and entry.backend == "gui" and entry.exe_path:
            if canonical_path(entry.exe_path) == source:
                return True
    return False


class _NoopKeepAlive:
    """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat."""

    def __init__(self) -> None:
        self.beats = self.failures = 0

    def start(self) -> "_NoopKeepAlive":
        return self

    def stop(self) -> None:
        pass


def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]:
    """Translate ida-tui's legacy first-open switches to Code Mode options.

    Code Mode has typed options for processor, natural loading address and file
    type.  It deliberately has no arbitrary command-line escape hatch; reject
    switches we cannot represent instead of silently loading a blob wrongly.
    """
    processor: str | None = None
    loading_address: int | None = None
    file_type: str | None = None
    unsupported: list[str] = []
    try:
        words = shlex.split(value or "", posix=os.name != "nt")
    except ValueError as exc:
        raise ValueError(f"invalid IDA load options: {exc}") from exc
    for word in words:
        if word.startswith("-p") and len(word) > 2:
            processor = word[2:]
        elif word.startswith("-b") and len(word) > 2:
            try:
                # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the
                # natural address, which is the safer public API.
                loading_address = int(word[2:], 16) << 4
            except ValueError as exc:
                raise ValueError(f"invalid IDA loading address: {word!r}") from exc
        elif word.startswith("-T") and len(word) > 2:
            file_type = word[2:]
        else:
            unsupported.append(word)
    if unsupported:
        joined = " ".join(unsupported)
        raise ValueError(
            "ida-codemode cannot represent arbitrary IDA load options: "
            f"{joined!r}; use processor/base/file type options instead"
        )
    return processor, loading_address, file_type


def _script(args: dict[str, Any], body: str) -> str:
    """Bind JSON arguments without interpolating user text into Python code."""
    encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":"))
    return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n"


# Rich flat-listing generation is the largest ida-domain gap in this port.
# ida-domain can enumerate heads and render plain disassembly, but it does not
# expose undefined runs, IDA colour spans, function banners, or expanded UDT
# members. Keep that IDAPython-only logic isolated in this one operation.
_HEADS = r'''
import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf
start = int(str(a["addr"]), 16)
count = max(1, min(int(a.get("count", 200)), 2000))
offset = max(0, int(a.get("offset", 0)))
annotate = bool(a.get("annotate", False))
seg = db.segments.get_at(start)
if seg is None:
    result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}}
else:
    lo, hi = int(seg.start_ea), int(seg.end_ea)
    if a.get("end"):
        hi = min(hi, int(str(a["end"]), 16))

    span_names = {
        "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"),
        "reg": ("SCOLOR_REG",),
        "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"),
        "str": ("SCOLOR_STRING",),
        "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME",
                 "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME",
                 "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"),
        "seg": ("SCOLOR_SEGNAME",),
        "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"),
        "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"),
        "err": ("SCOLOR_ERROR",),
    }
    tag_kinds = {}
    for kind, names in span_names.items():
        for name in names:
            value = getattr(ida_lines, name, None)
            if isinstance(value, str) and value:
                tag_kinds[value[0]] = kind
            elif isinstance(value, int):
                tag_kinds[chr(value)] = kind

    def spans(tagged):
        on, off, esc = "\x01", "\x02", "\x03"
        addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28))
        addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16))
        out, stack, buf = [], [], []
        def flush():
            if buf:
                out.append([stack[-1] if stack else "text", "".join(buf)])
                buf.clear()
        i = 0
        while i < len(tagged):
            ch = tagged[i]
            if ch == on and i + 1 < len(tagged):
                tag = tagged[i + 1]
                if tag == addr_tag:
                    i += 2 + addr_len
                    continue
                flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue
            if ch == off and i + 1 < len(tagged):
                flush()
                if stack: stack.pop()
                i += 2; continue
            if ch == esc and i + 1 < len(tagged):
                buf.append(tagged[i + 1]); i += 2; continue
            buf.append(ch); i += 1
        flush()
        collapsed, previous_space = [], False
        for kind, text in out:
            acc = []
            for ch in text:
                if ch.isspace():
                    if previous_space: continue
                    acc.append(" "); previous_space = True
                else:
                    acc.append(ch); previous_space = False
            if acc: collapsed.append([kind, "".join(acc)])
        if collapsed:
            collapsed[0][1] = collapsed[0][1].lstrip()
            collapsed[-1][1] = collapsed[-1][1].rstrip()
        return [[kind, text] for kind, text in collapsed if text]

    def row(ea):
        flags = ida_bytes.get_flags(ea)
        kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown")
        tagged = ida_lines.generate_disasm_line(ea, 0) or ""
        text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else ""
        item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text}
        if tagged:
            rich = spans(tagged)
            if " ".join("".join(x[1] for x in rich).split()) == text:
                item["spans"] = rich
        name = ida_name.get_ea_name(ea)
        if name: item["name"] = name
        return item

    def unknown_row(ea, size):
        if size <= 1: return row(ea)
        item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"}
        name = ida_name.get_ea_name(ea)
        if name: item["name"] = name
        return item

    def members(ea):
        tif = db.types.get_at(ea)
        if tif is None or not tif.is_udt(): return []
        answer = []
        for member in db.types.get_udt_members(tif):
            type_text = member.type.dstr() or ""
            text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "")
            answer.append({"ea": hex(ea + member.offset), "kind": "member",
                           "size": int(member.size), "text": text})
        return answer

    def is_unknown(ea):
        flags = ida_bytes.get_flags(ea)
        return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags))
    def run_end(ea):
        nxt = ida_bytes.next_head(ea, hi)
        return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi
    def advance(ea):
        if is_unknown(ea): return run_end(ea)
        nxt = ida_bytes.get_item_end(ea)
        return nxt if nxt > ea else ea + 1
    def rows_for(ea):
        if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)]
        fn = db.functions.get_at(ea) if annotate else None
        at_start = fn is not None and int(fn.start_ea) == ea
        answer = []
        if at_start:
            name = db.functions.get_name(fn) or f"sub_{ea:X}"
            answer += [
                {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""},
                {"ea": hex(ea), "kind": "sep", "size": 0,
                 "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15},
                {"ea": hex(ea), "kind": "funchdr", "size": 0,
                 "text": name + " proc", "name": name},
            ]
        item = row(ea)
        if at_start:
            item["name"] = None
        elif annotate and item["kind"] == "code" and item.get("name"):
            name = item["name"]
            answer.append({"ea": hex(ea), "kind": "label", "size": 0,
                           "text": name + ":", "name": name})
            item["name"] = None
        answer.append(item)
        if item["kind"] == "data": answer += members(ea)
        if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea):
            name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}"
            answer += [
                {"ea": hex(ea), "kind": "funchdr", "size": 0,
                 "text": name + " endp", "name": name},
                {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60},
            ]
        return answer

    ea = ida_bytes.get_item_head(start)
    if ea == ida_idaapi.BADADDR: ea = start
    for _ in range(offset):
        if ea >= hi: break
        ea = advance(ea)
    rows = []
    more = False
    while ea != ida_idaapi.BADADDR and ea < hi:
        if len(rows) >= count:
            more = True; break
        rows += rows_for(ea)
        ea = advance(ea)
    result = {"addr": a["addr"], "heads": rows,
              "cursor": {"next": hex(ea)} if more else {"done": True}}
result
'''


_DECOMP_MAP_HELPER = r'''
def line_map(cfunc):
    import ida_hexrays
    answer = []
    for sl in cfunc.get_pseudocode():
        tagged, eas, seen = sl.line, [], set()
        for x in range(len(tagged) + 1):
            head = ida_hexrays.ctree_item_t(); item = ida_hexrays.ctree_item_t(); tail = ida_hexrays.ctree_item_t()
            if not cfunc.get_line_item(tagged, x, False, head, item, tail): continue
            text = item.dstr() or ""
            try: ea = int(text.split(": ", 1)[0], 16)
            except (ValueError, IndexError): continue
            if ea not in seen: seen.add(ea); eas.append(ea)
        answer.append(eas)
    return answer
'''


_OPERATIONS: dict[str, str] = {
    "list_funcs": r'''
import fnmatch
queries = a.get("queries") or [{}]
q = queries[0]
offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500)))
pattern = str(q.get("filter") or "").lower()
if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*"
rows = []
for fn in db.functions.get_all():
    name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}"
    if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue
    rows.append({"addr": hex(int(fn.start_ea)), "name": name,
                 "size": int(fn.end_ea) - int(fn.start_ea)})
page = rows[offset:offset + count]
result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]}
result
''',
    "disasm": r'''
ea = int(str(a["addr"]), 16)
fn = db.functions.get_at(ea)
if fn is None:
    result = {"instructions": [], "total_instructions": 0, "instruction_count": 0}
else:
    instructions = list(db.functions.get_instructions(fn))
    limit = max(1, int(a.get("max_instructions", len(instructions) or 1)))
    rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)}
            for insn in instructions[:limit]]
    result = {"instructions": rows, "total_instructions": len(instructions),
              "instruction_count": len(instructions)}
result
''',
    "file_regions": r'''
import idaapi
rows = []
for seg in db.segments.get_all():
    try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea))
    except Exception: file_off = -1
    if file_off < 0 or file_off >= (1 << 48): file_off = -1
    rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)),
                 "file_off": file_off, "name": db.segments.get_name(seg) or ""})
result = {"regions": rows}
result
''',
    "read_raw": r'''
import ida_bytes
ea, size = int(str(a["addr"]), 16), max(0, int(a["size"]))
raw = ida_bytes.get_bytes(ea, size) or b""
raw = raw[:size] + b"\xff" * max(0, size - len(raw))
data = bytearray(raw)
for index, value in enumerate(data):
    if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0
result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)}
result
''',
    "get_bytes": r'''
rows = []
for region in a.get("regions", []):
    ea, size = int(str(region["addr"]), 16), int(region["size"])
    raw = db.bytes.get_bytes_at(ea, size) or b""
    rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)})
result = {"result": rows}
result
''',
    "search_structs": r'''
needle = str(a.get("filter") or "").lower()
rows = []
for tif in db.types.get_all():
    name = tif.get_type_name() or ""
    if not name or needle not in name.lower() or not tif.is_udt(): continue
    members = list(db.types.get_udt_members(tif))
    rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()),
                 "cardinality": len(members), "ordinal": int(tif.get_ordinal())})
result = {"result": rows}
result
''',
    "type_inspect": r'''
rows = []
for query in a.get("queries", []):
    name = str(query.get("name") or "")
    tif = db.types.get_by_name(name)
    if tif is None:
        rows.append({"name": name, "error": "type not found"}); continue
    members = [{"name": m.name, "type": m.type.dstr() or str(m.type),
                "offset": int(m.offset), "size": int(m.size)}
               for m in db.types.get_udt_members(tif)] if tif.is_udt() else []
    rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()),
                 "members": members})
result = {"result": rows}
result
''',
    "declare_type": r'''
import ida_typeinf
decls = a.get("decls", "")
if isinstance(decls, str): decls = [decls]
rows = []
for declaration in decls:
    try:
        errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration))
        rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})})
    except Exception as exc:
        rows.append({"ok": False, "error": str(exc)})
result = {"result": rows}
result
''',
    "del_type": r'''
import ida_typeinf
name = str(a["name"])
ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE))
result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})}
result
''',
    "func_types": r'''
import ida_typeinf
ea = int(str(a["addr"]), 16)
fn = db.functions.get_at(ea)
if fn is None:
    result = {"addr": a["addr"], "error": "no function at address"}
else:
    pseudo = db.pseudocode.decompile(fn)
    name = db.functions.get_name(fn) or ""
    tif = pseudo.get_func_type()
    try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else ""
    except Exception: prototype = tif.dstr() if tif else ""
    lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "",
              "is_arg": bool(var.is_arg)} for var in pseudo.local_variables]
    result = {"addr": hex(int(fn.start_ea)), "name": name,
              "prototype": (prototype or "").strip(), "lvars": lvars}
result
''',
    "set_lvar_type": r'''
import ida_typeinf
ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"])
fn = db.functions.get_at(ea)
if fn is None:
    result = {"error": "no function at address"}
else:
    pseudo = db.pseudocode.decompile(fn)
    var = pseudo.find_local_variable(variable)
    if var is None:
        result = {"error": f"local variable {variable!r} not found"}
    else:
        try:
            tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration)
            accepted = bool(var.set_type(tif))
            saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False
            result = {"addr": hex(int(fn.start_ea)), "variable": variable,
                      "type": declaration, "ok": accepted and saved}
        except Exception as exc:
            result = {"error": f"bad type {declaration!r}: {exc}"}
result
''',
    "set_type": r'''
from ida_domain.types import TypeApplyFlags
rows = []
for edit in a.get("edits", []):
    ea = int(str(edit["addr"]), 16)
    declaration = str(edit.get("signature") or edit.get("type") or "")
    try:
        ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE))
        rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})})
    except Exception as exc:
        rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
result = {"result": rows}
result
''',
    "data_type": r'''
ea = int(str(a["addr"]), 16)
try:
    tif = db.types.get_at(ea)
    fn = db.functions.get_at(ea)
    result = {"addr": hex(ea), "name": db.names.get_at(ea) or "",
              "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0,
              "is_func": bool(fn)}
except Exception as exc:
    result = {"addr": hex(ea), "error": str(exc)}
result
''',
    "force_recompile": r'''
import ida_hexrays
rows = []
for item in a.get("items", []):
    ea = int(str(item["addr"]), 16)
    ida_hexrays.mark_cfunc_dirty(ea, False)
    rows.append({"addr": hex(ea), "ok": True})
result = {"result": rows}
result
''',
    "undefine": r'''
import ida_bytes
rows = []
for item in a.get("items", []):
    ea = int(str(item["addr"]), 16)
    size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1))
    ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size))
    rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})})
result = {"result": rows}
result
''',
    "define_code": r'''
import ida_ua
rows = []
for item in a.get("items", []):
    ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea))
    rows.append({"addr": hex(ea), "ok": size > 0, "size": size,
                 **({} if size > 0 else {"error": "instruction did not decode"})})
result = {"result": rows}
result
''',
    "define_func": r'''
rows = []
for item in a.get("items", []):
    ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea))
    rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})})
result = {"result": rows}
result
''',
    "make_data": r'''
import ida_bytes, ida_idaapi, ida_typeinf
from ida_domain.types import TypeApplyFlags
rows = []
for item in a.get("items", []):
    ea, declaration = int(str(item["addr"]), 16), str(item["type"])
    try:
        tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration)
        size = max(1, int(tif.get_size()))
        saved_names = [(addr, name) for addr, name in db.names.get_all()
                       if ea <= int(addr) < ea + size]
        ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES,
                            max(size, int(ida_bytes.get_item_size(ea) or 1)))
        created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR))
        ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE))
        for address, name in saved_names:
            db.names.set_name(int(address), name)
        if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"])))
        rows.append({"addr": hex(ea), "ok": ok, "size": size,
                     **({} if ok else {"error": "IDA rejected the data type"})})
    except Exception as exc:
        rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
result = {"result": rows}
result
''',
    "make_string": r'''
from ida_domain.strings import StringType
ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0)))
kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32,
        "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C)
import ida_bytes
try:
    ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1)
except Exception:
    pass
try:
    ok = bool(db.bytes.create_string_at(ea, length or None, kind))
    text = db.bytes.get_string_at(ea) or "" if ok else ""
    result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text}
except Exception as exc:
    result = {"addr": hex(ea), "ok": False, "error": str(exc)}
result
''',
    "list_strings": r'''
from ida_domain.strings import StringListConfig
offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4)))
if offset == 0 or a.get("refresh"):
    from ida_domain.strings import StringType
    db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len,
                                        only_ascii_7bit=False))
items = list(db.strings.get_all())
page = items[offset:offset + count]
rows = []
for item in page:
    try: text = str(item)
    except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else ""
    rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name})
result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)}
result
''',
    "list_linkage": r'''
imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name}
           for item in db.imports.get_all_imports() if item.name]
exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)}
           for item in db.entries.get_all() if item.name]
result = {"imports": imports, "exports": exports,
          "n_imports": len(imports), "n_exports": len(exports)}
result
''',
    "lookup_funcs": r'''
rows = []
for query in a.get("queries", []):
    raw = str(query)
    try: ea = int(raw, 16)
    except ValueError:
        fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None
    else: fn = db.functions.get_at(ea)
    if fn is None:
        rows.append({"query": raw, "fn": None})
    else:
        rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)),
                     "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}",
                     "size": int(fn.end_ea) - int(fn.start_ea)}})
result = {"result": rows}
result
''',
    "resolve_names": r'''
import ida_idaapi, ida_name
rows = []
for query in a.get("queries", []):
    name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name)
    rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None})
result = {"result": rows}
result
''',
    "xref_types": r'''
queries = a.get("queries") or []
all_results = []
for query in queries:
    ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both"))
    refs = []
    if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea))
    if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea))
    rows, seen = [], set()
    for ref in refs:
        key = (int(ref.from_ea), int(ref.to_ea), int(ref.type))
        if query.get("dedup") and key in seen: continue
        seen.add(key)
        fn = db.functions.get_at(int(ref.from_ea))
        kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow
                else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower())
        row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)),
               "type": "code" if ref.is_code else "data", "kind": kind}
        if query.get("include_fn") and fn is not None:
            row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""}
        rows.append(row)
        if len(rows) >= int(query.get("count", 2000)): break
    all_results.append({"data": rows})
result = {"result": all_results}
result
''',
    "xref_query": r'''
queries = a.get("queries") or []
all_results = []
for query in queries:
    ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both"))
    refs = []
    if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea))
    if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea))
    rows = []
    for ref in refs[:int(query.get("count", 2000))]:
        fn = db.functions.get_at(int(ref.from_ea))
        row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)),
               "type": "code" if ref.is_code else "data"}
        if query.get("include_fn") and fn is not None:
            row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""}
        rows.append(row)
    all_results.append({"data": rows})
result = {"result": all_results}
result
''',
    "set_comments": r'''
rows = []
for item in a.get("items", []):
    ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "")
    try:
        if text: ok = bool(db.comments.set_at(ea, text))
        else: db.comments.delete_at(ea); ok = True
        rows.append({"addr": hex(ea), "ok": ok})
    except Exception as exc:
        rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
result = {"result": rows}
result
''',
    "rename": r'''
import ida_idaapi, ida_name, ida_typeinf
batch = a.get("batch") or {}
out = {}; ok_count = failed = 0
for category, edit in batch.items():
    try:
        if category == "func":
            ea, new = int(str(edit["addr"]), 16), str(edit["name"])
            fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new))
        elif category == "data":
            new = str(edit.get("new") or "")
            if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16)
            else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or "")))
            ok = bool(db.names.set_name(ea, new))
        elif category in ("local", "stack"):
            ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"])
            pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old)
            if var is None: ok = False
            else:
                var.set_user_name(new)
                ok = bool(pseudo.save_local_variable_info(var, save_name=True))
        else:
            raise ValueError(f"unsupported rename category: {category}")
        row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})}
    except Exception as exc:
        row = {"ok": False, "error": str(exc)}
    out[category] = [row]
    if row["ok"]: ok_count += 1
    else: failed += 1
out["summary"] = {"ok": ok_count, "failed": failed}
result = out
result
''',
}


_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r'''
ea = int(str(a["addr"]), 16)
fn = db.functions.get_at(ea)
if fn is None:
    result = {"error": f"no function at {ea:#x}"}
else:
    pseudo = db.pseudocode.decompile(fn)
    mapping = line_map(pseudo.raw_cfunc)
    plain = pseudo.to_text()
    marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "")
              for line, eas in zip(plain, mapping)]
    import ida_name
    refs, seen = [], set()
    for expr in pseudo.find_objects():
        target = int(expr.obj_ea)
        if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue
        seen.add(target)
        name = expr.obj_name or ida_name.get_name(target) or ""
        try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None
        except Exception: string = None
        refs.append({"addr": hex(target), "name": name, "string": string})
    result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs}
result
'''

_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r'''
ea = int(str(a["addr"]), 16)
fn = db.functions.get_at(ea)
if fn is None:
    result = {"error": f"no function at {ea:#x}"}
else:
    pseudo = db.pseudocode.decompile(fn)
    mapping = line_map(pseudo.raw_cfunc)
    result = {"addr": hex(int(fn.start_ea)),
              "lines": [{"ea": hex(eas[0]) if eas else None,
                         "eas": [hex(item) for item in eas]} for eas in mapping]}
result
'''

_OPERATIONS["define_code_run"] = r'''
import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi
ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000))
seg = ida_segment.getseg(ea)
if seg is None:
    result = {"addr": a["addr"], "error": "no segment", "count": 0}
else:
    start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea)
    while count < limit:
        if ea >= hi: stopped = "segment"; break
        flags = ida_bytes.get_flags(ea)
        if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break
        size = int(ida_ua.create_insn(ea))
        if size <= 0: stopped = "undecodable"; break
        count += 1
        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, ea) > 0:
            try: is_ret = bool(ida_idp.is_ret_insn(insn))
            except Exception: is_ret = False
            if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP):
                ea += size; stopped = "flow"; break
        ea += size
    result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped}
result
'''

_OPERATIONS["define_func_run"] = r'''
import ida_bytes, ida_funcs, ida_segment
ea = int(str(a["addr"]), 16)
fn = db.functions.get_at(ea)
if fn is not None and int(fn.start_ea) == ea:
    result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"}
else:
    automatic = bool(db.functions.create(ea))
    if not automatic:
        seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea
        while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)):
            nxt = int(ida_bytes.get_item_end(end))
            if nxt <= end: break
            end = nxt
        ok = bool(end > ea and ida_funcs.add_func(ea, end))
    else: ok = True
    fn = db.functions.get_at(ea)
    result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)),
               "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"}
              if ok and fn is not None else
              {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"})
result
'''

_OPERATIONS["set_thumb"] = r'''
import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs
ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T")
seg = ida_segment.getseg(ea)
if treg is None or treg < 0:
    result = {"addr": hex(ea), "error": "no T register (not an ARM database)"}
elif seg is None:
    result = {"addr": hex(ea), "error": "no segment"}
else:
    current = ida_segregs.get_sreg(ea, treg)
    current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current)
    want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1)
    changed = False
    if want and seg.bitness != 1:
        ida_segment.set_segm_addressing(seg, 1); changed = True
    size = max(int(ida_bytes.get_item_size(ea)), 2)
    ida_bytes.del_items(ea, 0, size)
    ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user))
    now = ida_segregs.get_sreg(ea, treg)
    result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok,
              "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed,
              "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)}
result
'''

_OPERATIONS["thumb_scan"] = r'''
import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua
lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16)
apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512))
treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo
while cursor + 4 <= hi and len(found) < limit:
    at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4
    if not value & 1: continue
    target = value & ~1; seg = ida_segment.getseg(target)
    if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue
    flags = ida_bytes.get_flags(target)
    if ida_bytes.is_data(flags): continue
    item = {"at": hex(at), "value": hex(value), "target": hex(target),
            "was_code": bool(ida_bytes.is_code(flags))}; found.append(item)
    if not apply: continue
    if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user)
    if not ida_bytes.is_code(ida_bytes.get_flags(target)):
        ida_bytes.del_items(target, 0, 2)
        if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue
    item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1
result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)}
result
'''

_OPERATIONS["decomp_error"] = r'''
import ida_hexrays, ida_ida
ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea)
result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()}
if fn is None:
    result["reason"] = "no function here"
else:
    try:
        failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure)
        if cfunc is not None: result["reason"] = ""
        else:
            result.update({"reason": failure.desc() or f"error {failure.code}",
                           "code": int(failure.code), "errea": hex(int(failure.errea))})
    except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}"
result
'''


class CodeModeClient:
    """A leased GUI/idalib database accessed through ``ida_codemode``."""

    def __init__(
        self,
        binary_path: str,
        *,
        ttl: int = 0,
        load_args: str = "",
        processor: str | None = None,
        loading_address: int | None = None,
        file_type: str | None = None,
        output_database: str | None = None,
        spawn: bool = True,
        new_database: bool = False,
    ) -> None:
        del ttl  # managed-worker lifetime is lease-based, not idle-TTL based
        self._path = os.path.abspath(os.path.expanduser(binary_path))
        parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args)
        self._processor = processor or parsed_processor
        self._loading_address = loading_address if loading_address is not None else parsed_address
        self._file_type = file_type or parsed_file_type
        self._output_database = output_database
        self._spawn = spawn
        self._new_database = new_database
        self._handle: DatabaseHandle | None = None
        self._last_entry: RegistryEntry | None = None
        self._connect_lock = threading.Lock()

    def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient":
        _require_codemode()
        with self._connect_lock:
            if self._handle is not None and self._handle.connected:
                return self
            if progress:
                progress(f"discovering Code Mode database for {os.path.basename(self._path)}…")
            try:
                # A Ctrl+L reload releases its current managed-worker lease, but
                # that worker remains registered during Code Mode's final-lease
                # grace period. Retry only that known handoff window. A GUI or
                # another long-lived client remains busy and yields a clear
                # failure rather than being modified underneath its owner.
                deadline = time.monotonic() + min(timeout, 60.0)
                while True:
                    try:
                        handle = DatabaseHandle.open(
                            self._path,
                            spawn=self._spawn,
                            timeout=max(0.1, timeout),
                            output_database=self._output_database,
                            processor=self._processor,
                            # DatabaseHandle calls this image_base and wants the
                            # natural (16-byte aligned) address; it does the
                            # conversion to IDA's paragraph-based -b itself.
                            image_base=self._loading_address,
                            file_type=self._file_type,
                            new_database=self._new_database,
                        )
                        break
                    except IdbBusy:
                        if not self._new_database or time.monotonic() >= deadline:
                            raise
                        if progress:
                            progress("waiting for the previous Code Mode lease to close…")
                        # Remember the record before managed shutdown withdraws
                        # its JSON. The lifetime lock remains held until IDA has
                        # actually closed the IDB; waiting on it avoids racing a
                        # replacement worker into the old process's file lock.
                        expected = canonical_path(
                            self._output_database or expected_idb_path(self._path)
                        )
                        owners = [item.entry for item in scan_instances(timeout=0.5)
                                  if item.entry.idb_key == idb_key(expected)]
                        if owners:
                            self._wait_for_entry_release(
                                owners[0], max(0.0, deadline - time.monotonic())
                            )
                        else:
                            time.sleep(0.2)
                if progress:
                    backend = handle.entry.backend
                    progress(f"attached to {backend} database; waiting for auto-analysis…")
                handle.wait_autoanalysis(timeout=timeout)
            except Exception as exc:  # normalize the dependency's transport errors
                raise self._connection_error(exc) from exc
            self._handle = handle
            self._last_entry = handle.entry
            return self

    @staticmethod
    def _connection_error(exc: BaseException) -> IDAConnectionError:
        return IDAConnectionError(str(exc) or type(exc).__name__)

    @property
    def connected(self) -> bool:
        return self._handle is not None and self._handle.connected

    @property
    def pid(self) -> int | None:
        return self._handle.entry.pid if self._handle is not None else None

    @property
    def backend(self) -> str | None:
        return self._handle.entry.backend if self._handle is not None else None

    def execute_python(self, code: str, *, timeout: float | None = None) -> Any:
        if not self.connected:
            self.connect()
        handle = self._handle
        if handle is None:
            raise IDAConnectionError("Code Mode database is not connected")
        try:
            response = handle.execute_python(code, timeout=timeout)
        except RemoteError as exc:
            details = exc.details or {}
            message = str(exc)
            if details.get("traceback"):
                message += f"\n{details['traceback']}"
            if exc.code == "operation_timeout":
                raise IDATimeoutError(message) from exc
            raise IDAToolError("execute_python", message) from exc
        except (InstanceDisconnectedError, ClientError) as exc:
            raise self._connection_error(exc) from exc
        if not isinstance(response, dict) or "result" not in response:
            raise IDAToolError("execute_python", "Code Mode returned an invalid execution result")
        return response["result"]

    def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any:
        """Execute one TUI domain operation through Code Mode."""
        if operation in ("idb_save", "save"):
            return self.save_database()
        if operation in ("server_health", "ping", "health", "state"):
            return self.health()
        body = _HEADS if operation == "heads" else _OPERATIONS.get(operation)
        if body is None:
            raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}")
        try:
            return self.execute_python(_script(args, body), timeout=timeout)
        except IDAToolError as exc:
            if exc.tool == "execute_python":
                raise IDAToolError(operation, exc.message) from exc
            raise

    # Temporary source compatibility for external drivers/tests that used the
    # old WorkerClient. Application code uses the accurately named invoke().
    call = invoke

    def save_database(self) -> dict[str, Any]:
        if not self.connected:
            self.connect()
        handle = self._handle
        if handle is None:
            raise IDAConnectionError("Code Mode database is not connected")
        try:
            return handle.save_database()
        except RemoteError as exc:
            raise IDAToolError("save_database", str(exc)) from exc
        except (InstanceDisconnectedError, ClientError) as exc:
            raise self._connection_error(exc) from exc

    def health(self) -> dict[str, Any]:
        if not self.connected:
            self.connect()
        assert self._handle is not None
        entry = self._handle.entry
        module = os.path.basename(entry.exe_path or entry.idb_path or self._path)
        return {
            "ok": self._handle.connected,
            "module": module,
            "backend": entry.backend,
            "record_id": entry.record_id,
            "input_path": entry.exe_path,
            "idb_path": entry.idb_path,
        }

    def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive:
        del interval
        return _NoopKeepAlive()

    def resolve_db(self) -> str:
        if not self.connected:
            self.connect()
        assert self._handle is not None
        return self._handle.entry.record_id

    def set_db(self, db: str | None) -> None:
        del db  # one handle is permanently bound to one registered database

    def list_sessions(self) -> list[Session]:
        if not self.connected:
            self.connect()
        assert self._handle is not None
        entry = self._handle.entry
        path = entry.exe_path or entry.idb_path or self._path
        return [Session(session_id=entry.record_id, filename=os.path.basename(path),
                        input_path=path, is_active=True)]

    def close(self, grace: float = 0.0) -> None:
        del grace
        with self._connect_lock:
            handle, self._handle = self._handle, None
        if handle is not None:
            self._last_entry = handle.entry
            handle.close()  # release our lease; never close a GUI/other client's DB

    @staticmethod
    def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool:
        _require_codemode()
        path = REGISTRY_DIR / f"{entry.record_id}.lock"
        deadline = time.monotonic() + max(0.0, timeout)
        while True:
            lock = FileLock(path)
            try:
                if lock.try_acquire():
                    return True
            except OSError:
                pass
            finally:
                lock.close()
            if time.monotonic() >= deadline:
                return False
            time.sleep(min(0.1, deadline - time.monotonic()))

    def wait_released(self, timeout: float = 45.0) -> bool:
        """Wait until a managed instance releases its lifetime lock.

        Normal application shutdown must not wait: another client may retain the
        worker. This is an explicit test/maintenance helper for deleting a
        temporary IDB safely after this client closes. GUI instances return
        ``False`` immediately because clients never own their lifetime.
        """
        entry = self._last_entry
        if entry is None or entry.backend != "idalib":
            return False
        return self._wait_for_entry_release(entry, timeout)

    def __enter__(self) -> "CodeModeClient":
        return self.connect()

    def __exit__(self, *exc) -> None:
        self.close()