ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/tools/sandboxtool.py

@@ -0,0 +1,703 @@
+"""
+AI-Parrot gVisor Sandbox Tool
+
+Secure Python execution tool using gVisor (runsc) for complete kernel-level isolation.
+This tool provides safe code execution for untrusted LLM-generated code.
+"""
+
+import asyncio
+import json
+import os
+import tempfile
+import uuid
+import shutil
+import subprocess
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, List, Tuple, Union
+from dataclasses import dataclass, field
+from contextlib import contextmanager
+import base64
+import pickle
+import pandas as pd
+import numpy as np
+
+from parrot.tools.abstract import AbstractTool, ToolResult
+from parrot.tools.pythonpandas import PythonPandasTool
+
+
+@dataclass
+class SandboxConfig:
+    """Configuration for gVisor sandbox"""
+    runtime: str = "runsc"
+    network: str = "none"  # Disable network by default
+    max_memory: str = "2G"
+    max_cpu: float = 2.0
+    timeout: int = 30
+    python_image: str = "python:3.11-slim"
+    enable_gpu: bool = False
+    mount_paths: List[str] = field(default_factory=list)
+    pip_packages: List[str] = field(default_factory=lambda: [
+        "pandas", "numpy", "matplotlib", "seaborn", "scikit-learn",
+        "scipy", "plotly", "jupyterlab", "ipykernel"
+    ])
+
+
+@dataclass
+class ExecutionResult:
+    """Result from sandbox execution"""
+    success: bool
+    stdout: str
+    stderr: str
+    exit_code: int
+    execution_time_ms: float
+    files_created: Dict[str, bytes] = field(default_factory=dict)
+    dataframes: Dict[str, pd.DataFrame] = field(default_factory=dict)
+    plots: List[Dict[str, Any]] = field(default_factory=list)
+    notebook_path: Optional[str] = None
+    error_message: Optional[str] = None
+
+
+class SandboxTool(AbstractTool):
+    """
+    Secure Python execution using gVisor sandbox.
+    
+    This tool provides kernel-level isolation for executing untrusted code
+    generated by LLMs, preventing escape attempts and resource exhaustion.
+    """
+    
+    name = "gvisor_python_sandbox"
+    description = "Execute Python code in a secure gVisor sandbox with kernel isolation"
+    
+    def __init__(
+        self,
+        config: Optional[SandboxConfig] = None,
+        workspace_dir: Optional[str] = None,
+        dataframes: Optional[Dict[str, pd.DataFrame]] = None,
+        enable_jupyter: bool = False,
+        logger: Optional[logging.Logger] = None
+    ):
+        """
+        Initialize gVisor sandbox tool.
+        
+        Args:
+            config: Sandbox configuration
+            workspace_dir: Directory for temporary files
+            dataframes: Initial dataframes to make available
+            enable_jupyter: Enable Jupyter notebook creation
+            logger: Logger instance
+        """
+        super().__init__()
+        self.config = config or SandboxConfig()
+        self.workspace_dir = Path(workspace_dir or tempfile.mkdtemp(prefix="gvisor_"))
+        self.workspace_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.dataframes = dataframes or {}
+        self.enable_jupyter = enable_jupyter
+        self.logger = logger or logging.getLogger(__name__)
+        
+        # Verify gVisor installation
+        self._verify_installation()
+        
+        # Prepare base container image
+        self._prepare_base_image()
+    
+    def _verify_installation(self):
+        """Verify gVisor and containerd are installed"""
+        try:
+            # Check runsc
+            result = subprocess.run(
+                ["runsc", "--version"],
+                capture_output=True,
+                text=True
+            )
+            if result.returncode != 0:
+                raise RuntimeError("runsc not found. Please install gVisor.")
+            
+            self.logger.info(f"gVisor version: {result.stdout.strip()}")
+            
+            # Check containerd
+            result = subprocess.run(
+                ["ctr", "--version"],
+                capture_output=True,
+                text=True
+            )
+            if result.returncode != 0:
+                self.logger.warning("containerd not found. Using Docker as fallback.")
+                self.use_docker = True
+            else:
+                self.use_docker = False
+                
+        except FileNotFoundError as e:
+            raise RuntimeError(f"Required tools not installed: {e}")
+    
+    def _prepare_base_image(self):
+        """Prepare base container image with required packages"""
+        dockerfile_content = f"""
+FROM {self.config.python_image}
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \\
+    gcc \\
+    g++ \\
+    make \\
+    libssl-dev \\
+    libffi-dev \\
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python packages
+RUN pip install --no-cache-dir \\
+    {' '.join(self.config.pip_packages)}
+
+# Create working directory
+WORKDIR /workspace
+
+# Create sandbox user (non-root)
+RUN useradd -m -s /bin/bash sandbox && \\
+    chown -R sandbox:sandbox /workspace
+
+USER sandbox
+
+# Set Python to unbuffered mode
+ENV PYTHONUNBUFFERED=1
+ENV PYTHONDONTWRITEBYTECODE=1
+"""
+        
+        # Build image
+        image_name = "ai-parrot-gvisor-sandbox"
+        dockerfile_path = self.workspace_dir / "Dockerfile"
+        dockerfile_path.write_text(dockerfile_content)
+        
+        try:
+            subprocess.run(
+                ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), "."],
+                cwd=self.workspace_dir,
+                check=True,
+                capture_output=True
+            )
+            self.image_name = image_name
+            self.logger.info(f"Built sandbox image: {image_name}")
+        except subprocess.CalledProcessError as e:
+            self.logger.error(f"Failed to build image: {e.stderr.decode()}")
+            raise
+    
+    @contextmanager
+    def _create_sandbox_container(self, code: str, session_id: str):
+        """Create and manage a gVisor sandbox container"""
+        container_name = f"sandbox-{session_id}"
+        exec_dir = self.workspace_dir / session_id
+        exec_dir.mkdir(parents=True, exist_ok=True)
+        
+        try:
+            # Prepare execution script
+            script_path = exec_dir / "execute.py"
+            script_path.write_text(self._prepare_execution_script(code))
+            
+            # Prepare data files
+            self._prepare_data_files(exec_dir)
+            
+            # Create container with gVisor runtime
+            run_cmd = [
+                "docker", "run",
+                "--runtime", self.config.runtime,  # Use gVisor runtime
+                "--name", container_name,
+                "-d",  # Detached
+                "--rm",  # Auto-remove
+                "--memory", self.config.max_memory,
+                "--cpus", str(self.config.max_cpu),
+                "--network", self.config.network,
+                "--security-opt", "no-new-privileges",
+                "--cap-drop", "ALL",  # Drop all capabilities
+                "-v", f"{exec_dir}:/workspace:ro",  # Read-only mount
+                "-v", f"{exec_dir}/output:/output:rw",  # Output directory
+                self.image_name,
+                "python", "/workspace/execute.py"
+            ]
+            
+            # Start container
+            subprocess.run(run_cmd, check=True, capture_output=True)
+            
+            yield container_name
+            
+        finally:
+            # Cleanup container
+            try:
+                subprocess.run(
+                    ["docker", "stop", container_name],
+                    capture_output=True,
+                    timeout=5
+                )
+            except:
+                subprocess.run(
+                    ["docker", "kill", container_name],
+                    capture_output=True
+                )
+    
+    def _prepare_execution_script(self, code: str) -> str:
+        """Prepare the Python script for execution in sandbox"""
+        return f'''
+import sys
+import os
+import json
+import pickle
+import base64
+import traceback
+import pandas as pd
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')  # Non-interactive backend
+import matplotlib.pyplot as plt
+from datetime import datetime
+
+# Load serialized dataframes
+dataframes_file = '/workspace/dataframes.pkl'
+if os.path.exists(dataframes_file):
+    with open(dataframes_file, 'rb') as f:
+        dataframes = pickle.load(f)
+        # Inject dataframes into global namespace
+        for name, df in dataframes.items():
+            globals()[name] = df
+            print(f"Loaded dataframe: {{name}} with shape {{df.shape}}")
+
+# Execution metadata
+execution_start = datetime.now()
+created_files = {{}}
+output_dataframes = {{}}
+plots = []
+
+# Redirect stdout to capture prints
+import io
+from contextlib import redirect_stdout, redirect_stderr
+
+stdout_capture = io.StringIO()
+stderr_capture = io.StringIO()
+
+# Custom figure save wrapper
+original_savefig = plt.savefig
+def wrapped_savefig(fname, **kwargs):
+    fname = f"/output/{{fname}}"
+    original_savefig(fname, **kwargs)
+    with open(fname, 'rb') as f:
+        created_files[os.path.basename(fname)] = base64.b64encode(f.read()).decode()
+    plots.append({{"type": "matplotlib", "filename": os.path.basename(fname)}})
+    return fname
+plt.savefig = wrapped_savefig
+
+try:
+    with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+        # Execute user code
+        exec("""
+{code}
+""")
+        
+        # Capture any created dataframes
+        for name, obj in list(globals().items()):
+            if isinstance(obj, pd.DataFrame) and name not in ['dataframes', 'pd']:
+                if name not in dataframes or not obj.equals(dataframes.get(name, pd.DataFrame())):
+                    output_dataframes[name] = obj
+    
+    execution_time = (datetime.now() - execution_start).total_seconds() * 1000
+    
+    # Save results
+    result = {{
+        "success": True,
+        "stdout": stdout_capture.getvalue(),
+        "stderr": stderr_capture.getvalue(),
+        "exit_code": 0,
+        "execution_time_ms": execution_time,
+        "files_created": created_files,
+        "dataframes": {{name: df.to_dict('records') for name, df in output_dataframes.items()}},
+        "plots": plots
+    }}
+    
+except Exception as e:
+    result = {{
+        "success": False,
+        "stdout": stdout_capture.getvalue(),
+        "stderr": stderr_capture.getvalue() + "\\n" + traceback.format_exc(),
+        "exit_code": 1,
+        "execution_time_ms": (datetime.now() - execution_start).total_seconds() * 1000,
+        "error_message": str(e)
+    }}
+
+# Write result
+with open('/output/result.json', 'w') as f:
+    json.dump(result, f)
+'''
+    
+    def _prepare_data_files(self, exec_dir: Path):
+        """Prepare data files for the sandbox"""
+        # Create output directory
+        output_dir = exec_dir / "output"
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Serialize dataframes
+        if self.dataframes:
+            with open(exec_dir / "dataframes.pkl", 'wb') as f:
+                pickle.dump(self.dataframes, f)
+    
+    async def execute(
+        self,
+        code: str,
+        description: str = "Execute Python code in gVisor sandbox",
+        create_notebook: bool = False,
+        **kwargs
+    ) -> ToolResult:
+        """
+        Execute Python code in gVisor sandbox.
+        
+        Args:
+            code: Python code to execute
+            description: Description of execution
+            create_notebook: Create Jupyter notebook of execution
+            **kwargs: Additional parameters
+            
+        Returns:
+            ToolResult with execution results
+        """
+        session_id = str(uuid.uuid4())[:8]
+        self.logger.info(f"Starting sandbox execution: {session_id}")
+        
+        try:
+            # Execute in sandbox
+            result = await self._execute_in_sandbox(code, session_id)
+            
+            # Optionally create notebook
+            if create_notebook and self.enable_jupyter:
+                notebook_path = await self._create_notebook(code, result, session_id)
+                result.notebook_path = notebook_path
+            
+            # Format response
+            response = self._format_response(result)
+            
+            return ToolResult(
+                result=response,
+                status="success" if result.success else "error",
+                error=result.error_message,
+                metadata={
+                    "session_id": session_id,
+                    "execution_time_ms": result.execution_time_ms,
+                    "files_created": list(result.files_created.keys()),
+                    "dataframes_created": list(result.dataframes.keys()),
+                    "notebook_path": result.notebook_path
+                }
+            )
+            
+        except Exception as e:
+            self.logger.error(f"Sandbox execution failed: {e}")
+            return ToolResult(
+                result=None,
+                status="error",
+                error=str(e)
+            )
+        finally:
+            # Cleanup session files
+            if not kwargs.get("keep_files", False):
+                self._cleanup_session(session_id)
+    
+    async def _execute_in_sandbox(self, code: str, session_id: str) -> ExecutionResult:
+        """Execute code in gVisor sandbox container"""
+        exec_dir = self.workspace_dir / session_id
+        
+        with self._create_sandbox_container(code, session_id) as container_name:
+            # Wait for execution with timeout
+            try:
+                result = subprocess.run(
+                    ["docker", "wait", container_name],
+                    capture_output=True,
+                    text=True,
+                    timeout=self.config.timeout
+                )
+                
+                # Get logs
+                logs = subprocess.run(
+                    ["docker", "logs", container_name],
+                    capture_output=True,
+                    text=True
+                )
+                
+            except subprocess.TimeoutExpired:
+                # Kill container on timeout
+                subprocess.run(["docker", "kill", container_name], capture_output=True)
+                return ExecutionResult(
+                    success=False,
+                    stdout="",
+                    stderr="Execution timeout exceeded",
+                    exit_code=-1,
+                    execution_time_ms=self.config.timeout * 1000,
+                    error_message=f"Code execution exceeded {self.config.timeout}s timeout"
+                )
+        
+        # Read results
+        result_file = exec_dir / "output" / "result.json"
+        if result_file.exists():
+            with open(result_file, 'r') as f:
+                result_data = json.load(f)
+            
+            # Decode file contents
+            files_created = {}
+            for filename, content in result_data.get("files_created", {}).items():
+                files_created[filename] = base64.b64decode(content)
+            
+            # Reconstruct dataframes
+            dataframes = {}
+            for name, records in result_data.get("dataframes", {}).items():
+                dataframes[name] = pd.DataFrame(records)
+            
+            return ExecutionResult(
+                success=result_data["success"],
+                stdout=result_data["stdout"],
+                stderr=result_data["stderr"],
+                exit_code=result_data["exit_code"],
+                execution_time_ms=result_data["execution_time_ms"],
+                files_created=files_created,
+                dataframes=dataframes,
+                plots=result_data.get("plots", []),
+                error_message=result_data.get("error_message")
+            )
+        else:
+            return ExecutionResult(
+                success=False,
+                stdout=logs.stdout,
+                stderr=logs.stderr,
+                exit_code=result.returncode,
+                execution_time_ms=0,
+                error_message="Failed to read execution results"
+            )
+    
+    async def _create_notebook(
+        self,
+        code: str,
+        result: ExecutionResult,
+        session_id: str
+    ) -> str:
+        """Create Jupyter notebook from execution"""
+        import nbformat
+        from nbformat.v4 import new_notebook, new_code_cell, new_markdown_cell
+        
+        # Create notebook
+        nb = new_notebook()
+        
+        # Add header cell
+        nb.cells.append(new_markdown_cell(
+            f"# gVisor Sandbox Execution\n"
+            f"Session: {session_id}\n"
+            f"Execution Time: {result.execution_time_ms:.2f}ms"
+        ))
+        
+        # Add code cell
+        code_cell = new_code_cell(code)
+        if result.stdout:
+            code_cell.outputs = [{
+                "output_type": "stream",
+                "name": "stdout",
+                "text": result.stdout
+            }]
+        nb.cells.append(code_cell)
+        
+        # Add results cell
+        if result.dataframes:
+            nb.cells.append(new_markdown_cell("## Created DataFrames"))
+            for name, df in result.dataframes.items():
+                nb.cells.append(new_code_cell(f"# DataFrame: {name}\n{df.to_string()}"))
+        
+        if result.plots:
+            nb.cells.append(new_markdown_cell("## Generated Plots"))
+            for plot in result.plots:
+                nb.cells.append(new_markdown_cell(f"![{plot['filename']}]({plot['filename']})"))
+        
+        # Save notebook
+        notebook_path = self.workspace_dir / f"notebook_{session_id}.ipynb"
+        with open(notebook_path, 'w') as f:
+            nbformat.write(nb, f)
+        
+        return str(notebook_path)
+    
+    def _format_response(self, result: ExecutionResult) -> str:
+        """Format execution result for display"""
+        parts = []
+        
+        if result.success:
+            parts.append("✓ Code executed successfully")
+            parts.append(f"Execution time: {result.execution_time_ms:.2f}ms")
+        else:
+            parts.append("✗ Code execution failed")
+            parts.append(f"Error: {result.error_message}")
+        
+        if result.stdout:
+            parts.append("\n--- Output ---")
+            parts.append(result.stdout)
+        
+        if result.stderr and not result.success:
+            parts.append("\n--- Errors ---")
+            parts.append(result.stderr)
+        
+        if result.files_created:
+            parts.append(f"\n--- Files Created ({len(result.files_created)}) ---")
+            for filename in result.files_created.keys():
+                parts.append(f"  • {filename}")
+        
+        if result.dataframes:
+            parts.append(f"\n--- DataFrames Created ({len(result.dataframes)}) ---")
+            for name, df in result.dataframes.items():
+                parts.append(f"  • {name}: {df.shape[0]} rows × {df.shape[1]} columns")
+        
+        if result.notebook_path:
+            parts.append(f"\n--- Notebook ---")
+            parts.append(f"  📓 {result.notebook_path}")
+        
+        return "\n".join(parts)
+    
+    def _cleanup_session(self, session_id: str):
+        """Clean up session files"""
+        session_dir = self.workspace_dir / session_id
+        if session_dir.exists():
+            shutil.rmtree(session_dir)
+    
+    def add_dataframe(self, name: str, df: pd.DataFrame):
+        """Add a dataframe to be available in the sandbox"""
+        self.dataframes[name] = df
+    
+    def set_timeout(self, timeout: int):
+        """Set execution timeout in seconds"""
+        self.config.timeout = timeout
+    
+    def set_memory_limit(self, memory: str):
+        """Set memory limit (e.g., '1G', '512M')"""
+        self.config.max_memory = memory
+    
+    def cleanup(self):
+        """Clean up all resources"""
+        if self.workspace_dir.exists():
+            shutil.rmtree(self.workspace_dir)
+
+
+class SandboxPandasTool(SandboxTool):
+    """
+    Specialized version for Pandas operations with enhanced data handling.
+    Drop-in replacement for PythonPandasTool with security.
+    """
+    
+    name = "sandbox_pandas_tool"
+    description = "Secure Pandas operations in gVisor sandbox"
+    
+    def __init__(self, dataframes: Optional[Dict[str, pd.DataFrame]] = None, **kwargs):
+        """Initialize with Pandas-specific configuration"""
+        config = SandboxConfig(
+            pip_packages=[
+                "pandas", "numpy", "matplotlib", "seaborn", 
+                "plotly", "scikit-learn", "scipy", "statsmodels",
+                "openpyxl", "xlrd", "pyarrow", "fastparquet"
+            ]
+        )
+        super().__init__(config=config, dataframes=dataframes, **kwargs)
+    
+    async def analyze_dataframe(
+        self,
+        df_name: str,
+        analysis_type: str = "summary"
+    ) -> ToolResult:
+        """
+        Perform automated DataFrame analysis.
+        
+        Args:
+            df_name: Name of dataframe to analyze
+            analysis_type: Type of analysis (summary, correlation, distribution)
+        """
+        if df_name not in self.dataframes:
+            return ToolResult(
+                result=None,
+                status="error", 
+                error=f"DataFrame '{df_name}' not found"
+            )
+        
+        analysis_code = {
+            "summary": f"""
+import pandas as pd
+df = {df_name}
+print("=== DataFrame Summary ===")
+print(f"Shape: {{df.shape}}")
+print(f"\\nData Types:\\n{{df.dtypes}}")
+print(f"\\nSummary Statistics:\\n{{df.describe()}}")
+print(f"\\nMissing Values:\\n{{df.isnull().sum()}}")
+print(f"\\nMemory Usage:\\n{{df.memory_usage(deep=True)}}")
+            """,
+            
+            "correlation": f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+df = {df_name}
+numeric_cols = df.select_dtypes(include=['number']).columns
+if len(numeric_cols) > 1:
+    corr_matrix = df[numeric_cols].corr()
+    print("=== Correlation Matrix ===")
+    print(corr_matrix)
+    
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0)
+    plt.title('Correlation Heatmap')
+    plt.tight_layout()
+    plt.savefig('correlation_heatmap.png')
+else:
+    print("Not enough numeric columns for correlation analysis")
+            """,
+            
+            "distribution": f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+df = {df_name}
+numeric_cols = df.select_dtypes(include=['number']).columns
+
+for i, col in enumerate(numeric_cols[:6]):  # Limit to first 6 columns
+    plt.figure(figsize=(12, 4))
+    
+    plt.subplot(1, 3, 1)
+    df[col].hist(bins=30, edgecolor='black')
+    plt.title(f'{{col}} - Histogram')
+    
+    plt.subplot(1, 3, 2)
+    df.boxplot(column=col)
+    plt.title(f'{{col}} - Boxplot')
+    
+    plt.subplot(1, 3, 3)
+    df[col].plot(kind='kde')
+    plt.title(f'{{col}} - KDE')
+    
+    plt.tight_layout()
+    plt.savefig(f'distribution_{{col}}.png')
+    plt.close()
+    
+print(f"Generated distribution plots for {{len(numeric_cols)}} numeric columns")
+            """
+        }
+        
+        code = analysis_code.get(analysis_type, analysis_code["summary"])
+        return await self.execute(
+            code,
+            description=f"Analyze DataFrame '{df_name}' - {analysis_type}"
+        )
+
+
+# Factory function for easy instantiation
+def create_sandbox_tool(
+    tool_type: str = "sandbox",
+    **kwargs
+) -> Union[SandboxTool, SandboxPandasTool]:
+    """
+    Factory function to create gVisor tools.
+    
+    Args:
+        tool_type: Type of tool ('sandbox' or 'pandas')
+        **kwargs: Configuration parameters
+    
+    Returns:
+        Configured gVisor tool instance
+    """
+    if tool_type == "pandas":
+        return SandboxPandasTool(**kwargs)
+    else:
+        return SandboxTool(**kwargs)