shadowcat 2.0.0__py3-none-any.whl

agent/core/controller.py

@@ -0,0 +1,638 @@
+"""Agent controller with lifecycle management, pause/resume, and session persistence."""
+
+import asyncio
+import json
+import logging
+import os
+import re
+import sys
+from enum import Enum
+from typing import Any, ClassVar
+
+from agent.core.backend import (
+    AgentBackend,
+    AgentMessage,
+    AnthropicBackend,
+    ClaudeSDKBackend,
+    MessageType,
+    OpenRouterBackend,
+)
+from agent.core.config import ShadowCatConfig
+from agent.core.events import Event, EventBus, EventType
+from agent.core.profiling import mark, span
+from agent.core.session import SessionStatus, SessionStore
+from agent.parsing.tool_parser import ToolOutputParser
+
+# RAG is optional — the module pulls in sentence-transformers / chromadb /
+# langchain at import time, none of which ship in the default install. If any
+# of them is missing, set CTFKnowledgeBase=None and let the existing null-checks
+# below skip the RAG path.
+try:
+    from agent.rag_module.knowledge_base import CTFKnowledgeBase
+except ImportError:
+    CTFKnowledgeBase = None  # type: ignore[assignment,misc]
+
+
+class AgentState(Enum):
+    """Simple 5-state model for agent lifecycle."""
+
+    IDLE = "idle"
+    RUNNING = "running"
+    PAUSED = "paused"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+
+class AgentController:
+    """
+    Central orchestrator with lifecycle management.
+
+    Features:
+    - Framework-agnostic via AgentBackend
+    - Pause/resume/stop control
+    - Instruction injection
+    - Session persistence
+    """
+
+    # Flag detection patterns
+    FLAG_PATTERNS: ClassVar[list[str]] = [
+        r"(?:flag|FLAG|HTB|THM|CTF|picoCTF|HackTheBox)\{[^\}\s]+\}",
+        r"\b[a-f0-9]{32}\b",  # 32-char hex (HTB user/root flags)
+    ]
+
+    def __init__(
+        self,
+        config: ShadowCatConfig,
+        backend: AgentBackend | None = None,
+        session_store: SessionStore | None = None,
+        events: EventBus | None = None,
+    ):
+        """Initialize controller.
+
+        Args:
+            config: ShadowCat configuration
+            backend: Optional custom backend (defaults to ClaudeCodeBackend)
+            session_store: Optional custom session store
+            events: Optional custom event bus
+        """
+        self.config = config
+        self.backend = backend
+        self.sessions = session_store or SessionStore()
+        self.events = events or EventBus.get()
+
+        # State management
+        self._state = AgentState.IDLE
+        self._pause_requested = False
+        self._stop_requested = False
+        self._resume_event = asyncio.Event()
+        self._pending_instruction: str | None = None
+        self._tool_history: list[tuple[str, str]] = []
+
+        self._tool_parser = ToolOutputParser()
+        if CTFKnowledgeBase is None:
+            self._kb = None  # RAG deps not installed
+        else:
+            try:
+                # Reuse the process-wide singleton instead of constructing a
+                # second KB: the SentenceTransformer load is "heavy (hundreds of
+                # ms to several seconds)" (see knowledge_base.get_shared_kb), and
+                # the http_request tool in backend.py already loads it via
+                # get_shared_kb(). Two instances = the embedding model loaded
+                # twice at startup. Share one.
+                from agent.rag_module.knowledge_base import get_shared_kb
+
+                self._kb = get_shared_kb()
+            except Exception:
+                self._kb = None  # RAG init failed (e.g. embedding model download), degrade
+
+        # Subscribe to user events
+        self.events.subscribe(EventType.USER_COMMAND, self._on_user_command)
+        self.events.subscribe(EventType.USER_INPUT, self._on_user_input)
+
+    @property
+    def state(self) -> AgentState:
+        """Get current agent state."""
+        return self._state
+
+    def _set_state(
+        self,
+        state: AgentState,
+        details: str = "",
+        target: str | None = None,
+        task: str | None = None,
+    ) -> None:
+        """Update state and emit event.
+
+        Args:
+            state: New agent state
+            details: Optional details about the state
+            target: Optional target for session tracking (used by Langfuse)
+            task: Optional full task description for session tracking (used by Langfuse)
+        """
+        self._state = state
+        self.events.emit_state(state.value, details, target=target, task=task)
+
+    def _looks_like_loop(self) -> bool:
+        """Detect if the exact same tool with the exact same arguments was called 3 times in a row."""
+        if len(self._tool_history) < 3:
+            return False
+        last_three = self._tool_history[-3:]
+        return last_three[0] == last_three[1] == last_three[2]
+
+    # === Control Methods (called from TUI) ===
+
+    def pause(self) -> bool:
+        """Request pause at next safe point.
+
+        Returns:
+            True if pause request was accepted
+        """
+        if self._state == AgentState.RUNNING:
+            self._pause_requested = True
+            return True
+        return False
+
+    def resume(self, instruction: str | None = None) -> bool:
+        """Resume from paused state.
+
+        Args:
+            instruction: Optional instruction to inject on resume
+
+        Returns:
+            True if resume request was accepted
+        """
+        if self._state == AgentState.PAUSED:
+            self._pending_instruction = instruction
+            self._pause_requested = False
+            self._resume_event.set()
+            return True
+        return False
+
+    def stop(self) -> bool:
+        """Request stop.
+
+        Returns:
+            True (stop is always accepted)
+        """
+        self._stop_requested = True
+        self._resume_event.set()  # Unblock if paused
+        return True
+
+    def inject(self, instruction: str) -> bool:
+        """Queue instruction for next pause point.
+
+        Args:
+            instruction: Instruction to inject
+
+        Returns:
+            True if instruction was queued
+        """
+        if self._state in (AgentState.RUNNING, AgentState.PAUSED):
+            self._pending_instruction = instruction
+            if self._state == AgentState.RUNNING:
+                self._pause_requested = True
+            return True
+        return False
+
+    # === Event Handlers ===
+
+    def _on_user_command(self, event: Event) -> None:
+        """Handle user command events."""
+        cmd = event.data.get("command")
+        if cmd == "pause":
+            self.pause()
+        elif cmd == "resume":
+            self.resume()
+        elif cmd == "stop":
+            self.stop()
+
+    def _on_user_input(self, event: Event) -> None:
+        """Handle user input events."""
+        text = event.data.get("text", "")
+        if text:
+            self.inject(text)
+
+    # === Main Execution ===
+
+    async def run(self, task: str, resume_session_id: str | None = None) -> dict[str, Any]:
+        """Run agent with full lifecycle management.
+
+        Args:
+            task: Task description for the agent
+            resume_session_id: Optional session ID to resume
+
+        Returns:
+            Result dictionary with success, output, flags, etc.
+        """
+        # Reset state
+        self._pause_requested = False
+        self._stop_requested = False
+        self._resume_event.clear()
+
+        # Create or resume session
+        if resume_session_id:
+            session = self.sessions.load(resume_session_id)
+            if not session:
+                return {
+                    "success": False,
+                    "error": f"Session {resume_session_id} not found",
+                }
+            # Update task if resuming
+            if not task:
+                task = session.task
+        else:
+            session = self.sessions.create(
+                target=self.config.target,
+                task=task,
+                model=self.config.llm_model,
+            )
+
+        # Create backend if needed.
+        # Auth-mode-driven backend selection:
+        #   anthropic            -> AnthropicBackend  (Option 2, direct API key)
+        #   claude_subscription  -> ClaudeSDKBackend  (Option 4, OAuth via subscription)
+        #   anything else        -> OpenRouterBackend (Option 1 default; Option 3 local
+        #                                              also goes here — point MODEL at
+        #                                              the local OpenAI-compatible URL)
+        # See scripts/config.sh for how SHADOWCAT_AUTH_MODE gets set.
+        if self.backend is None:
+            from agent.prompts.pentesting import get_ctf_prompt
+
+            auth_mode = (os.getenv("SHADOWCAT_AUTH_MODE") or "openrouter").lower()
+            system_prompt = get_ctf_prompt(self.config.custom_instruction)
+            working_directory = str(self.config.working_directory)
+            model = self.config.llm_model
+
+            if auth_mode == "claude_subscription":
+                self.backend = ClaudeSDKBackend(
+                    working_directory=working_directory,
+                    system_prompt=system_prompt,
+                    model=model,
+                )
+            elif auth_mode == "anthropic":
+                self.backend = AnthropicBackend(
+                    working_directory=working_directory,
+                    system_prompt=system_prompt,
+                    model=model,
+                )
+            else:
+                self.backend = OpenRouterBackend(
+                    working_directory=working_directory,
+                    system_prompt=system_prompt,
+                    model=model,
+                )
+
+        try:
+            self._set_state(
+                AgentState.RUNNING,
+                "Connecting...",
+                target=self.config.target,
+                task=task,
+            )
+
+            # Connect (or resume)
+            with span("backend.connect/resume"):
+                if resume_session_id and self.backend.supports_resume:
+                    backend_session = session.backend_session_id or resume_session_id
+                    await self.backend.resume(backend_session)
+                    self.events.emit_message(f"Resumed session {resume_session_id}", "info")
+                else:
+                    await self.backend.connect()
+
+            # Store backend session ID if available
+            if self.backend.session_id:
+                self.sessions.set_backend_session_id(self.backend.session_id)
+
+            # ═══════════════════════════════════════════════════════════════
+            # Stage A: Planner pre-pass
+            # ───────────────────────────────────────────────────────────────
+            # One cheap LLM call BEFORE the Senior loop starts. Augments the
+            # task with a JSON exploit chain. Falls back silently on any
+            # failure so the Senior runs exactly as today if anything breaks.
+            #
+            # Heavy logging at every decision point so the next "Planner is
+            # silent" incident is diagnosable from a single log file. After
+            # the previous absent-file incident, observability beats clean
+            # code: every branch traces to stderr AND to logging.info.
+            # ═══════════════════════════════════════════════════════════════
+            planner_model = os.getenv("PLANNER_MODEL", "").strip()
+
+            def _planner_trace(msg: str) -> None:
+                logging.getLogger(__name__).info("[PLANNER] %s", msg)
+                print(f"[PLANNER] {msg}", file=sys.stderr, flush=True)
+
+            # f-strings can't combine !r with a conditional inside a single
+            # placeholder, so compute the display strings first.
+            _planner_model_display = repr(planner_model) if planner_model else "(UNSET - will skip)"
+            _resume_display = (
+                repr(resume_session_id) if resume_session_id else "(none - fresh session)"
+            )
+            _planner_trace(f"Planner Model detected: {_planner_model_display}")
+            _planner_trace(f"resume_session_id: {_resume_display}")
+
+            if not planner_model:
+                _planner_trace(
+                    "Planner attempt: SKIPPED (PLANNER_MODEL env var is empty "
+                    "or unset). To enable, set PLANNER_MODEL in .env.auth or "
+                    "export it before `make connect`. Continuing un-planned."
+                )
+            elif resume_session_id:
+                _planner_trace(
+                    "Planner attempt: SKIPPED (resuming existing session; "
+                    "original plan is already in the conversation history). "
+                    "Continuing un-planned."
+                )
+            else:
+                _planner_trace(f"Planner attempt: STARTING with model={planner_model!r}")
+                try:
+                    from agent.core.planner import (
+                        format_plan_for_injection,
+                        gather_planning_context,
+                        generate_plan,
+                    )
+
+                    # Pre-plan enrichment: fingerprint a web target (PHANTOM)
+                    # and retrieve framework CVEs (RAG) so the FIRST plan is
+                    # grounded in the real stack, not parametric guessing.
+                    #
+                    # PERF: enrichment adds a SECOND serial LLM round-trip
+                    # (crystallize_fingerprint) before generate_plan. On a slow
+                    # reasoning model that doubles pre-plan latency, so it's now
+                    # opt-in via PLANNER_ENRICHMENT (default off) — the common
+                    # path makes a single planner call. Set PLANNER_ENRICHMENT=1
+                    # to trade ~one extra round-trip for a stack-grounded plan.
+                    # Best-effort either way — returns (None, None) on failure.
+                    enrichment_on = bool(os.getenv("PLANNER_ENRICHMENT", "").strip())
+                    if enrichment_on:
+                        with span("planner.enrichment (PHANTOM+RAG)"):
+                            phantom_findings, rag_context = await gather_planning_context(
+                                target=self.config.target,
+                                fingerprint_model=planner_model,
+                                kb=self._kb,
+                            )
+                    else:
+                        phantom_findings, rag_context = None, None
+                        _planner_trace(
+                            "Pre-plan enrichment: SKIPPED (PLANNER_ENRICHMENT "
+                            "unset). Single-call planning. Set PLANNER_ENRICHMENT=1 "
+                            "to fingerprint the target + RAG before planning "
+                            "(adds one LLM round-trip)."
+                        )
+                    _planner_trace(
+                        "Pre-plan enrichment: "
+                        f"phantom={'yes' if phantom_findings else 'no'}, "
+                        f"rag={'yes' if rag_context else 'no'}"
+                    )
+
+                    with span(f"planner.generate_plan [{planner_model}]"):
+                        plan, plan_cost = await generate_plan(
+                            task=task,
+                            model=planner_model,
+                            phantom_findings=phantom_findings,
+                            rag_context=rag_context,
+                        )
+                    if plan_cost > 0:
+                        self.sessions.add_cost(plan_cost)
+                    injection = format_plan_for_injection(plan)
+                    # Append (not prepend) so operator hint at top of task
+                    # outranks the planner on conflict.
+                    task = task + "\n\n" + injection
+
+                    snippet = json.dumps(
+                        {
+                            "hypothesis": plan.overall_hypothesis[:120],
+                            "step_count": len(plan.steps),
+                            "first_step": plan.steps[0].goal[:120] if plan.steps else None,
+                        },
+                        default=str,
+                    )
+                    _planner_trace(f"Planner attempt: SUCCESS — {snippet}")
+                    _planner_trace(f"Injection output (first 300 chars): {injection[:300]!r}")
+                    _planner_trace(
+                        f"Augmented task: original_len={len(task) - len(injection) - 2}, "
+                        f"new_len={len(task)}, plan_cost=${plan_cost:.4f}"
+                    )
+                    self.events.emit_message(
+                        f"Planner produced a {len(plan.steps)}-step plan "
+                        f"(model={planner_model}, cost=${plan_cost:.4f})",
+                        "info",
+                    )
+                except Exception as exc:
+                    # Planner is best-effort — never fatal. Log loudly so the
+                    # operator sees WHY it skipped (vs. silent failure).
+                    _planner_trace(
+                        f"Planner attempt: FAILURE — {exc.__class__.__name__}: {exc}. "
+                        f"Falling back to un-planned baseline."
+                    )
+                    self.events.emit_message(
+                        f"Planner skipped ({exc.__class__.__name__}: {exc}); running un-planned",
+                        "info",
+                    )
+
+            # Send initial query (now potentially augmented with the plan)
+            mark("agent loop: dispatching first query (pre-pass complete)")
+            await self.backend.query(task)
+            self.sessions.update_status(SessionStatus.RUNNING)
+
+            # Process messages with pause/stop handling
+            output_parts: list[str] = []
+            flags_found: list[str] = []
+            backend_error: str | None = None
+
+            async for msg in self.backend.receive_messages():
+                # Backend-reported failure (timeout, rate limit, oversized
+                # prompt, auth failure). Capture the message, break out of
+                # the loop, and let the post-loop check below raise so the
+                # outer try/except sets ERROR state and the spinner stops.
+                if msg.type == MessageType.ERROR:
+                    backend_error = (
+                        msg.content if isinstance(msg.content, str) else str(msg.content)
+                    )
+                    self.events.emit_message(backend_error, "error")
+                    break
+
+                # Check stop request
+                if self._stop_requested:
+                    self._set_state(AgentState.IDLE, "Stopped by user")
+                    self.sessions.update_status(SessionStatus.PAUSED)
+                    break
+
+                # Check pause request (between messages = safe point)
+                if self._pause_requested:
+                    self._pause_requested = False
+                    self._set_state(AgentState.PAUSED, "Paused - waiting for input")
+                    self.sessions.update_status(SessionStatus.PAUSED)
+
+                    # Wait for resume
+                    await self._resume_event.wait()
+                    self._resume_event.clear()
+
+                    if self._stop_requested:
+                        break
+
+                    # Resume with pending instruction
+                    self._set_state(AgentState.RUNNING, "Resumed")
+                    self.sessions.update_status(SessionStatus.RUNNING)
+
+                    if self._pending_instruction:
+                        self.sessions.add_instruction(self._pending_instruction)
+                        self.events.emit_message(
+                            f"Injecting: {self._pending_instruction[:50]}...", "info"
+                        )
+                        await self.backend.query(self._pending_instruction)
+                        self._pending_instruction = None
+
+                # Process message by type
+                await self._process_message(msg, output_parts, flags_found)
+
+            # Backend signalled failure (timeout/auth/rate-limit/etc.) via an
+            # ERROR message. Surface as a clean failure rather than reporting
+            # success with whatever partial output we accumulated.
+            if backend_error:
+                self._set_state(AgentState.ERROR, backend_error)
+                self.sessions.set_error(backend_error)
+                self.sessions.update_status(SessionStatus.ERROR)
+                return {
+                    "success": False,
+                    "error": backend_error,
+                    "output": "\n".join(output_parts),
+                    "flags_found": flags_found,
+                    "session_id": session.session_id,
+                    "cost_usd": session.total_cost_usd,
+                }
+
+            # Completed successfully
+            if not self._stop_requested:
+                self._set_state(AgentState.COMPLETED)
+                self.sessions.update_status(SessionStatus.COMPLETED)
+
+            return {
+                "success": True,
+                "output": "\n".join(output_parts),
+                "flags_found": flags_found,
+                "session_id": session.session_id,
+                "cost_usd": session.total_cost_usd,
+            }
+
+        except Exception as e:
+            self._set_state(AgentState.ERROR, str(e))
+            self.sessions.set_error(str(e))
+            self.sessions.update_status(SessionStatus.ERROR)
+            return {"success": False, "error": str(e)}
+
+        finally:
+            if self.backend:
+                await self.backend.disconnect()
+
+    async def _process_message(
+        self, msg: AgentMessage, output_parts: list[str], flags_found: list[str]
+    ) -> None:
+        """Process a single agent message.
+
+        Args:
+            msg: Message to process
+            output_parts: List to append text output to
+            flags_found: List to append found flags to
+        """
+        if msg.type == MessageType.TEXT:
+            output_parts.append(msg.content)
+            self.events.emit_message(msg.content)
+
+            # Detect flags
+            detected = self._detect_flags(msg.content)
+            for flag in detected:
+                if flag not in flags_found:
+                    flags_found.append(flag)
+                    self.sessions.add_flag(flag, msg.content[:200])
+                    self.events.emit_flag(flag, msg.content[:200])
+
+        elif msg.type == MessageType.TOOL_START:
+            tool_name = msg.tool_name or "unknown"
+            tool_args = msg.tool_args or {}
+
+            self.events.emit_tool(
+                status="start",
+                name=tool_name,
+                args=tool_args,
+            )
+
+            # Loop Detection & Replan
+            import json
+
+            args_str = json.dumps(tool_args, sort_keys=True)
+            self._tool_history.append((tool_name, args_str))
+
+            if self._looks_like_loop():
+                self.inject(
+                    "LOOP DETECTED: you have run the same command 3 times. "
+                    'Output a JSON object: {"why_stuck": "...", "new_hypothesis": "..."}. '
+                    "Then execute a new chain."
+                )
+
+        elif msg.type == MessageType.TOOL_RESULT:
+            tool_name = msg.tool_name or "unknown"
+            raw_content = msg.content or ""
+            try:
+                summarized = self._tool_parser.summarize_output(tool_name, raw_content)
+            except Exception as e:
+                logging.getLogger(__name__).warning(
+                    "Tool parser failed for %s: %s. Falling back to raw output.",
+                    tool_name,
+                    e,
+                )
+                summarized = raw_content
+
+            # Scan raw tool output for flags — the LLM may not echo a flag it
+            # already saw in `cat flag.txt` output, so we'd otherwise miss it
+            # and trigger a needless retry cycle. Use raw_content (not the
+            # summarized JSON) so we don't depend on summarizer behaviour.
+            for flag in self._detect_flags(raw_content):
+                if flag not in flags_found:
+                    flags_found.append(flag)
+                    context_snippet = raw_content[:200] if isinstance(raw_content, str) else ""
+                    self.sessions.add_flag(flag, context_snippet)
+                    self.events.emit_flag(flag, context_snippet)
+
+            self.events.emit_tool(
+                status="complete",
+                name=tool_name,
+                result=summarized,
+            )
+
+            if tool_name.lower() == "nmap" and self._kb is not None:
+                try:
+                    parsed = self._tool_parser.parse_nmap(raw_content)
+                    for entry in parsed.get("open_ports", []):
+                        service = entry.get("service", "unknown")
+                        port = str(entry.get("port", ""))
+                        chunks = await self._kb.query_technique(service, port)
+                        if chunks:
+                            rag_context = (
+                                f"[RAG] Techniques for {service} on port {port}:\n"
+                                + "\n---\n".join(chunks)
+                            )
+                            self.inject(rag_context)
+                except Exception as exc:
+                    logging.getLogger(__name__).debug("RAG query failed: %s", exc)
+
+        elif msg.type == MessageType.RESULT:
+            cost = msg.metadata.get("cost_usd", 0)
+            if cost > 0:
+                self.sessions.add_cost(cost)
+
+    def _detect_flags(self, text: str) -> list[str]:
+        """Detect potential flags in text.
+
+        Args:
+            text: Text to search for flags
+
+        Returns:
+            List of detected flag strings
+        """
+        flags = []
+        for pattern in self.FLAG_PATTERNS:
+            for match in re.finditer(pattern, text, re.IGNORECASE):
+                flag = match.group(0)
+                if flag not in flags:
+                    flags.append(flag)
+        return flags