ida-code 0.2.1__py3-none-any.whl

ida_code/session.py

@@ -0,0 +1,293 @@
+import atexit
+import enum
+import logging
+import os
+import sys
+
+from ida_code.config import IDA_INSTALL_DIR
+
+log = logging.getLogger(__name__)
+
+# Add idalib's Python package to sys.path so `idapro` can be imported
+# without requiring manual `pip install`.
+_idalib_python = IDA_INSTALL_DIR / "idalib" / "python"
+if _idalib_python.is_dir() and str(_idalib_python) not in sys.path:
+    sys.path.insert(0, str(_idalib_python))
+
+# Set IDADIR so idapro finds the IDA install directory
+# without requiring `py-activate-idalib.py`.
+os.environ.setdefault("IDADIR", str(IDA_INSTALL_DIR))
+
+try:
+    import idapro
+except ImportError as e:
+    raise ImportError(
+        f"Could not import idapro from {_idalib_python}. "
+        f"Set IDA_INSTALL_DIR to your IDA Pro 9.2+ installation directory "
+        f"(currently {IDA_INSTALL_DIR}). "
+        f"Original error: {e}"
+    ) from e
+
+
+class State(enum.Enum):
+    NO_DATABASE = "no_database"
+    DATABASE_OPEN = "database_open"
+
+
+_state = State.NO_DATABASE
+_db_path: str | None = None
+_db_file_path: str | None = None  # actual .i64/.idb path on disk
+_orphaned: bool = False  # True when database file vanished from disk
+
+
+def get_state() -> State:
+    return _state
+
+
+def _find_database_file(path: str) -> str | None:
+    """Find the .i64/.idb database file for *path*.
+
+    Checks both naming conventions IDA uses: replace-suffix and append.
+    Returns the path if found, else None.
+    """
+    from pathlib import Path
+    p = Path(path)
+    for ext in (".i64", ".idb"):
+        for candidate in (p.with_suffix(ext), Path(str(p) + ext)):
+            if candidate.is_file():
+                return str(candidate)
+    return None
+
+
+def require_open() -> None:
+    """Raise ``ToolError`` if no database is usable.
+
+    Checks both the state enum *and* that the database file still exists
+    on disk. If the file has been moved or deleted, resets state to
+    ``NO_DATABASE`` and raises a descriptive error instead of letting
+    idalib segfault.
+    """
+    from fastmcp.exceptions import ToolError
+
+    global _state, _db_file_path, _orphaned
+
+    if _state == State.NO_DATABASE:
+        raise ToolError("No database is open. Call open_database first.")
+
+    if _db_file_path is not None and not os.path.isfile(_db_file_path):
+        log.warning(
+            "Database file missing: %s — resetting state", _db_file_path
+        )
+        _state = State.NO_DATABASE
+        _orphaned = True
+        _db_file_path = None
+        from ida_code.executor import reset
+        reset()
+        raise ToolError(
+            "The database file has been moved or deleted. "
+            "Call open_database with the new path."
+        )
+
+
+def info() -> dict:
+    """Return a summary dict of the current database.
+
+    Raises ``ToolError`` if no database is open.
+    """
+    require_open()
+    return _collect_summary(_db_path or "<unknown>")
+
+
+def open(
+    path: str,
+    auto_analysis: bool = True,
+    overwrite: bool = False,
+    timeout: int = 0,
+    arch: str | None = None,
+) -> dict:
+    """Open a binary/database via idalib. Returns a summary dict.
+
+    *timeout* limits auto-analysis wait time in seconds (0 = unlimited).
+    When the timeout expires, the database remains open with partial analysis.
+
+    *arch* selects a specific architecture slice from a fat (universal) Mach-O
+    binary (e.g. "arm64e", "x86_64"). The slice is extracted to a temporary
+    thin file before opening. Ignored for non-fat binaries.
+
+    Raises ``ToolError`` on failure.
+    """
+    from fastmcp.exceptions import ToolError
+
+    from ida_code import macho
+
+    global _state, _db_path, _db_file_path, _orphaned
+
+    _orphaned = False
+
+    # Close any existing database first.
+    if _state == State.DATABASE_OPEN:
+        close()
+
+    # If an architecture is requested, extract the slice from a fat Mach-O.
+    open_path = path
+    if arch:
+        try:
+            open_path = macho.extract_slice(path, arch)
+        except ValueError:
+            # Not a fat binary or arch not found.
+            available = macho.list_architectures(path)
+            if available:
+                raise ToolError(
+                    f"Architecture '{arch}' not found. "
+                    f"Available: {available}"
+                )
+            log.warning(
+                "arch='%s' requested but %s is not a fat Mach-O; "
+                "opening as-is",
+                arch, path,
+            )
+
+    if overwrite:
+        _remove_existing_databases(open_path)
+
+    use_polling = auto_analysis and timeout > 0
+    run_auto = auto_analysis and not use_polling
+
+    log.info(
+        "Opening database: %s (auto_analysis=%s, overwrite=%s, timeout=%d, arch=%s)",
+        open_path, auto_analysis, overwrite, timeout, arch,
+    )
+    rc = idapro.open_database(open_path, run_auto)
+    if rc != 0:
+        log.error("open_database failed with code %d for %s", rc, open_path)
+        raise ToolError(f"open_database returned code {rc}")
+
+    _state = State.DATABASE_OPEN
+    _db_path = open_path
+    _db_file_path = _find_database_file(open_path)
+
+    timed_out = False
+    if use_polling:
+        timed_out = _wait_for_analysis(timeout)
+
+    log.info("Database opened successfully: %s", open_path)
+
+    # Reset executor namespace for the new database.
+    from ida_code.executor import reset
+    reset()
+
+    summary = _collect_summary(open_path)
+    if arch:
+        summary["arch"] = arch
+        summary["original_path"] = path
+    summary["warning"] = (
+        "Auto-analysis timed out \u2014 results may be incomplete."
+        if timed_out
+        else None
+    )
+    return summary
+
+
+def _wait_for_analysis(timeout: int) -> bool:
+    """Poll until auto-analysis finishes or *timeout* seconds elapse.
+
+    Returns True if the timeout expired (analysis incomplete).
+    """
+    import time
+    import ida_auto
+    import ida_funcs
+
+    deadline = time.monotonic() + timeout
+    interval = 0.5  # seconds between polls
+    last_count = 0
+
+    while not ida_auto.auto_is_ok():
+        now = time.monotonic()
+        if now >= deadline:
+            log.warning("Auto-analysis timed out after %ds", timeout)
+            return True
+
+        count = ida_funcs.get_func_qty()
+        if count != last_count:
+            log.info("Auto-analysis in progress: %d functions so far", count)
+            last_count = count
+
+        remaining = deadline - now
+        time.sleep(min(interval, remaining))
+
+    log.info("Auto-analysis completed")
+    return False
+
+
+def _collect_summary(path: str) -> dict:
+    import ida_ida
+    import ida_entry
+    import ida_segment
+    import idautils
+
+    processor = ida_ida.inf_get_procname()
+    bits = 64 if ida_ida.inf_is_64bit() else (32 if ida_ida.inf_is_32bit() else 16)
+
+    # Segments
+    segments = []
+    for seg_ea in idautils.Segments():
+        seg = ida_segment.getseg(seg_ea)
+        name = ida_segment.get_segm_name(seg)
+        segments.append({
+            "name": name,
+            "start": f"{seg.start_ea:#x}",
+            "end": f"{seg.end_ea:#x}",
+        })
+
+    # Entry points (capped to avoid huge responses for symbol-heavy binaries)
+    entry_count = ida_entry.get_entry_qty()
+    max_entries = 20
+    entry_points = []
+    for i in range(min(entry_count, max_entries)):
+        ordinal = ida_entry.get_entry_ordinal(i)
+        ea = ida_entry.get_entry(ordinal)
+        name = ida_entry.get_entry_name(ordinal)
+        entry_points.append({"name": name, "address": f"{ea:#x}"})
+
+    # Function count
+    func_count = sum(1 for _ in idautils.Functions())
+
+    return {
+        "path": path,
+        "processor": processor,
+        "bits": bits,
+        "function_count": func_count,
+        "segments": segments,
+        "entry_point_count": entry_count,
+        "entry_points": entry_points,
+    }
+
+
+def _remove_existing_databases(path: str) -> None:
+    """Remove existing IDA database files so a fresh analysis starts."""
+    from pathlib import Path
+    p = Path(path)
+    for ext in (".i64", ".idb"):
+        for candidate in {p.with_suffix(ext), Path(str(p) + ext)}:
+            if candidate.is_file():
+                candidate.unlink()
+
+
+def close() -> None:
+    """Close the current database."""
+    global _state, _db_path, _db_file_path, _orphaned
+    if _state == State.DATABASE_OPEN:
+        log.info("Closing database")
+        if not _orphaned:
+            idapro.close_database()
+        _state = State.NO_DATABASE
+        _db_path = None
+        _db_file_path = None
+        _orphaned = False
+        from ida_code.executor import reset
+        reset()
+
+
+@atexit.register
+def _cleanup():
+    close()

