ida-code 0.2.1__py3-none-any.whl

ida_code/__init__.py

@@ -0,0 +1,2 @@
+# Import session first to ensure idapro is loaded before any other ida_* imports.
+from ida_code import session  # noqa: F401

ida_code/_search_utils.py

@@ -0,0 +1,33 @@
+"""Shared search utilities for doc_search and example_search.
+
+Provides word-boundary-aware matching: "set" matches "set_name" and
+"ida_name.set_name" but not "reset" or "offset".
+
+Boundaries are: start-of-string, underscore, dot, whitespace.
+"""
+
+import re
+from functools import lru_cache
+
+
+@lru_cache(maxsize=128)
+def _boundary_pattern(term: str) -> re.Pattern:
+    """Regex that matches term at an underscore/dot/whitespace boundary."""
+    escaped = re.escape(term)
+    return re.compile(rf"(?:^|[_.\s]){escaped}", re.IGNORECASE)
+
+
+def term_matches(term: str, text: str) -> bool:
+    """Check if term appears in text at a word boundary.
+
+    Boundaries are: start-of-string, underscore, dot, whitespace.
+    Fast path: rejects via substring check before running regex.
+
+    Dotted terms (e.g. "ida_funcs.get_func") use plain substring
+    matching since the dot is already specific enough.
+    """
+    if term not in text.lower():
+        return False
+    if "." in term:
+        return True  # dotted terms: substring match is sufficient
+    return _boundary_pattern(term).search(text) is not None

ida_code/comments.py

