r3d-agent 0.2.0__py3-none-any.whl

r3d/__init__.py

@@ -0,0 +1,3 @@
+"""R3D — Autonomous AI Penetration Testing CLI Agent."""
+
+__version__ = "0.2.0"

r3d/__main__.py

@@ -0,0 +1,6 @@
+"""Allow running R3D via `python -m r3d`."""
+
+from r3d.cli import app
+
+if __name__ == "__main__":
+    app()

r3d/agent.py

@@ -0,0 +1,638 @@
+"""ReAct loop: plan → execute → observe → analyze → repeat (§4/§5 agent.py).
+
+The core agent loop. Includes ToolCallState, _stream_with_failover,
+and proper message building for all provider types.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+import signal
+import sys
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, AsyncIterator
+
+from r3d import redact
+from r3d.audit import AuditEvent, AuditLog
+from r3d.context import Session
+from r3d.errors import FatalProviderError, RateLimitError, TransientProviderError, with_retry
+from r3d.logger import Logger
+from r3d.planner import Planner
+from r3d.providers.base import (
+    ProviderEvent,
+    StreamDone,
+    TextDelta,
+    ToolCallComplete,
+    ToolCallDelta,
+    ToolCallStart,
+)
+from r3d.safety import SafetyChecker, SafetyDecision
+from r3d.system_prompt import build_system_prompt
+from r3d.tools.registry import ToolRegistry
+
+if TYPE_CHECKING:
+    from r3d.config import Config
+    from r3d.providers.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ToolCallState:
+    """Tracks a pending tool call during streaming."""
+    id: str
+    name: str
+    args_json: str = ""
+
+
+@dataclass
+class AgentResult:
+    goal: str
+    steps_taken: int
+    elapsed_seconds: float
+    stop_reason: str  # done|step_limit|time_limit|token_limit|interrupted|provider_error
+    summary: str
+    session_id: str
+    findings: list = field(default_factory=list)
+    tool_calls: int = 0
+    tool_errors: int = 0
+
+
+class Agent:
+    """The core ReAct agent."""
+
+    def __init__(
+        self,
+        provider: "LLMProvider",
+        tools: ToolRegistry,
+        session: Session,
+        log: Logger,
+        config: "Config",
+        safety: SafetyChecker,
+        audit: AuditLog,
+        failover_provider: "LLMProvider | None" = None,
+    ):
+        self._provider = provider
+        self._failover = failover_provider
+        self._tools = tools
+        self._session = session
+        self._logger = log
+        self._config = config
+        self._safety = safety
+        self._audit = audit
+        self._interrupted = False
+        self._total_tool_calls = 0
+        self._total_tool_errors = 0
+        self._replans_used = 0
+        self._created_files: list[str] = []  # track files created by agent
+
+    # ------------------------------------------------------------------ #
+    #  Main entry point                                                    #
+    # ------------------------------------------------------------------ #
+
+    async def run(self, goal: str) -> AgentResult:
+        """Run the agent to completion."""
+        try:
+            signal.signal(signal.SIGINT, self._graceful_shutdown)
+        except (OSError, ValueError):
+            pass
+
+        # Snapshot files before run — for cleanup later
+        self._pre_snapshot = self._snapshot_files()
+
+        self._audit.record(AuditEvent(
+            ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            session_id=self._session.session_id,
+            kind="session", tool="-", target=goal[:80],
+            decision="started",
+        ))
+
+        # Phase 1: Planning — goal as first user message
+        self._session.add_message("user", f"Goal: {goal}")
+
+        if self._config.enable_planner:
+            try:
+                planner = Planner(self._provider, self._logger)
+                plan = await with_retry(
+                    lambda: planner.make_plan(goal),
+                    max_retries=self._config.max_retries,
+                    base_delay=self._config.retry_base_delay,
+                )
+                self._logger.plan_display(plan)
+                self._session.plan = plan.steps
+                self._session.add_message("assistant", plan.raw_text)
+                self._session.add_message(
+                    "user",
+                    f"ดำเนินการตามแผนข้างต้นเพื่อบรรลุเป้าหมาย: {goal}",
+                )
+                self._session.save()
+            except Exception as e:
+                self._logger.warn(f"Planning failed: {e} — proceeding without plan")
+
+        stop_reason = "done"
+        full_text = ""
+
+        # Phase 2: ReAct loop (runs until <DONE>, time, or token/step limit)
+        max_steps = self._config.max_steps or 0
+        steps_since_last_tool = 0
+        while True:
+            # Check step limit
+            if max_steps and self._session.step_count >= max_steps:
+                stop_reason = "step_limit"
+                break
+
+            reason = self._check_stop()
+            if reason:
+                stop_reason = reason
+                break
+
+            self._session.step_count += 1
+            elapsed = time.time() - self._session.created_at
+            self._logger.step_banner(
+                self._session.step_count,
+                elapsed, self._session.estimated_tokens,
+            )
+
+            if self._session.estimated_tokens > self._config.max_tokens * 0.8:
+                self._session.compact()
+                self._logger.warn("Context compacted — oldest messages removed.")
+
+            messages = self._build_messages()
+            system = build_system_prompt(self._tools.list_tools(), self._config)
+            schemas = self._tools.get_schemas(
+                "anthropic" if self._config.provider == "anthropic" else "openai"
+            )
+
+            # Stream with failover
+            full_text = ""
+            tool_calls_this_step: list[tuple[str, str, dict, str]] = []
+            pending: dict[str, ToolCallState] = {}
+
+            self._logger.start_thinking("กำลังคิด")
+
+            try:
+                async for event in self._stream_with_failover(
+                    messages, system, schemas
+                ):
+                    if isinstance(event, TextDelta):
+                        self._logger.stream_text(event.text)
+                        full_text += event.text
+
+                    elif isinstance(event, ToolCallStart):
+                        self._logger.stop_thinking()
+                        self._logger.end_stream()
+                        pending[event.id] = ToolCallState(
+                            id=event.id, name=event.name
+                        )
+
+                    elif isinstance(event, ToolCallDelta):
+                        if event.id in pending:
+                            pending[event.id].args_json += event.json_fragment
+
+                    elif isinstance(event, ToolCallComplete):
+                        result_str = await self._handle_tool_call(
+                            event.id, event.name, event.args
+                        )
+                        tool_calls_this_step.append(
+                            (event.id, event.name, event.args, result_str)
+                        )
+
+                    elif isinstance(event, StreamDone):
+                        self._session.estimated_tokens += (
+                            event.input_tokens + event.output_tokens
+                        )
+
+            except FatalProviderError as e:
+                self._logger.stop_thinking()
+                self._logger.error(f"Provider error (fatal): {e}")
+                stop_reason = "provider_error"
+                break
+            except Exception as e:
+                self._logger.stop_thinking()
+                self._logger.error(f"Provider error: {e}")
+                stop_reason = "provider_error"
+                break
+
+            self._logger.stop_thinking()
+            self._logger.end_stream()
+
+            # ── Text-based tool call fallback ──────────────────────
+            # When models don't support native function calling, they
+            # write tool calls as text: [Tool Call: name] {json_args}
+            # Parse and execute these as real tool calls.
+            if not tool_calls_this_step:
+                text_calls = self._parse_text_tool_calls(full_text)
+                if text_calls:
+                    for tc_id, tc_name, tc_args in text_calls:
+                        result_str = await self._handle_tool_call(
+                            tc_id, tc_name, tc_args
+                        )
+                        tool_calls_this_step.append(
+                            (tc_id, tc_name, tc_args, result_str)
+                        )
+                else:
+                    # No parseable tool calls found — check if model output
+                    # contains only a <DONE> sentinel or natural language.
+                    # If no <DONE> and no tool calls, we'll nudge below.
+                    pass
+
+            # Persist
+            self._session.add_assistant_turn(full_text, tool_calls_this_step)
+            self._session.updated_at = time.time()
+            self._session.save()
+
+            # Sentinels
+            if "<DONE>" in full_text:
+                stop_reason = "done"
+                break
+
+            if (m := re.search(r"<REPLAN:\s*(.+?)>", full_text)) and self._replans_left():
+                try:
+                    planner = Planner(self._provider, self._logger)
+                    plan = await planner.replan(
+                        goal, self._session.plan, [m.group(1)]
+                    )
+                    self._logger.plan_display(plan)
+                    self._session.plan = plan.steps
+                    self._replans_used += 1
+                except Exception:
+                    pass
+
+            # Nudge if stalled — remind to use tools AND record findings
+            if not tool_calls_this_step and "<DONE>" not in full_text:
+                findings_count = len(self._session.findings)
+                if findings_count == 0:
+                    self._session.add_message(
+                        "user",
+                        "ดำเนินการต่อ ใช้ tools เพื่อทำงาน "
+                        "สำคัญ: เมื่อพบข้อมูลที่น่าสนใจหรือช่องโหว่ ให้ใช้ record_finding บันทึกทันที "
+                        "อย่าลืมบันทึก findings ก่อนส่ง <DONE>",
+                    )
+                else:
+                    self._session.add_message(
+                        "user",
+                        "ดำเนินการต่อ ใช้ tools เพื่อทำงาน หรือส่ง <DONE> เมื่อเสร็จสิ้น",
+                    )
+
+        # Summary
+        result = AgentResult(
+            goal=goal,
+            steps_taken=self._session.step_count,
+            elapsed_seconds=time.time() - self._session.created_at,
+            stop_reason=stop_reason,
+            summary=self._extract_summary(full_text),
+            session_id=self._session.session_id,
+            findings=self._session.findings,
+            tool_calls=self._total_tool_calls,
+            tool_errors=self._total_tool_errors,
+        )
+        self._logger.session_summary(result)
+        self._session.save()
+
+        self._audit.record(AuditEvent(
+            ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            session_id=self._session.session_id,
+            kind="session", tool="-", target=goal[:80],
+            decision="completed", note=stop_reason,
+        ))
+
+        # Cleanup files created by the agent
+        self._cleanup_files()
+
+        return result
+
+    # ------------------------------------------------------------------ #
+    #  Stream with failover (§12 retry.py)                                 #
+    # ------------------------------------------------------------------ #
+
+    async def _stream_with_failover(
+        self,
+        messages: list[dict],
+        system: str,
+        schemas: list[dict],
+    ) -> AsyncIterator[ProviderEvent]:
+        """Try primary provider, then failover on transient/rate errors."""
+        providers = [self._provider]
+        if self._failover:
+            providers.append(self._failover)
+
+        for i, provider in enumerate(providers):
+            try:
+                stream = provider.stream_with_tools(
+                    messages, system, schemas,
+                    temperature=self._config.temperature,
+                    max_tokens=self._config.max_output_tokens,
+                )
+                async for event in stream:
+                    yield event
+                return  # success
+            except FatalProviderError:
+                raise  # never retry fatal
+            except (RateLimitError, TransientProviderError) as e:
+                if i < len(providers) - 1:
+                    self._logger.warn(f"Primary failed ({e}), switching to failover")
+                    self._audit.record(AuditEvent(
+                        ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                        session_id=self._session.session_id,
+                        kind="session", tool="-",
+                        decision="allowed", note=f"failover: {e}",
+                    ))
+                    continue
+                raise
+
+    # ------------------------------------------------------------------ #
+    #  Tool call handler                                                   #
+    # ------------------------------------------------------------------ #
+
+    async def _handle_tool_call(
+        self, tc_id: str, name: str, args: dict
+    ) -> str:
+        """Safety → audit → execute → audit."""
+        args_summary = redact.scrub_str(json.dumps(args, default=str))[:120]
+        self._logger.tool_call_box(name, args)  # pass full dict, not truncated string
+        self._total_tool_calls += 1
+
+        # 1) Safety
+        decision = SafetyDecision(allowed=True, reason="no check needed")
+        if name == "run_cmd":
+            decision = self._safety.check_command(args.get("cmd", ""))
+        elif name == "file" and args.get("action") == "write":
+            decision = self._safety.check_file_write(args.get("path", ""))
+
+        if not decision.allowed:
+            self._logger.warn(f"BLOCKED: {decision.reason}")
+            self._audit.record(AuditEvent(
+                ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                session_id=self._session.session_id,
+                kind="denied", tool=name, args_summary=args_summary,
+                decision="denied", risk_level=decision.risk_level,
+                matched_pattern=decision.matched_pattern,
+            ))
+            self._total_tool_errors += 1
+            return f"ERROR: Blocked by safety policy — {decision.reason}"
+
+        if decision.requires_confirm:
+            if not self._logger.confirm_prompt(
+                f"[{decision.risk_level.upper()}] Allow: {args_summary[:80]}?"
+            ):
+                self._audit.record(AuditEvent(
+                    ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                    session_id=self._session.session_id,
+                    kind="denied", tool=name, args_summary=args_summary,
+                    decision="denied",
+                ))
+                return "ERROR: User declined this command."
+            self._audit.record(AuditEvent(
+                ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                session_id=self._session.session_id,
+                kind="confirm", tool=name, args_summary=args_summary,
+                decision="confirmed",
+            ))
+
+        # 2) Audit intent
+        self._audit.record(AuditEvent(
+            ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            session_id=self._session.session_id,
+            kind="intent", tool=name, args_summary=args_summary,
+            decision="allowed", risk_level=decision.risk_level,
+        ))
+
+        # 3) Execute
+        t0 = time.time()
+        try:
+            tool_result = await self._tools.dispatch(name, args)
+            elapsed = time.time() - t0
+            output = tool_result.format_output(max_lines=100)
+            self._logger.tool_result_box(
+                name, output, tool_result.truncated, elapsed
+            )
+            self._audit.record(AuditEvent(
+                ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                session_id=self._session.session_id,
+                kind="result", tool=name,
+                return_code=tool_result.return_code,
+                elapsed_ms=elapsed * 1000, decision="allowed",
+            ))
+            return output
+        except Exception as e:
+            err = f"ERROR: {type(e).__name__}: {e}"
+            self._logger.tool_error_box(name, err)
+            self._total_tool_errors += 1
+            self._audit.record(AuditEvent(
+                ts=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                session_id=self._session.session_id,
+                kind="result", tool=name,
+                return_code=-1, decision="allowed",
+            ))
+            return err
+
+    # ------------------------------------------------------------------ #
+    #  Message building                                                    #
+    # ------------------------------------------------------------------ #
+
+    def _build_messages(self) -> list[dict]:
+        """Build messages for the provider.
+
+        Ensures:
+          1. First message is always role=user
+          2. Last message is always role=user
+          3. Internal roles converted for API compatibility
+          4. No consecutive same-role messages
+        """
+        messages = list(self._session.messages)
+        if not messages:
+            messages.append({"role": "user", "content": f"Goal: {self._session.goal}"})
+
+        if messages[0].get("role") != "user":
+            messages.insert(0, {"role": "user", "content": f"Goal: {self._session.goal}"})
+
+        if messages[-1].get("role") != "user":
+            messages.append({
+                "role": "user",
+                "content": "ดำเนินการต่อ ใช้ tools เพื่อทำงาน หรือส่ง <DONE> เมื่อเสร็จสิ้น",
+            })
+
+        # Convert internal roles
+        filtered = []
+        for m in messages:
+            role = m.get("role", "")
+            if role in ("user", "assistant", "system"):
+                filtered.append(m)
+            elif role == "tool_call":
+                name = m.get("name", "?")
+                args = m.get("args", {})
+                filtered.append({
+                    "role": "assistant",
+                    "content": f"[Tool Call: {name}] {json.dumps(args, default=str)[:200]}",
+                })
+            elif role == "tool_result":
+                content = m.get("content", "")
+                filtered.append({
+                    "role": "user",
+                    "content": f"[Tool Result] {content[:500]}",
+                })
+
+        # Merge consecutive same-role
+        merged: list[dict] = []
+        for m in filtered:
+            if merged and merged[-1]["role"] == m["role"]:
+                merged[-1]["content"] += "\n" + m.get("content", "")
+            else:
+                merged.append(dict(m))
+
+        return merged
+
+    # ------------------------------------------------------------------ #
+    #  Stop / shutdown                                                     #
+    # ------------------------------------------------------------------ #
+
+    def _check_stop(self) -> str | None:
+        """Only stop on interrupt or token limit.
+
+        No time limit, no step limit — the agent runs until:
+          1. Model emits <DONE>
+          2. User presses Ctrl+C (interrupted)
+          3. Token budget exceeded
+          4. Provider error (handled in the stream loop)
+        """
+        if self._interrupted:
+            return "interrupted"
+        if self._session.estimated_tokens >= self._config.max_tokens:
+            return "token_limit"
+        return None
+
+    def _replans_left(self) -> bool:
+        return self._replans_used < self._config.max_replans
+
+    def _graceful_shutdown(self, sig: int, frame: object) -> None:
+        if self._interrupted:
+            sys.exit(1)
+        self._interrupted = True
+        self._logger.warn("Interrupted — finishing current step...")
+        self._session.save()
+
+    @staticmethod
+    def _extract_summary(text: str) -> str:
+        cleaned = text.replace("<DONE>", "").strip()
+        # Remove text-based tool call artifacts
+        cleaned = re.sub(r'\[Tool Call:.*?\]\s*\{[^}]*\}', '', cleaned)
+        cleaned = cleaned.strip()
+        if len(cleaned) > 500:
+            return cleaned[:500] + "..."
+        return cleaned or "(no summary)"
+
+    # ------------------------------------------------------------------ #
+    #  File tracking & cleanup                                             #
+    # ------------------------------------------------------------------ #
+
+    def _snapshot_files(self) -> set[str]:
+        """No-op — kept for compatibility. Cleanup now uses fixed /tmp/r3d."""
+        return set()
+
+    def _cleanup_files(self) -> None:
+        """Remove agent workdir — all agent files live here."""
+        import shutil
+        from pathlib import Path
+        from r3d.tools.shell import AGENT_WORKDIR
+
+        workdir = AGENT_WORKDIR
+        if not workdir.exists():
+            return
+
+        files: list[str] = []
+        try:
+            for f in workdir.rglob("*"):
+                if f.is_file():
+                    files.append(f.name)
+        except Exception:
+            pass
+
+        if not files:
+            return
+
+        try:
+            shutil.rmtree(workdir, ignore_errors=True)
+            self._logger.success(f"🧹 ลบไฟล์ที่ agent สร้างใน {workdir}: {len(files)} ไฟล์")
+            for name in files[:10]:
+                self._logger.console.print(f"     🗑️  {name}", style="r3d.dim")
+            if len(files) > 10:
+                self._logger.console.print(
+                    f"     … และอีก {len(files) - 10} ไฟล์", style="r3d.dim"
+                )
+        except Exception as e:
+            self._logger.warn(f"ลบ {workdir} ไม่สำเร็จ: {e}")
+
+    # ------------------------------------------------------------------ #
+    #  Text-based tool call parser                                         #
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _parse_text_tool_calls(text: str) -> list[tuple[str, str, dict]]:
+        """Parse tool calls written as text by models without native function calling.
+
+        Recognizes patterns like:
+          [Tool Call: run_cmd] {"cmd": "nmap ...", "timeout": 30}
+          [tool_call: file] {"action": "read", "path": "test.txt"}
+          ```json\n{"tool": "run_cmd", "cmd": "..."}\n```
+
+        Returns list of (id, name, args) tuples.
+        """
+        calls: list[tuple[str, str, dict]] = []
+        call_id = 0
+
+        # Pattern 1: [Tool Call: name] {json}
+        pattern1 = re.compile(
+            r'\[Tool\s*Call:\s*(\w+)\]\s*(\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\})',
+            re.IGNORECASE | re.DOTALL,
+        )
+        for m in pattern1.finditer(text):
+            name = m.group(1).strip()
+            raw_json = m.group(2).strip()
+            try:
+                args = json.loads(raw_json)
+                call_id += 1
+                calls.append((f"text_call_{call_id}", name, args))
+            except json.JSONDecodeError:
+                continue
+
+        # Pattern 2: ```json {"tool": "name", ...} ```  or  {"name": "run_cmd", "cmd": ...}
+        if not calls:
+            pattern2 = re.compile(
+                r'```(?:json)?\s*(\{[^`]+\})\s*```',
+                re.DOTALL,
+            )
+            for m in pattern2.finditer(text):
+                raw_json = m.group(1).strip()
+                try:
+                    obj = json.loads(raw_json)
+                    if isinstance(obj, dict):
+                        name = obj.pop("tool", obj.pop("name", ""))
+                        if name and isinstance(name, str):
+                            call_id += 1
+                            calls.append((f"text_call_{call_id}", name, obj))
+                except json.JSONDecodeError:
+                    continue
+
+        # Pattern 3: Inline JSON with known tool names
+        if not calls:
+            known_tools = ("run_cmd", "file", "record_finding", "exploit",
+                           "payload", "open_url", "web_search", "post")
+            for tool_name in known_tools:
+                # Match: run_cmd({"cmd": "..."}) or run_cmd {"cmd": "..."}
+                pat = re.compile(
+                    rf'\b{tool_name}\s*\(?\s*(\{{[^{{}}]*(?:\{{[^{{}}]*\}}[^{{}}]*)*\}})\s*\)?',
+                    re.DOTALL,
+                )
+                for m in pat.finditer(text):
+                    raw_json = m.group(1).strip()
+                    try:
+                        args = json.loads(raw_json)
+                        call_id += 1
+                        calls.append((f"text_call_{call_id}", tool_name, args))
+                    except json.JSONDecodeError:
+                        continue
+
+        return calls