ida_code/snapshots.py

@@ -0,0 +1,110 @@
+import logging
+import os
+
+from fastmcp.exceptions import ToolError
+
+from ida_code import session
+
+log = logging.getLogger(__name__)
+
+
+def _build_snapshot_list():
+    """Return list of snapshot_t from the snapshot tree."""
+    import ida_loader
+
+    root = ida_loader.snapshot_t()
+    ida_loader.build_snapshot_tree(root)
+    return list(root.children)
+
+
+def _find_by_id(snapshots, id_str: str):
+    """Find a snapshot by its string ID. Raises ToolError if not found."""
+    for ss in snapshots:
+        if str(ss.id) == id_str:
+            return ss
+    raise ToolError(f"Snapshot with id '{id_str}' not found.")
+
+
+def _to_dict(ss) -> dict:
+    """Convert a snapshot_t to a plain dict."""
+    return {
+        "id": str(ss.id),
+        "desc": ss.desc,
+        "filename": ss.filename,
+    }
+
+
+def list_snapshots() -> dict:
+    """List all snapshots for the current database."""
+    session.require_open()
+    snapshots = _build_snapshot_list()
+    return {
+        "snapshots": [_to_dict(ss) for ss in snapshots],
+        "count": len(snapshots),
+    }
+
+
+def create_snapshot(desc: str = "") -> dict:
+    """Create a new database snapshot."""
+    session.require_open()
+
+    import ida_kernwin
+    import ida_loader
+
+    ss = ida_loader.snapshot_t()
+    ss.desc = desc[:ida_loader.MAX_DATABASE_DESCRIPTION]
+
+    ok, err = ida_kernwin.take_database_snapshot(ss)
+    if not ok:
+        raise ToolError(f"Failed to create snapshot: {err}")
+
+    log.info("Created snapshot id=%s desc=%r", ss.id, ss.desc)
+    return _to_dict(ss)
+
+
+def restore_snapshot(snapshot_id: str) -> dict:
+    """Restore the database to a previous snapshot."""
+    session.require_open()
+
+    import ida_kernwin
+
+    snapshots = _build_snapshot_list()
+    ss = _find_by_id(snapshots, snapshot_id)
+
+    # restore_database_snapshot takes (snapshot, callback, userdata).
+    # In idalib headless mode the restore is synchronous; the callback
+    # receives (error_msg_or_empty, userdata).
+    restore_err = []
+
+    def _cb(err_msg, ud):
+        if err_msg:
+            restore_err.append(err_msg)
+
+    ida_kernwin.restore_database_snapshot(ss, _cb, None)
+
+    if restore_err:
+        raise ToolError(f"Failed to restore snapshot: {restore_err[0]}")
+
+    # Database state changed — reset executor namespace.
+    from ida_code.executor import reset
+    reset()
+
+    log.info("Restored snapshot id=%s", snapshot_id)
+    return {"status": "restored", "id": snapshot_id}
+
+
+def remove_snapshot(snapshot_id: str) -> dict:
+    """Remove a snapshot by deleting its file from disk."""
+    session.require_open()
+
+    snapshots = _build_snapshot_list()
+    ss = _find_by_id(snapshots, snapshot_id)
+
+    filename = ss.filename
+    try:
+        os.remove(filename)
+    except OSError as e:
+        raise ToolError(f"Failed to remove snapshot file '{filename}': {e}")
+
+    log.info("Removed snapshot id=%s file=%s", snapshot_id, filename)
+    return {"status": "removed", "id": snapshot_id, "filename": filename}