@@ -0,0 +1,191 @@
+"""Comment management (regular, repeatable, function, anterior, posterior)."""
+
+import logging
+
+from fastmcp.exceptions import ToolError
+
+from ida_code import session
+
+log = logging.getLogger(__name__)
+
+_COMMENT_TYPES = {"regular", "repeatable", "function", "anterior", "posterior"}
+
+
+def _validate_comment_type(comment_type: str, allow_empty: bool = False) -> None:
+    """Raise ToolError if comment_type is not recognized."""
+    if allow_empty and comment_type == "":
+        return
+    if comment_type not in _COMMENT_TYPES:
+        allowed = ", ".join(sorted(_COMMENT_TYPES))
+        if allow_empty:
+            allowed += ', or "" for all types'
+        raise ToolError(
+            f"Invalid comment_type '{comment_type}'. Must be one of: {allowed}"
+        )
+
+
+def _get_func(ea: int):
+    """Resolve address to func_t, raise ToolError if not in a function."""
+    import ida_funcs
+
+    pfn = ida_funcs.get_func(ea)
+    if pfn is None:
+        raise ToolError(f"Address {ea:#x} is not within a recognized function.")
+    return pfn
+
+
+def _get_anterior(ea: int) -> str:
+    """Collect all anterior extra comment lines at ea."""
+    import ida_lines
+
+    lines = []
+    idx = 0
+    while True:
+        line = ida_lines.get_extra_cmt(ea, ida_lines.E_PREV + idx)
+        if line is None:
+            break
+        lines.append(line)
+        idx += 1
+    return "\n".join(lines)
+
+
+def _get_posterior(ea: int) -> str:
+    """Collect all posterior extra comment lines at ea."""
+    import ida_lines
+
+    lines = []
+    idx = 0
+    while True:
+        line = ida_lines.get_extra_cmt(ea, ida_lines.E_NEXT + idx)
+        if line is None:
+            break
+        lines.append(line)
+        idx += 1
+    return "\n".join(lines)
+
+
+def get_comment(ea: int, comment_type: str = "") -> dict:
+    """Get comment(s) at an address.
+
+    When *comment_type* is empty, returns all non-empty comment types.
+    When a specific type is given, returns just that type.
+    """
+    session.require_open()
+    _validate_comment_type(comment_type, allow_empty=True)
+
+    import idc
+    import ida_funcs
+
+    if comment_type == "":
+        result: dict = {"address": f"{ea:#x}"}
+
+        regular = idc.get_cmt(ea, 0) or ""
+        if regular:
+            result["regular"] = regular
+
+        repeatable = idc.get_cmt(ea, 1) or ""
+        if repeatable:
+            result["repeatable"] = repeatable
+
+        pfn = ida_funcs.get_func(ea)
+        if pfn is not None:
+            func_cmt = ida_funcs.get_func_cmt(pfn, 0) or ""
+            if func_cmt:
+                result["function"] = func_cmt
+
+        anterior = _get_anterior(ea)
+        if anterior:
+            result["anterior"] = anterior
+
+        posterior = _get_posterior(ea)
+        if posterior:
+            result["posterior"] = posterior
+
+        return result
+
+    # Specific type requested.
+    if comment_type == "regular":
+        comment = idc.get_cmt(ea, 0) or ""
+    elif comment_type == "repeatable":
+        comment = idc.get_cmt(ea, 1) or ""
+    elif comment_type == "function":
+        pfn = _get_func(ea)
+        comment = ida_funcs.get_func_cmt(pfn, 0) or ""
+    elif comment_type == "anterior":
+        comment = _get_anterior(ea)
+    else:  # posterior
+        comment = _get_posterior(ea)
+
+    return {"address": f"{ea:#x}", "comment_type": comment_type, "comment": comment}
+
+
+def set_comment(ea: int, comment: str, comment_type: str = "regular") -> dict:
+    """Set a comment at an address."""
+    session.require_open()
+    _validate_comment_type(comment_type)
+
+    import idc
+    import ida_funcs
+    import ida_lines
+
+    if comment_type == "regular":
+        idc.set_cmt(ea, comment, 0)
+    elif comment_type == "repeatable":
+        idc.set_cmt(ea, comment, 1)
+    elif comment_type == "function":
+        pfn = _get_func(ea)
+        ida_funcs.set_func_cmt(pfn, comment, 0)
+    elif comment_type == "anterior":
+        lines = comment.split("\n")
+        for i, line in enumerate(lines):
+            ida_lines.update_extra_cmt(ea, ida_lines.E_PREV + i, line)
+    else:  # posterior
+        lines = comment.split("\n")
+        for i, line in enumerate(lines):
+            ida_lines.update_extra_cmt(ea, ida_lines.E_NEXT + i, line)
+
+    log.info("Set %s comment at %#x", comment_type, ea)
+    return {
+        "address": f"{ea:#x}",
+        "comment_type": comment_type,
+        "comment": comment,
+        "status": "updated",
+    }
+
+
+def delete_comment(ea: int, comment_type: str = "regular") -> dict:
+    """Delete a comment at an address."""
+    session.require_open()
+    _validate_comment_type(comment_type)
+
+    import idc
+    import ida_funcs
+    import ida_lines
+
+    if comment_type == "regular":
+        idc.set_cmt(ea, "", 0)
+    elif comment_type == "repeatable":
+        idc.set_cmt(ea, "", 1)
+    elif comment_type == "function":
+        pfn = _get_func(ea)
+        ida_funcs.set_func_cmt(pfn, "", 0)
+    elif comment_type == "anterior":
+        # Count existing lines first, then delete all.
+        count = 0
+        while ida_lines.get_extra_cmt(ea, ida_lines.E_PREV + count) is not None:
+            count += 1
+        for i in range(count):
+            ida_lines.del_extra_cmt(ea, ida_lines.E_PREV + i)
+    else:  # posterior
+        count = 0
+        while ida_lines.get_extra_cmt(ea, ida_lines.E_NEXT + count) is not None:
+            count += 1
+        for i in range(count):
+            ida_lines.del_extra_cmt(ea, ida_lines.E_NEXT + i)
+
+    log.info("Deleted %s comment at %#x", comment_type, ea)
+    return {
+        "address": f"{ea:#x}",
+        "comment_type": comment_type,
+        "status": "deleted",
+    }

ida_code/config.py

@@ -0,0 +1,9 @@
+import os
+from pathlib import Path
+
+IDA_INSTALL_DIR = Path(os.environ.get("IDA_INSTALL_DIR", "/opt/ida-pro-9.2"))
+IDA_DOCS_DIR = IDA_INSTALL_DIR / "docs"
+IDA_PYTHON_DIR = IDA_INSTALL_DIR / "python"
+IDA_EXAMPLES_DIR = IDA_PYTHON_DIR / "examples"
+LOG_LEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+MCP_AUTH_TOKEN = os.environ.get("MCP_AUTH_TOKEN", "")

