astraagent 2.25.6

package/astra/tools/python.py

@@ -0,0 +1,238 @@
+"""
+Python Code Execution Tool for AstraAgent.
+Safely execute Python code in an isolated environment.
+"""
+
+import ast
+import sys
+import io
+import traceback
+import asyncio
+from typing import List, Dict, Any, Optional
+from contextlib import redirect_stdout, redirect_stderr
+
+from astra.tools.base import Tool, ToolParameter, ToolResult, ToolCategory, RiskLevel
+
+
+class PythonTool(Tool):
+    """Tool for executing Python code."""
+    
+    name = "python"
+    description = "Execute Python code and return the result. Supports data analysis, calculations, and scripting."
+    category = ToolCategory.CODE_EXECUTION
+    risk_level = RiskLevel.MEDIUM
+    
+    # Dangerous imports/operations to block
+    BLOCKED_IMPORTS = [
+        'subprocess', 'os.system', 'os.popen', 'os.spawn',
+        'pty', 'pexpect', 'paramiko', 'fabric'
+    ]
+    
+    BLOCKED_BUILTINS = ['eval', 'exec', 'compile', '__import__']
+    
+    def __init__(self, 
+                 timeout: int = 30,
+                 max_memory_mb: int = 512,
+                 allow_imports: List[str] = None,
+                 sandbox_mode: bool = True):
+        super().__init__()
+        self.timeout = timeout
+        self.max_memory_mb = max_memory_mb
+        self.allow_imports = allow_imports
+        self.sandbox_mode = sandbox_mode
+        
+        # Persistent namespace for REPL-like behavior
+        self._namespace: Dict[str, Any] = {}
+    
+    @property
+    def parameters(self) -> List[ToolParameter]:
+        return [
+            ToolParameter(
+                name="code",
+                description="Python code to execute",
+                type="string",
+                required=True
+            ),
+            ToolParameter(
+                name="timeout",
+                description="Execution timeout in seconds",
+                type="integer",
+                required=False,
+                default=30
+            ),
+            ToolParameter(
+                name="reset_namespace",
+                description="Clear the execution namespace before running",
+                type="boolean",
+                required=False,
+                default=False
+            )
+        ]
+    
+    def _validate_code(self, code: str) -> Optional[str]:
+        """Validate code for safety issues."""
+        if not self.sandbox_mode:
+            return None
+        
+        try:
+            tree = ast.parse(code)
+        except SyntaxError as e:
+            return f"Syntax error: {e}"
+        
+        for node in ast.walk(tree):
+            # Check for dangerous imports
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    if alias.name in self.BLOCKED_IMPORTS:
+                        return f"Import blocked for safety: {alias.name}"
+            
+            elif isinstance(node, ast.ImportFrom):
+                if node.module in self.BLOCKED_IMPORTS:
+                    return f"Import blocked for safety: {node.module}"
+            
+            # Check for dangerous function calls
+            elif isinstance(node, ast.Call):
+                if isinstance(node.func, ast.Name):
+                    if node.func.id in self.BLOCKED_BUILTINS:
+                        return f"Function blocked for safety: {node.func.id}"
+        
+        return None
+    
+    def _create_safe_namespace(self) -> Dict[str, Any]:
+        """Create a restricted namespace for code execution."""
+        import math
+        import json
+        import re
+        from datetime import datetime, timedelta
+        
+        safe_builtins = {
+            'abs': abs, 'all': all, 'any': any, 'bin': bin, 'bool': bool,
+            'chr': chr, 'dict': dict, 'dir': dir, 'divmod': divmod,
+            'enumerate': enumerate, 'filter': filter, 'float': float,
+            'format': format, 'frozenset': frozenset, 'getattr': getattr,
+            'hasattr': hasattr, 'hash': hash, 'hex': hex, 'id': id,
+            'int': int, 'isinstance': isinstance, 'issubclass': issubclass,
+            'iter': iter, 'len': len, 'list': list, 'map': map, 'max': max,
+            'min': min, 'next': next, 'oct': oct, 'ord': ord, 'pow': pow,
+            'print': print, 'range': range, 'repr': repr, 'reversed': reversed,
+            'round': round, 'set': set, 'slice': slice, 'sorted': sorted,
+            'str': str, 'sum': sum, 'tuple': tuple, 'type': type, 'zip': zip,
+        }
+        
+        namespace = {
+            '__builtins__': safe_builtins,
+            'math': math,
+            'json': json,
+            're': re,
+            'datetime': datetime,
+            'timedelta': timedelta,
+        }
+        
+        # Try to import common data science libraries if available
+        try:
+            import numpy as np
+            namespace['np'] = np
+            namespace['numpy'] = np
+        except ImportError:
+            pass
+        
+        try:
+            import pandas as pd
+            namespace['pd'] = pd
+            namespace['pandas'] = pd
+        except ImportError:
+            pass
+        
+        return namespace
+    
+    async def execute(self, code: str, timeout: int = None, 
+                      reset_namespace: bool = False) -> ToolResult:
+        """Execute Python code."""
+        
+        # Validate code
+        validation_error = self._validate_code(code)
+        if validation_error:
+            return ToolResult.error_result(validation_error)
+        
+        # Reset namespace if requested
+        if reset_namespace:
+            self._namespace = {}
+        
+        # Prepare namespace
+        exec_namespace = self._create_safe_namespace()
+        exec_namespace.update(self._namespace)
+        
+        # Capture output
+        stdout_capture = io.StringIO()
+        stderr_capture = io.StringIO()
+        
+        execution_timeout = timeout or self.timeout
+        result_value = None
+        
+        def run_code():
+            nonlocal result_value
+            try:
+                # Try to evaluate as expression first
+                tree = ast.parse(code, mode='eval')
+                result_value = eval(compile(tree, '<string>', 'eval'), exec_namespace)
+            except SyntaxError:
+                # Not an expression, execute as statements
+                exec(code, exec_namespace)
+                # Check if there's a 'result' variable
+                if 'result' in exec_namespace:
+                    result_value = exec_namespace['result']
+        
+        try:
+            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+                # Run with timeout
+                await asyncio.wait_for(
+                    asyncio.get_event_loop().run_in_executor(None, run_code),
+                    timeout=execution_timeout
+                )
+            
+            # Update persistent namespace (excluding builtins)
+            for key, value in exec_namespace.items():
+                if key != '__builtins__':
+                    self._namespace[key] = value
+            
+            # Build output
+            stdout_str = stdout_capture.getvalue()
+            stderr_str = stderr_capture.getvalue()
+            
+            output_parts = []
+            if result_value is not None:
+                output_parts.append(f"Result: {repr(result_value)}")
+            if stdout_str:
+                output_parts.append(f"Output:\n{stdout_str}")
+            if stderr_str:
+                output_parts.append(f"Stderr:\n{stderr_str}")
+            
+            output = "\n".join(output_parts) if output_parts else "Code executed successfully (no output)"
+            
+            return ToolResult.success_result(
+                output=output,
+                metadata={
+                    'result': result_value,
+                    'stdout': stdout_str,
+                    'stderr': stderr_str,
+                    'namespace_vars': list(self._namespace.keys())
+                }
+            )
+            
+        except asyncio.TimeoutError:
+            return ToolResult.error_result(
+                f"Code execution timed out after {execution_timeout} seconds"
+            )
+        except Exception as e:
+            error_msg = traceback.format_exc()
+            return ToolResult.error_result(
+                f"Execution error: {str(e)}\n{error_msg}"
+            )
+    
+    def reset(self):
+        """Reset the execution namespace."""
+        self._namespace = {}
+    
+    def get_namespace_vars(self) -> List[str]:
+        """Get list of variables in the namespace."""
+        return list(self._namespace.keys())