ida_code/structures.py

@@ -0,0 +1,227 @@
+"""Structure (struct/union) management via IDA type info library."""
+
+import logging
+import re
+
+from fastmcp.exceptions import ToolError
+
+from ida_code import session
+
+log = logging.getLogger(__name__)
+
+_NAME_RE = re.compile(r"(?:struct|union)\s+(\w+)\s*\{")
+
+
+def _extract_name(c_code: str) -> str:
+    """Extract the struct/union name from a C definition string.
+
+    Raises ToolError if no name can be found.
+    """
+    m = _NAME_RE.search(c_code)
+    if not m:
+        raise ToolError(
+            "Could not extract struct/union name from definition. "
+            "Expected format: 'struct name { ... };' or 'union name { ... };'"
+        )
+    return m.group(1)
+
+
+def _get_struct_tinfo(name: str):
+    """Look up a struct/union by name, returning (tinfo, ordinal).
+
+    Raises ToolError if the type is not found.
+    """
+    import ida_typeinf
+
+    til = ida_typeinf.get_idati()
+    tif = ida_typeinf.tinfo_t()
+    if not tif.get_named_type(til, name):
+        raise ToolError(f"Structure '{name}' not found.")
+    ordinal = ida_typeinf.get_type_ordinal(til, name)
+    return tif, ordinal
+
+
+def _annotated_definition(tif, ordinal: int) -> str:
+    """Build a C definition with ``/* offset */`` comments on each member."""
+    import ida_typeinf
+    import idc
+
+    # Get base definition from IDA.
+    if ordinal <= 0:
+        return ""
+    raw = idc.print_decls(ordinal, 0) or ""
+    raw = raw.strip()
+    if not raw:
+        return ""
+
+    # Extract member offsets.
+    udt = ida_typeinf.udt_type_data_t()
+    if not tif.get_udt_details(udt):
+        return raw
+
+    offsets = {}  # member name -> offset in bytes
+    for i in range(udt.size()):
+        udm = udt.at(i)
+        offsets[udm.name] = udm.offset // 8
+
+    # Annotate each member line inside the braces.
+    lines = raw.split("\n")
+    out = []
+    for line in lines:
+        stripped = line.strip()
+        # Skip opening/closing braces and empty lines.
+        if stripped in ("{", "};", "") or stripped.startswith("struct ") or stripped.startswith("union "):
+            out.append(line)
+            continue
+        # Try to match a member name in this line.
+        for mname, off in offsets.items():
+            if mname in stripped:
+                line = f"{line}  /* {off:#x} */"
+                break
+        out.append(line)
+    return "\n".join(out)
+
+
+def _struct_to_dict(name: str) -> dict:
+    """Convert a named struct/union to a result dict with C definition."""
+    import ida_typeinf
+
+    tif, ordinal = _get_struct_tinfo(name)
+    size = tif.get_size()
+    is_union = tif.is_union()
+    member_count = tif.get_udt_nmembers()
+
+    # Get alignment from UDT details.
+    alignment = 0
+    udt = ida_typeinf.udt_type_data_t()
+    if tif.get_udt_details(udt):
+        alignment = udt.effalign
+
+    definition = _annotated_definition(tif, ordinal)
+
+    return {
+        "name": name,
+        "size": size,
+        "is_union": is_union,
+        "alignment": alignment,
+        "member_count": member_count,
+        "definition": definition,
+    }
+
+
+def list_structures(offset: int = 0, limit: int = 50, filter: str = "") -> dict:
+    """List structures in the database with pagination."""
+    session.require_open()
+
+    import ida_typeinf
+    import idautils
+
+    all_structs = list(idautils.Structs())
+    til = ida_typeinf.get_idati()
+    structures = []
+    skipped = 0
+
+    for ordinal, sid, name in all_structs:
+        if filter and filter.lower() not in name.lower():
+            continue
+        if skipped < offset:
+            skipped += 1
+            continue
+
+        tif = ida_typeinf.tinfo_t()
+        if tif.get_named_type(til, name):
+            size = tif.get_size()
+            member_count = tif.get_udt_nmembers()
+            udt = ida_typeinf.udt_type_data_t()
+            alignment = udt.effalign if tif.get_udt_details(udt) else 0
+        else:
+            size = 0
+            member_count = 0
+            alignment = 0
+
+        structures.append({
+            "name": name,
+            "size": size,
+            "alignment": alignment,
+            "member_count": member_count,
+        })
+        if len(structures) >= limit:
+            break
+
+    return {
+        "structures": structures,
+        "total": len(all_structs),
+        "showing": len(structures),
+        "offset": offset,
+        "filter": filter,
+    }
+
+
+def get_structure(name: str) -> dict:
+    """Get detailed info about a structure by name."""
+    session.require_open()
+    return _struct_to_dict(name)
+
+
+def create_structure(definition: str) -> dict:
+    """Create a new structure from a C definition string."""
+    session.require_open()
+
+    import ida_typeinf
+    import idc
+
+    name = _extract_name(definition)
+
+    # Check it doesn't already exist.
+    til = ida_typeinf.get_idati()
+    tif = ida_typeinf.tinfo_t()
+    if tif.get_named_type(til, name):
+        raise ToolError(f"Structure '{name}' already exists. Use edit_structure to modify it.")
+
+    result = idc.parse_decls(definition, idc.PT_SIL)
+    if result != 0:
+        raise ToolError(f"Failed to parse definition (error code {result}). Check C syntax.")
+
+    log.info("Created structure '%s'", name)
+    return _struct_to_dict(name)
+
+
+def edit_structure(definition: str) -> dict:
+    """Edit an existing structure by replacing its definition."""
+    session.require_open()
+
+    import ida_typeinf
+    import idc
+
+    name = _extract_name(definition)
+
+    # Check it exists.
+    til = ida_typeinf.get_idati()
+    tif = ida_typeinf.tinfo_t()
+    if not tif.get_named_type(til, name):
+        raise ToolError(f"Structure '{name}' not found. Use create_structure to create it.")
+
+    result = idc.parse_decls(definition, idc.PT_SIL | idc.PT_REPLACE)
+    if result != 0:
+        raise ToolError(f"Failed to parse definition (error code {result}). Check C syntax.")
+
+    log.info("Edited structure '%s'", name)
+    return _struct_to_dict(name)
+
+
+def delete_structure(name: str) -> dict:
+    """Delete a structure by name."""
+    session.require_open()
+
+    import ida_typeinf
+
+    # Check it exists.
+    til = ida_typeinf.get_idati()
+    tif = ida_typeinf.tinfo_t()
+    if not tif.get_named_type(til, name):
+        raise ToolError(f"Structure '{name}' not found.")
+
+    ida_typeinf.del_named_type(til, name, ida_typeinf.NTF_TYPE)
+
+    log.info("Deleted structure '%s'", name)
+    return {"status": "deleted", "name": name}