ida_code/doc_search.py

@@ -0,0 +1,255 @@
+import json
+import logging
+import re
+from html.parser import HTMLParser
+from pathlib import Path
+
+from ida_code._search_utils import term_matches
+from ida_code.config import IDA_DOCS_DIR, IDA_PYTHON_DIR
+
+log = logging.getLogger(__name__)
+
+# Lazily-built indexes.
+_html_docs: list[tuple[str, str, str]] | None = None  # (title, clean_text, location)
+_py_chunks: list[tuple[str, str, str]] | None = None   # (name, body, source_file)
+
+
+class _HTMLStripper(HTMLParser):
+    """Strip HTML tags, keep text content."""
+
+    def __init__(self):
+        super().__init__()
+        self._parts: list[str] = []
+
+    def handle_data(self, data: str):
+        self._parts.append(data)
+
+    def get_text(self) -> str:
+        return "".join(self._parts)
+
+
+def _strip_html(html: str) -> str:
+    s = _HTMLStripper()
+    s.feed(html)
+    return s.get_text()
+
+
+def _load_html_docs() -> list[tuple[str, str, str]]:
+    index_path = IDA_DOCS_DIR / "search" / "search_index.json"
+    with open(index_path) as f:
+        data = json.load(f)
+
+    entries = []
+    for doc in data["docs"]:
+        title = doc.get("title", "")
+        text = _strip_html(doc.get("text", ""))
+        location = doc.get("location", "")
+        if title or text:
+            entries.append((title, text, location))
+    return entries
+
+
+def _load_py_chunks() -> list[tuple[str, str, str]]:
+    chunks = []
+    for py_file in sorted(IDA_PYTHON_DIR.glob("ida_*.py")):
+        _parse_py_file(py_file, chunks)
+    # Also include idautils.py and idc.py if present.
+    for name in ("idautils.py", "idc.py"):
+        p = IDA_PYTHON_DIR / name
+        if p.exists():
+            _parse_py_file(p, chunks)
+    return chunks
+
+
+def _parse_py_file(path: Path, chunks: list[tuple[str, str, str]]):
+    """Split a Python file into chunks at top-level def/class boundaries."""
+    source_name = path.name
+    try:
+        lines = path.read_text(errors="replace").splitlines()
+    except OSError:
+        return
+
+    # Find lines that start a new def or class at the top level (no indentation
+    # or class-level indentation for methods).
+    boundary_pattern = re.compile(r"^(def |class )")
+    boundaries: list[int] = []
+    for i, line in enumerate(lines):
+        if boundary_pattern.match(line):
+            boundaries.append(i)
+
+    if not boundaries:
+        return
+
+    for idx, start in enumerate(boundaries):
+        end = boundaries[idx + 1] if idx + 1 < len(boundaries) else len(lines)
+        block = lines[start:end]
+
+        # Extract name from the first line.
+        first_line = block[0]
+        m = re.match(r"(?:def|class)\s+(\w+)", first_line)
+        if not m:
+            continue
+        name = m.group(1)
+
+        # Skip SWIG internals.
+        if name.startswith("_swig"):
+            continue
+
+        # Keep a reasonable amount of context: signature + docstring + a few lines.
+        body = "\n".join(block[:40])
+        chunks.append((name, body, source_name))
+
+
+def _ensure_indexes():
+    global _html_docs, _py_chunks
+    if _html_docs is None:
+        log.info("Loading HTML docs index from %s", IDA_DOCS_DIR)
+        _html_docs = _load_html_docs()
+        log.info("Loaded %d HTML doc entries", len(_html_docs))
+    if _py_chunks is None:
+        log.info("Loading Python API chunks from %s", IDA_PYTHON_DIR)
+        _py_chunks = _load_py_chunks()
+        log.info("Loaded %d Python API chunks", len(_py_chunks))
+
+
+def search(
+    query: str,
+    max_results: int = 5,
+    max_snippet_length: int = 150,
+    include_examples: bool = True,
+) -> dict:
+    """Search IDA docs and Python API sources. Returns structured dict."""
+    _ensure_indexes()
+
+    terms = query.lower().split()
+    if not terms:
+        return {"query": query, "results": []}
+    log.debug("Searching for terms: %s", terms)
+
+    results: list[tuple[float, str, str, str]] = []  # (score, title, snippet, source)
+
+    # Search HTML docs.
+    for title, text, location in _html_docs:
+        score = _score(terms, title, text)
+        if score > 0:
+            snippet = _excerpt(text, terms, max_len=max_snippet_length)
+            results.append((score, title, snippet, f"docs: {location}"))
+
+    # Search Python API chunks (name field gets higher weight).
+    for name, body, source_file in _py_chunks:
+        score = _score_py(terms, name, body)
+        if score > 0:
+            snippet = _excerpt(body, terms, max_len=max_snippet_length)
+            results.append((score, name, snippet, f"python: {source_file}"))
+
+    # Sort by score descending.
+    results.sort(key=lambda r: r[0], reverse=True)
+    results = results[:max_results]
+
+    result = {
+        "query": query,
+        "results": [
+            {"source": source, "title": title, "snippet": snippet, "score": score}
+            for score, title, snippet, source in results
+        ],
+    }
+
+    # Cross-link: append matching examples if available.
+    if include_examples:
+        from ida_code.example_search import search as _search_examples
+
+        ex_results = _search_examples(query, max_results=2, max_snippet_lines=5)
+        if ex_results["results"]:
+            result["related_examples"] = ex_results["results"]
+
+    return result
+
+
+def _score(terms: list[str], title: str, text: str) -> float:
+    """Score a document against search terms with field weighting.
+
+    Title matches score 4.0, body matches score 1.0.
+    All-terms-match bonus: 1.5x multiplier.
+    """
+    total = 0.0
+    matched_terms = 0
+
+    title_lower = title.lower()
+    text_lower = text.lower()
+
+    for term in terms:
+        term_score = 0.0
+
+        # Title match (high value)
+        if term_matches(term, title_lower):
+            term_score = max(term_score, 4.0)
+
+        # Body match (lower value)
+        if term_matches(term, text_lower):
+            term_score = max(term_score, 1.0)
+
+        if term_score > 0:
+            matched_terms += 1
+        total += term_score
+
+    # All-terms-match bonus
+    if len(terms) > 1 and matched_terms == len(terms):
+        total *= 1.5
+
+    return total
+
+
+def _score_py(terms: list[str], name: str, body: str) -> float:
+    """Score a Python API chunk with name-weighted scoring.
+
+    Name matches score 5.0, body matches score 1.0.
+    """
+    total = 0.0
+    matched_terms = 0
+
+    name_lower = name.lower()
+    body_lower = body.lower()
+
+    for term in terms:
+        term_score = 0.0
+
+        # Name match (highest value — this IS the API definition)
+        if term_matches(term, name_lower):
+            term_score = max(term_score, 5.0)
+
+        # Body match
+        if term_matches(term, body_lower):
+            term_score = max(term_score, 1.0)
+
+        if term_score > 0:
+            matched_terms += 1
+        total += term_score
+
+    if len(terms) > 1 and matched_terms == len(terms):
+        total *= 1.5
+
+    return total
+
+
+def _excerpt(text: str, terms: list[str], max_len: int = 300) -> str:
+    """Extract a snippet around the first matching term."""
+    text_lower = text.lower()
+    best_pos = len(text)
+    for t in terms:
+        # Use simple substring find for excerpt positioning
+        pos = text_lower.find(t)
+        if pos != -1 and pos < best_pos:
+            best_pos = pos
+
+    start = max(0, best_pos - 50)
+    end = start + max_len
+    snippet = text[start:end].strip()
+
+    if start > 0:
+        snippet = "..." + snippet
+    if end < len(text):
+        snippet = snippet + "..."
+
+    # Collapse whitespace.
+    snippet = re.sub(r"\s+", " ", snippet)
+    return snippet