apisec-code-bolt 0.1.0__py3-none-any.whl

apisec_code_bolt/query/ast_cache.py

@@ -0,0 +1,182 @@
+"""
+Cached file reader and AST parser for query answering.
+
+Avoids re-reading and re-parsing the same files across multiple
+verification questions in a single analysis session. Uses content hashes
+to skip re-parsing when file content hasn't changed.
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class ASTCache:
+    """Cache parsed ASTs and source text for project files.
+
+    Uses content hashing so unchanged files are not re-parsed on repeat runs.
+    """
+
+    def __init__(self, project_root: Path, max_entries: int = 500):
+        self._root = project_root.resolve()
+        self._source_cache: dict[str, str] = {}
+        self._ast_cache: dict[str, ast.Module | None] = {}
+        self._content_hashes: dict[str, str] = {}
+        self._max = max_entries
+
+    @property
+    def project_root(self) -> Path:
+        return self._root
+
+    def get_source(self, file_path: str) -> str | None:
+        """Read and cache the source text of a file.
+
+        Accepts both absolute and project-relative paths.
+        Uses content hashing to detect changes.
+        """
+        resolved = self._resolve(file_path)
+        if not resolved:
+            return None
+
+        key = str(resolved)
+
+        try:
+            raw = resolved.read_bytes()
+        except OSError as e:
+            logger.debug("Cannot read %s: %s", resolved, e)
+            return None
+
+        content_hash = hashlib.sha256(raw).hexdigest()
+        cached_hash = self._content_hashes.get(key)
+
+        if cached_hash == content_hash and key in self._source_cache:
+            return self._source_cache[key]
+
+        # File changed or first read — invalidate AST cache entry too
+        self._ast_cache.pop(key, None)
+
+        text = raw.decode("utf-8", errors="replace")
+        self._evict_if_needed()
+        self._source_cache[key] = text
+        self._content_hashes[key] = content_hash
+        return text
+
+    def get_ast(self, file_path: str) -> ast.Module | None:
+        """Parse and cache the AST for a Python file."""
+        resolved = self._resolve(file_path)
+        if not resolved:
+            return None
+
+        key = str(resolved)
+        if key in self._ast_cache:
+            return self._ast_cache[key]
+
+        source = self.get_source(file_path)
+        if source is None:
+            return None
+
+        try:
+            tree = ast.parse(source, filename=str(resolved))
+        except SyntaxError as e:
+            logger.debug("Cannot parse %s: %s", resolved, e)
+            self._ast_cache[key] = None
+            return None
+
+        self._ast_cache[key] = tree
+        return tree
+
+    def get_lines(self, file_path: str, start: int, end: int) -> list[str] | None:
+        """Return lines [start, end] (1-indexed inclusive) from a file."""
+        source = self.get_source(file_path)
+        if source is None:
+            return None
+        lines = source.splitlines()
+        s = max(0, start - 1)
+        e = min(len(lines), end)
+        return lines[s:e]
+
+    def get_function_node(
+        self,
+        file_path: str,
+        function_name: str,
+    ) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
+        """Find a function definition by name in a file's AST."""
+        tree = self.get_ast(file_path)
+        if tree is None:
+            return None
+
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                if node.name == function_name:
+                    return node
+        return None
+
+    def get_class_node(self, file_path: str, class_name: str) -> ast.ClassDef | None:
+        """Find a class definition by name in a file's AST."""
+        tree = self.get_ast(file_path)
+        if tree is None:
+            return None
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef) and node.name == class_name:
+                return node
+        return None
+
+    def get_function_source(self, file_path: str, function_name: str) -> str | None:
+        """Get the raw source text of a function."""
+        fn = self.get_function_node(file_path, function_name)
+        if fn is None:
+            return None
+
+        source = self.get_source(file_path)
+        if source is None:
+            return None
+
+        lines = source.splitlines()
+        start = fn.lineno - 1
+        end = fn.end_lineno or fn.lineno
+        return "\n".join(lines[start:end])
+
+    def invalidate(self, file_path: str | None = None) -> None:
+        """Clear cached entries. None clears all."""
+        if file_path is None:
+            self._source_cache.clear()
+            self._ast_cache.clear()
+            self._content_hashes.clear()
+            return
+
+        resolved = self._resolve(file_path)
+        if resolved:
+            key = str(resolved)
+            self._source_cache.pop(key, None)
+            self._ast_cache.pop(key, None)
+            self._content_hashes.pop(key, None)
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _resolve(self, file_path: str) -> Path | None:
+        p = Path(file_path)
+        resolved = p.resolve() if p.is_absolute() else (self._root / p).resolve()
+
+        if not resolved.is_relative_to(self._root):
+            logger.warning("Path traversal blocked: %s (resolves outside project root)", file_path)
+            return None
+
+        if not resolved.is_file():
+            logger.debug("File not found: %s", resolved)
+            return None
+        return resolved
+
+    def _evict_if_needed(self) -> None:
+        while len(self._source_cache) >= self._max:
+            oldest_key = next(iter(self._source_cache))
+            del self._source_cache[oldest_key]
+            self._ast_cache.pop(oldest_key, None)
+            self._content_hashes.pop(oldest_key, None)