package/astra/tools/shell.py

@@ -0,0 +1,183 @@
+"""
+Shell Command Tool for AstraAgent.
+Executes shell commands with safety guardrails.
+"""
+
+import asyncio
+import os
+import platform
+import shlex
+import subprocess
+from typing import List, Optional
+
+from astra.tools.base import Tool, ToolParameter, ToolResult, ToolCategory, RiskLevel
+
+
+class ShellTool(Tool):
+    """Tool for executing shell commands."""
+    
+    name = "shell"
+    description = "Execute shell commands on the system. Use with caution."
+    category = ToolCategory.SHELL
+    risk_level = RiskLevel.HIGH
+    
+    # Dangerous patterns to block
+    DANGEROUS_PATTERNS = [
+        "rm -rf /",
+        "rm -rf /*",
+        "mkfs",
+        "dd if=",
+        ":(){ :|:& };:",
+        "> /dev/sda",
+        "chmod -R 777 /",
+        "wget | sh",
+        "curl | sh",
+        "format c:",
+        "del /f /s /q c:\\",
+        "rd /s /q c:\\",
+    ]
+    
+    def __init__(self, allowed_commands: List[str] = None, blocked_commands: List[str] = None):
+        super().__init__()
+        self.allowed_commands = allowed_commands  # None = all allowed
+        self.blocked_commands = blocked_commands or []
+        self.is_windows = platform.system() == "Windows"
+    
+    @property
+    def parameters(self) -> List[ToolParameter]:
+        return [
+            ToolParameter(
+                name="command",
+                description="The shell command to execute",
+                type="string",
+                required=True
+            ),
+            ToolParameter(
+                name="working_dir",
+                description="Working directory for the command",
+                type="string",
+                required=False,
+                default=None
+            ),
+            ToolParameter(
+                name="timeout",
+                description="Timeout in seconds",
+                type="integer",
+                required=False,
+                default=60
+            ),
+            ToolParameter(
+                name="capture_output",
+                description="Whether to capture stdout/stderr",
+                type="boolean",
+                required=False,
+                default=True
+            )
+        ]
+    
+    def _is_dangerous(self, command: str) -> bool:
+        """Check if command is dangerous."""
+        cmd_lower = command.lower()
+        for pattern in self.DANGEROUS_PATTERNS + self.blocked_commands:
+            if pattern.lower() in cmd_lower:
+                return True
+        return False
+    
+    def _is_allowed(self, command: str) -> bool:
+        """Check if command is in allowed list."""
+        if self.allowed_commands is None:
+            return True
+        
+        cmd_parts = command.split()
+        if not cmd_parts:
+            return False
+        
+        base_cmd = cmd_parts[0]
+        return base_cmd in self.allowed_commands
+    
+    async def execute(self, command: str, working_dir: str = None, 
+                      timeout: int = 60, capture_output: bool = True) -> ToolResult:
+        """Execute a shell command."""
+        
+        # Safety checks
+        if self._is_dangerous(command):
+            return ToolResult.error_result(
+                f"Command blocked for safety: {command[:100]}..."
+            )
+        
+        if not self._is_allowed(command):
+            return ToolResult.error_result(
+                f"Command not in allowed list: {command.split()[0]}"
+            )
+        
+        # Resolve working directory
+        if working_dir:
+            if not os.path.isdir(working_dir):
+                return ToolResult.error_result(f"Working directory not found: {working_dir}")
+        else:
+            working_dir = os.getcwd()
+        
+        try:
+            # Execute command
+            if self.is_windows:
+                # Windows: use shell=True
+                process = await asyncio.create_subprocess_shell(
+                    command,
+                    stdout=asyncio.subprocess.PIPE if capture_output else None,
+                    stderr=asyncio.subprocess.PIPE if capture_output else None,
+                    cwd=working_dir
+                )
+            else:
+                # Unix: parse and execute
+                args = shlex.split(command)
+                process = await asyncio.create_subprocess_exec(
+                    *args,
+                    stdout=asyncio.subprocess.PIPE if capture_output else None,
+                    stderr=asyncio.subprocess.PIPE if capture_output else None,
+                    cwd=working_dir
+                )
+            
+            # Wait for completion with timeout
+            try:
+                stdout, stderr = await asyncio.wait_for(
+                    process.communicate(),
+                    timeout=timeout
+                )
+            except asyncio.TimeoutError:
+                process.kill()
+                return ToolResult.error_result(
+                    f"Command timed out after {timeout} seconds"
+                )
+            
+            # Decode output
+            stdout_str = stdout.decode('utf-8', errors='replace') if stdout else ""
+            stderr_str = stderr.decode('utf-8', errors='replace') if stderr else ""
+            
+            # Build result
+            if process.returncode == 0:
+                return ToolResult.success_result(
+                    output=stdout_str,
+                    metadata={
+                        'stderr': stderr_str,
+                        'return_code': process.returncode,
+                        'command': command,
+                        'working_dir': working_dir
+                    }
+                )
+            else:
+                return ToolResult.error_result(
+                    error=f"Command failed with code {process.returncode}:\n{stderr_str or stdout_str}",
+                    metadata={
+                        'stdout': stdout_str,
+                        'stderr': stderr_str,
+                        'return_code': process.returncode,
+                        'command': command
+                    }
+                )
+                
+        except FileNotFoundError:
+            return ToolResult.error_result(f"Command not found: {command.split()[0]}")
+        except PermissionError:
+            return ToolResult.error_result(f"Permission denied for command: {command}")
+        except Exception as e:
+            return ToolResult.error_result(f"Error executing command: {str(e)}")