aiptx 2.0.2__py3-none-any.whl

aipt_v2/tools/cloud/scoutsuite_tool.py

@@ -0,0 +1,359 @@
+"""
+ScoutSuite Integration - Multi-Cloud Security Auditing
+
+ScoutSuite is an open-source multi-cloud security auditing tool
+that assesses the security posture of cloud environments.
+
+Supports:
+- AWS (Amazon Web Services)
+- Azure (Microsoft Azure)
+- GCP (Google Cloud Platform)
+- Alibaba Cloud
+- Oracle Cloud Infrastructure
+
+Usage:
+    from aipt_v2.tools.cloud import run_scoutsuite
+
+    findings = await run_scoutsuite("aws", profile="production")
+"""
+
+import asyncio
+import json
+import os
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+
+from aipt_v2.tools.cloud.cloud_config import CloudConfig, get_cloud_config
+
+
+@dataclass
+class ScoutSuiteConfig:
+    """ScoutSuite configuration."""
+    provider: str = "aws"  # aws, azure, gcp, aliyun, oci
+
+    # AWS options
+    aws_profile: str = "default"
+    aws_regions: List[str] = field(default_factory=list)  # Empty = all
+
+    # Azure options
+    azure_subscription_id: str = ""
+    azure_tenant_id: str = ""
+
+    # GCP options
+    gcp_project_id: str = ""
+    gcp_service_account: str = ""
+
+    # Scanning options
+    services: List[str] = field(default_factory=list)  # Empty = all
+    skip_services: List[str] = field(default_factory=list)
+    max_workers: int = 25
+    no_browser: bool = True
+
+    # Output options
+    output_dir: str = "./scoutsuite_results"
+    report_name: str = ""
+    timestamp: bool = True
+
+
+@dataclass
+class ScoutSuiteResult:
+    """Result of a ScoutSuite scan."""
+    provider: str
+    status: str
+    started_at: str
+    finished_at: str
+    duration: float
+    findings_count: int
+    flagged_items: int
+    report_path: str
+    summary: Dict[str, int]
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class ScoutSuiteTool:
+    """
+    ScoutSuite wrapper for multi-cloud security auditing.
+
+    Provides a Python interface to the ScoutSuite CLI tool
+    for automated cloud security assessments.
+    """
+
+    def __init__(self, config: Optional[ScoutSuiteConfig] = None):
+        """
+        Initialize ScoutSuite tool.
+
+        Args:
+            config: ScoutSuite configuration
+        """
+        self.config = config or ScoutSuiteConfig()
+        self.output_dir = Path(self.config.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self._installed = None
+
+    async def check_installed(self) -> bool:
+        """Check if ScoutSuite is installed."""
+        if self._installed is not None:
+            return self._installed
+
+        try:
+            process = await asyncio.create_subprocess_shell(
+                "scout --version",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE
+            )
+            stdout, _ = await process.communicate()
+            self._installed = process.returncode == 0
+        except Exception:
+            self._installed = False
+
+        return self._installed
+
+    def _build_command(self) -> str:
+        """Build ScoutSuite command from configuration."""
+        cmd_parts = ["scout", self.config.provider]
+
+        # Add provider-specific options
+        if self.config.provider == "aws":
+            if self.config.aws_profile:
+                cmd_parts.extend(["--profile", self.config.aws_profile])
+            if self.config.aws_regions:
+                cmd_parts.extend(["--regions", ",".join(self.config.aws_regions)])
+
+        elif self.config.provider == "azure":
+            if self.config.azure_subscription_id:
+                cmd_parts.extend(["--subscription-id", self.config.azure_subscription_id])
+            if self.config.azure_tenant_id:
+                cmd_parts.extend(["--tenant-id", self.config.azure_tenant_id])
+
+        elif self.config.provider == "gcp":
+            if self.config.gcp_project_id:
+                cmd_parts.extend(["--project-id", self.config.gcp_project_id])
+            if self.config.gcp_service_account:
+                cmd_parts.extend(["--service-account", self.config.gcp_service_account])
+
+        # Add service filtering
+        if self.config.services:
+            cmd_parts.extend(["--services", ",".join(self.config.services)])
+        if self.config.skip_services:
+            cmd_parts.extend(["--skip", ",".join(self.config.skip_services)])
+
+        # Add general options
+        cmd_parts.extend(["--max-workers", str(self.config.max_workers)])
+        cmd_parts.extend(["--report-dir", str(self.output_dir)])
+
+        if self.config.report_name:
+            cmd_parts.extend(["--report-name", self.config.report_name])
+
+        if self.config.no_browser:
+            cmd_parts.append("--no-browser")
+
+        if self.config.timestamp:
+            cmd_parts.append("--timestamp")
+
+        return " ".join(cmd_parts)
+
+    async def scan(self, timeout: int = 3600) -> ScoutSuiteResult:
+        """
+        Run ScoutSuite scan.
+
+        Args:
+            timeout: Scan timeout in seconds
+
+        Returns:
+            ScoutSuiteResult with findings summary
+        """
+        if not await self.check_installed():
+            raise RuntimeError("ScoutSuite is not installed. Install with: pip install scoutsuite")
+
+        started_at = datetime.now(timezone.utc).isoformat()
+        start_time = asyncio.get_event_loop().time()
+
+        cmd = self._build_command()
+        print(f"[*] Running: {cmd}")
+
+        # Set up environment
+        env = os.environ.copy()
+
+        process = await asyncio.create_subprocess_shell(
+            cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=env
+        )
+
+        try:
+            stdout, stderr = await asyncio.wait_for(
+                process.communicate(),
+                timeout=timeout
+            )
+        except asyncio.TimeoutError:
+            process.kill()
+            raise TimeoutError(f"ScoutSuite scan timed out after {timeout}s")
+
+        finished_at = datetime.now(timezone.utc).isoformat()
+        duration = asyncio.get_event_loop().time() - start_time
+
+        # Parse results
+        report_path = self._find_latest_report()
+        findings_count = 0
+        flagged_items = 0
+        summary = {"danger": 0, "warning": 0, "info": 0}
+
+        if report_path and report_path.exists():
+            results = self._parse_results(report_path)
+            findings_count = results.get("findings_count", 0)
+            flagged_items = results.get("flagged_items", 0)
+            summary = results.get("summary", summary)
+
+        status = "completed" if process.returncode == 0 else "failed"
+
+        return ScoutSuiteResult(
+            provider=self.config.provider,
+            status=status,
+            started_at=started_at,
+            finished_at=finished_at,
+            duration=duration,
+            findings_count=findings_count,
+            flagged_items=flagged_items,
+            report_path=str(report_path) if report_path else "",
+            summary=summary,
+            metadata={
+                "command": cmd,
+                "return_code": process.returncode,
+                "stderr": stderr.decode() if process.returncode != 0 else ""
+            }
+        )
+
+    def _find_latest_report(self) -> Optional[Path]:
+        """Find the latest ScoutSuite results file."""
+        results_dir = self.output_dir / "scoutsuite-results"
+        if results_dir.exists():
+            js_files = list(results_dir.glob("scoutsuite_results*.js"))
+            if js_files:
+                return max(js_files, key=lambda f: f.stat().st_mtime)
+
+        # Try alternative location
+        js_files = list(self.output_dir.glob("**/scoutsuite_results*.js"))
+        if js_files:
+            return max(js_files, key=lambda f: f.stat().st_mtime)
+
+        return None
+
+    def _parse_results(self, results_file: Path) -> Dict[str, Any]:
+        """Parse ScoutSuite results JavaScript file."""
+        try:
+            content = results_file.read_text()
+
+            # Remove JS variable assignment
+            if "scoutsuite_results =" in content:
+                content = content.split("scoutsuite_results =", 1)[1].strip()
+                if content.endswith(";"):
+                    content = content[:-1]
+
+            data = json.loads(content)
+
+            # Count findings
+            findings_count = 0
+            flagged_items = 0
+            summary = {"danger": 0, "warning": 0, "info": 0}
+
+            services = data.get("services", {})
+            for service_data in services.values():
+                service_findings = service_data.get("findings", {})
+                for finding in service_findings.values():
+                    findings_count += 1
+                    flagged = finding.get("flagged_items", 0)
+                    flagged_items += flagged
+
+                    if flagged > 0:
+                        level = finding.get("level", "warning")
+                        if level in summary:
+                            summary[level] += flagged
+
+            return {
+                "findings_count": findings_count,
+                "flagged_items": flagged_items,
+                "summary": summary,
+                "account_id": data.get("account_id", ""),
+                "last_run": data.get("last_run", {})
+            }
+
+        except Exception as e:
+            print(f"[!] Error parsing ScoutSuite results: {e}")
+            return {"findings_count": 0, "flagged_items": 0, "summary": {}}
+
+    def get_findings(self) -> List[Dict[str, Any]]:
+        """Get parsed findings from the latest scan."""
+        report_path = self._find_latest_report()
+        if not report_path:
+            return []
+
+        try:
+            content = report_path.read_text()
+            if "scoutsuite_results =" in content:
+                content = content.split("scoutsuite_results =", 1)[1].strip()
+                if content.endswith(";"):
+                    content = content[:-1]
+
+            data = json.loads(content)
+            findings = []
+
+            services = data.get("services", {})
+            for service_name, service_data in services.items():
+                for finding_id, finding_data in service_data.get("findings", {}).items():
+                    if finding_data.get("flagged_items", 0) > 0:
+                        findings.append({
+                            "provider": self.config.provider,
+                            "service": service_name,
+                            "id": finding_id,
+                            "level": finding_data.get("level", "warning"),
+                            "description": finding_data.get("description", ""),
+                            "rationale": finding_data.get("rationale", ""),
+                            "remediation": finding_data.get("remediation", ""),
+                            "flagged_items": finding_data.get("flagged_items", 0),
+                            "items": finding_data.get("items", []),
+                            "compliance": finding_data.get("compliance", [])
+                        })
+
+            return findings
+
+        except Exception as e:
+            print(f"[!] Error getting findings: {e}")
+            return []
+
+
+# Convenience function
+async def run_scoutsuite(
+    provider: str = "aws",
+    profile: Optional[str] = None,
+    regions: Optional[List[str]] = None,
+    services: Optional[List[str]] = None,
+    output_dir: str = "./scoutsuite_results",
+    timeout: int = 3600
+) -> ScoutSuiteResult:
+    """
+    Run ScoutSuite scan.
+
+    Args:
+        provider: Cloud provider (aws, azure, gcp)
+        profile: AWS profile name (for AWS)
+        regions: Specific regions to scan
+        services: Specific services to scan
+        output_dir: Output directory for reports
+        timeout: Scan timeout in seconds
+
+    Returns:
+        ScoutSuiteResult
+    """
+    config = ScoutSuiteConfig(
+        provider=provider,
+        aws_profile=profile or "default",
+        aws_regions=regions or [],
+        services=services or [],
+        output_dir=output_dir
+    )
+
+    tool = ScoutSuiteTool(config)
+    return await tool.scan(timeout=timeout)

aipt_v2/tools/executor.py

@@ -0,0 +1,307 @@
+from __future__ import annotations
+
+import inspect
+import os
+from typing import Any
+
+import httpx
+
+
+if os.getenv("AIPT_SANDBOX_MODE", "false").lower() == "false":
+    from aipt_v2.runtime import get_runtime
+
+from .argument_parser import convert_arguments
+from .registry import (
+    get_tool_by_name,
+    get_tool_names,
+    needs_agent_state,
+    should_execute_in_sandbox,
+)
+
+
+async def execute_tool(tool_name: str, agent_state: Any | None = None, **kwargs: Any) -> Any:
+    execute_in_sandbox = should_execute_in_sandbox(tool_name)
+    sandbox_mode = os.getenv("AIPT_SANDBOX_MODE", "false").lower() == "true"
+
+    if execute_in_sandbox and not sandbox_mode:
+        return await _execute_tool_in_sandbox(tool_name, agent_state, **kwargs)
+
+    return await _execute_tool_locally(tool_name, agent_state, **kwargs)
+
+
+async def _execute_tool_in_sandbox(tool_name: str, agent_state: Any, **kwargs: Any) -> Any:
+    if not hasattr(agent_state, "sandbox_id") or not agent_state.sandbox_id:
+        raise ValueError("Agent state with a valid sandbox_id is required for sandbox execution.")
+
+    if not hasattr(agent_state, "sandbox_token") or not agent_state.sandbox_token:
+        raise ValueError(
+            "Agent state with a valid sandbox_token is required for sandbox execution."
+        )
+
+    if (
+        not hasattr(agent_state, "sandbox_info")
+        or "tool_server_port" not in agent_state.sandbox_info
+    ):
+        raise ValueError(
+            "Agent state with a valid sandbox_info containing tool_server_port is required."
+        )
+
+    runtime = get_runtime()
+    tool_server_port = agent_state.sandbox_info["tool_server_port"]
+    server_url = await runtime.get_sandbox_url(agent_state.sandbox_id, tool_server_port)
+    request_url = f"{server_url}/execute"
+
+    agent_id = getattr(agent_state, "agent_id", "unknown")
+
+    request_data = {
+        "agent_id": agent_id,
+        "tool_name": tool_name,
+        "kwargs": kwargs,
+    }
+
+    headers = {
+        "Authorization": f"Bearer {agent_state.sandbox_token}",
+        "Content-Type": "application/json",
+    }
+
+    async with httpx.AsyncClient(trust_env=False) as client:
+        try:
+            response = await client.post(
+                request_url, json=request_data, headers=headers, timeout=None
+            )
+            response.raise_for_status()
+            response_data = response.json()
+            if response_data.get("error"):
+                raise RuntimeError(f"Sandbox execution error: {response_data['error']}")
+            return response_data.get("result")
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 401:
+                raise RuntimeError("Authentication failed: Invalid or missing sandbox token") from e
+            raise RuntimeError(f"HTTP error calling tool server: {e.response.status_code}") from e
+        except httpx.RequestError as e:
+            raise RuntimeError(f"Request error calling tool server: {e}") from e
+
+
+async def _execute_tool_locally(tool_name: str, agent_state: Any | None, **kwargs: Any) -> Any:
+    tool_func = get_tool_by_name(tool_name)
+    if not tool_func:
+        raise ValueError(f"Tool '{tool_name}' not found")
+
+    converted_kwargs = convert_arguments(tool_func, kwargs)
+
+    if needs_agent_state(tool_name):
+        if agent_state is None:
+            raise ValueError(f"Tool '{tool_name}' requires agent_state but none was provided.")
+        result = tool_func(agent_state=agent_state, **converted_kwargs)
+    else:
+        result = tool_func(**converted_kwargs)
+
+    return await result if inspect.isawaitable(result) else result
+
+
+def validate_tool_availability(tool_name: str | None) -> tuple[bool, str]:
+    if tool_name is None:
+        return False, "Tool name is missing"
+
+    if tool_name not in get_tool_names():
+        return False, f"Tool '{tool_name}' is not available"
+
+    return True, ""
+
+
+async def execute_tool_with_validation(
+    tool_name: str | None, agent_state: Any | None = None, **kwargs: Any
+) -> Any:
+    is_valid, error_msg = validate_tool_availability(tool_name)
+    if not is_valid:
+        return f"Error: {error_msg}"
+
+    assert tool_name is not None
+
+    try:
+        result = await execute_tool(tool_name, agent_state, **kwargs)
+    except Exception as e:  # noqa: BLE001
+        error_str = str(e)
+        if len(error_str) > 500:
+            error_str = error_str[:500] + "... [truncated]"
+        return f"Error executing {tool_name}: {error_str}"
+    else:
+        return result
+
+
+async def execute_tool_invocation(tool_inv: dict[str, Any], agent_state: Any | None = None) -> Any:
+    tool_name = tool_inv.get("toolName")
+    tool_args = tool_inv.get("args", {})
+
+    return await execute_tool_with_validation(tool_name, agent_state, **tool_args)
+
+
+def _check_error_result(result: Any) -> tuple[bool, Any]:
+    is_error = False
+    error_payload: Any = None
+
+    if (isinstance(result, dict) and "error" in result) or (
+        isinstance(result, str) and result.strip().lower().startswith("error:")
+    ):
+        is_error = True
+        error_payload = result
+
+    return is_error, error_payload
+
+
+def _update_tracer_with_result(
+    tracer: Any, execution_id: Any, is_error: bool, result: Any, error_payload: Any
+) -> None:
+    if not tracer or not execution_id:
+        return
+
+    try:
+        if is_error:
+            tracer.update_tool_execution(execution_id, "error", error_payload)
+        else:
+            tracer.update_tool_execution(execution_id, "completed", result)
+    except (ConnectionError, RuntimeError) as e:
+        error_msg = str(e)
+        if tracer and execution_id:
+            tracer.update_tool_execution(execution_id, "error", error_msg)
+        raise
+
+
+def _format_tool_result(tool_name: str, result: Any) -> tuple[str, list[dict[str, Any]]]:
+    images: list[dict[str, Any]] = []
+
+    screenshot_data = extract_screenshot_from_result(result)
+    if screenshot_data:
+        images.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/png;base64,{screenshot_data}"},
+            }
+        )
+        result_str = remove_screenshot_from_result(result)
+    else:
+        result_str = result
+
+    if result_str is None:
+        final_result_str = f"Tool {tool_name} executed successfully"
+    else:
+        final_result_str = str(result_str)
+        if len(final_result_str) > 10000:
+            start_part = final_result_str[:4000]
+            end_part = final_result_str[-4000:]
+            final_result_str = start_part + "\n\n... [middle content truncated] ...\n\n" + end_part
+
+    observation_xml = (
+        f"<tool_result>\n<tool_name>{tool_name}</tool_name>\n"
+        f"<result>{final_result_str}</result>\n</tool_result>"
+    )
+
+    return observation_xml, images
+
+
+async def _execute_single_tool(
+    tool_inv: dict[str, Any],
+    agent_state: Any | None,
+    tracer: Any | None,
+    agent_id: str,
+) -> tuple[str, list[dict[str, Any]], bool]:
+    tool_name = tool_inv.get("toolName", "unknown")
+    args = tool_inv.get("args", {})
+    execution_id = None
+    should_agent_finish = False
+
+    if tracer:
+        execution_id = tracer.log_tool_execution_start(agent_id, tool_name, args)
+
+    try:
+        result = await execute_tool_invocation(tool_inv, agent_state)
+
+        is_error, error_payload = _check_error_result(result)
+
+        if (
+            tool_name in ("finish_scan", "agent_finish")
+            and not is_error
+            and isinstance(result, dict)
+        ):
+            if tool_name == "finish_scan":
+                should_agent_finish = result.get("scan_completed", False)
+            elif tool_name == "agent_finish":
+                should_agent_finish = result.get("agent_completed", False)
+
+        _update_tracer_with_result(tracer, execution_id, is_error, result, error_payload)
+
+    except (ConnectionError, RuntimeError, ValueError, TypeError, OSError) as e:
+        error_msg = str(e)
+        if tracer and execution_id:
+            tracer.update_tool_execution(execution_id, "error", error_msg)
+        raise
+
+    observation_xml, images = _format_tool_result(tool_name, result)
+    return observation_xml, images, should_agent_finish
+
+
+def _get_tracer_and_agent_id(agent_state: Any | None) -> tuple[Any | None, str]:
+    try:
+        from aipt_v2.telemetry.tracer import get_global_tracer
+
+        tracer = get_global_tracer()
+        agent_id = agent_state.agent_id if agent_state else "unknown_agent"
+    except (ImportError, AttributeError):
+        tracer = None
+        agent_id = "unknown_agent"
+
+    return tracer, agent_id
+
+
+async def process_tool_invocations(
+    tool_invocations: list[dict[str, Any]],
+    conversation_history: list[dict[str, Any]],
+    agent_state: Any | None = None,
+) -> bool:
+    observation_parts: list[str] = []
+    all_images: list[dict[str, Any]] = []
+    should_agent_finish = False
+
+    tracer, agent_id = _get_tracer_and_agent_id(agent_state)
+
+    for tool_inv in tool_invocations:
+        observation_xml, images, tool_should_finish = await _execute_single_tool(
+            tool_inv, agent_state, tracer, agent_id
+        )
+        observation_parts.append(observation_xml)
+        all_images.extend(images)
+
+        if tool_should_finish:
+            should_agent_finish = True
+
+    if all_images:
+        content = [{"type": "text", "text": "Tool Results:\n\n" + "\n\n".join(observation_parts)}]
+        content.extend(all_images)
+        conversation_history.append({"role": "user", "content": content})
+    else:
+        observation_content = "Tool Results:\n\n" + "\n\n".join(observation_parts)
+        conversation_history.append({"role": "user", "content": observation_content})
+
+    return should_agent_finish
+
+
+def extract_screenshot_from_result(result: Any) -> str | None:
+    if not isinstance(result, dict):
+        return None
+
+    screenshot = result.get("screenshot")
+    if isinstance(screenshot, str) and screenshot:
+        return screenshot
+
+    return None
+
+
+def remove_screenshot_from_result(result: Any) -> Any:
+    if not isinstance(result, dict):
+        return result
+
+    result_copy = result.copy()
+    if "screenshot" in result_copy:
+        result_copy["screenshot"] = "[Image data extracted - see attached image]"
+
+    return result_copy