apisec_code_bolt/query/executor.py

@@ -0,0 +1,283 @@
+"""
+Query executor: orchestrates the verification loop.
+
+Supports two modes:
+  - Connected: probe uploads manifest, polls engine for questions,
+    answers them, and submits responses over HTTP.
+  - Air-gapped: reads questions from a JSON file, answers them,
+    and writes responses to another JSON file.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from ..cloud.client import CloudClient, QuestionItem
+from .ast_cache import ASTCache
+from .handlers import AnswerOutput, QuestionInput, get_handler
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExecutionStats:
+    rounds_completed: int = 0
+    questions_answered: int = 0
+    errors: int = 0
+    elapsed_seconds: float = 0.0
+
+
+@dataclass
+class ExecutionResult:
+    analysis_id: str | None = None
+    final_status: str = "unknown"
+    stats: ExecutionStats = field(default_factory=ExecutionStats)
+    error: str | None = None
+
+
+class QueryExecutor:
+    """Orchestrate the verification question/answer loop."""
+
+    def __init__(
+        self,
+        project_root: Path,
+        max_batches: int = 10,
+        poll_timeout: int = 30,
+        max_wait: int = 300,
+    ):
+        self._root = project_root.resolve()
+        self._max_batches = max_batches
+        self._poll_timeout = poll_timeout
+        self._max_wait = max_wait
+        self._cache = ASTCache(self._root)
+
+    # ------------------------------------------------------------------
+    # Connected mode
+    # ------------------------------------------------------------------
+
+    def run_connected(
+        self,
+        client: CloudClient,
+        manifest_data: dict[str, Any],
+    ) -> ExecutionResult:
+        """Upload manifest to cloud, then run the verification loop."""
+        start = time.monotonic()
+        result = ExecutionResult()
+
+        try:
+            analysis_id, status = client.upload_manifest(manifest_data)
+            result.analysis_id = analysis_id
+            logger.info("Upload complete: analysis_id=%s status=%s", analysis_id, status)
+
+            if status == "complete":
+                result.final_status = "complete"
+                result.stats.elapsed_seconds = time.monotonic() - start
+                return result
+
+            self._run_loop(client, analysis_id, result)
+
+        except Exception as e:
+            result.error = str(e)
+            result.final_status = "error"
+            logger.error("Connected execution failed: %s", e)
+
+        result.stats.elapsed_seconds = time.monotonic() - start
+        return result
+
+    def _run_loop(
+        self,
+        client: CloudClient,
+        analysis_id: str,
+        result: ExecutionResult,
+    ) -> None:
+        """Poll for questions, answer them, submit, repeat."""
+        for round_num in range(1, self._max_batches + 1):
+            logger.info("Polling for questions (round %d)...", round_num)
+
+            poll_result = client.poll_questions(
+                analysis_id,
+                poll_timeout=self._poll_timeout,
+                max_wait=self._max_wait,
+            )
+
+            if poll_result.status == "complete":
+                result.final_status = "complete"
+                logger.info("Analysis complete after %d rounds", round_num - 1)
+                return
+
+            if poll_result.status in ("failed", "timeout"):
+                result.final_status = poll_result.status
+                return
+
+            if poll_result.status == "verifying":
+                # LLM verification running, no on-demand questions yet — keep polling
+                logger.debug("LLM verifying, no questions yet — will poll again")
+                time.sleep(2)
+                continue
+
+            if poll_result.status != "questions_ready" or not poll_result.questions:
+                logger.warning("Unexpected poll status: %s", poll_result.status)
+                result.final_status = poll_result.status
+                return
+
+            answers = self._answer_batch(poll_result.questions)
+            result.stats.rounds_completed = round_num
+            result.stats.questions_answered += len(answers)
+
+            serialized = [
+                {
+                    "question_id": a.question_id,
+                    "result": a.result,
+                    "evidence": a.evidence,
+                    "confidence": a.confidence,
+                    "details": a.details,
+                }
+                for a in answers
+            ]
+
+            submit_result = client.submit_answers(
+                analysis_id,
+                poll_result.batch_id or "",
+                serialized,
+            )
+
+            if not submit_result.more_questions:
+                if submit_result.status == "complete":
+                    result.final_status = "complete"
+                    logger.info("Analysis complete after round %d", round_num)
+                    return
+                if submit_result.status == "verifying":
+                    # LLM verification in progress — keep polling for on-demand questions
+                    logger.info(
+                        "LLM verification in progress, continuing to poll for on-demand questions"
+                    )
+                    continue
+                result.final_status = submit_result.status
+                logger.info("No more questions; status=%s", submit_result.status)
+                return
+
+        result.final_status = "max_batches_reached"
+        logger.warning("Max batches (%d) reached", self._max_batches)
+
+    # ------------------------------------------------------------------
+    # Air-gapped mode
+    # ------------------------------------------------------------------
+
+    def run_airgapped(
+        self,
+        questions_file: Path,
+        output_file: Path,
+    ) -> ExecutionResult:
+        """Read questions from file, answer them, write answers to file."""
+        start = time.monotonic()
+        result = ExecutionResult()
+
+        try:
+            raw = json.loads(questions_file.read_text())
+
+            questions_data = raw if isinstance(raw, list) else raw.get("questions", [])
+            batch_id = raw.get("batch_id") if isinstance(raw, dict) else None
+            analysis_id = raw.get("analysis_id") if isinstance(raw, dict) else None
+            result.analysis_id = analysis_id
+
+            question_items = [
+                QuestionItem(
+                    id=q["id"],
+                    type=q["type"],
+                    gate_id=q.get("gate_id", ""),
+                    finding_id=q.get("finding_id", ""),
+                    round_number=q.get("round_number", 1),
+                    target_file=q.get("target_file"),
+                    target_function=q.get("target_function"),
+                    target_line=q.get("target_line"),
+                    params=q.get("params", {}),
+                )
+                for q in questions_data
+            ]
+
+            answers = self._answer_batch(question_items)
+            result.stats.questions_answered = len(answers)
+            result.stats.rounds_completed = 1
+
+            output = {
+                "batch_id": batch_id,
+                "analysis_id": analysis_id,
+                "answers": [
+                    {
+                        "question_id": a.question_id,
+                        "result": a.result,
+                        "evidence": a.evidence,
+                        "confidence": a.confidence,
+                        "details": a.details,
+                    }
+                    for a in answers
+                ],
+            }
+
+            output_file.parent.mkdir(parents=True, exist_ok=True)
+            output_file.write_text(json.dumps(output, indent=2, default=str))
+
+            result.final_status = "complete"
+            logger.info(
+                "Air-gapped: answered %d questions -> %s",
+                len(answers),
+                output_file,
+            )
+
+        except Exception as e:
+            result.error = str(e)
+            result.final_status = "error"
+            logger.error("Air-gapped execution failed: %s", e)
+
+        result.stats.elapsed_seconds = time.monotonic() - start
+        return result
+
+    # ------------------------------------------------------------------
+    # Shared question answering
+    # ------------------------------------------------------------------
+
+    def _answer_batch(self, questions: list[QuestionItem]) -> list[AnswerOutput]:
+        answers: list[AnswerOutput] = []
+
+        for q in questions:
+            question_input = QuestionInput(
+                id=q.id,
+                type=q.type,
+                gate_id=q.gate_id,
+                finding_id=q.finding_id,
+                target_file=q.target_file,
+                target_function=q.target_function,
+                target_line=q.target_line,
+                params=q.params,
+            )
+
+            handler = get_handler(q.type)
+            if handler is None:
+                answers.append(
+                    AnswerOutput(
+                        question_id=q.id,
+                        result="error",
+                        details=f"No handler for question type: {q.type}",
+                    )
+                )
+                continue
+
+            try:
+                answer = handler.handle(question_input, self._cache)
+                answers.append(answer)
+            except Exception as e:
+                logger.error("Handler error for question %s: %s", q.id, e)
+                answers.append(
+                    AnswerOutput(
+                        question_id=q.id,
+                        result="error",
+                        details=str(e),
+                    )
+                )
+
+        return answers