tools/code_execution_tool.py

"""
Secure Code Execution Tool for GAIA Agent
Provides safe Python code execution with mathematical computation capabilities.

Features:
- Secure sandboxed execution environment
- Mathematical libraries (numpy, scipy, sympy, pandas)
- Timeout and resource management
- Result validation and formatting
- Security restrictions and input sanitization
"""

import os
import sys
import ast
import subprocess
import tempfile
import time
import signal
import logging
import traceback
import re
from typing import Dict, Any, Optional, Union, List
from pathlib import Path
import json

# Mathematical and scientific computing libraries
try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False

try:
    import pandas as pd
    PANDAS_AVAILABLE = True
except ImportError:
    PANDAS_AVAILABLE = False

try:
    import scipy
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False

try:
    import sympy as sp
    SYMPY_AVAILABLE = True
except ImportError:
    SYMPY_AVAILABLE = False

try:
    import matplotlib
    matplotlib.use('Agg')  # Non-interactive backend
    import matplotlib.pyplot as plt
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

logger = logging.getLogger(__name__)


class SecurityError(Exception):
    """Raised when code contains potentially dangerous operations."""
    pass


class ExecutionTimeoutError(Exception):
    """Raised when code execution exceeds timeout limit."""
    pass


class CodeSecurityValidator:
    """Validates Python code for security risks before execution."""
    
    # Dangerous imports and functions to block
    BLOCKED_IMPORTS = {
        'os', 'sys', 'subprocess', 'shutil', 'glob', 'pickle', 'marshal',
        'importlib', '__import__', 'eval', 'exec', 'compile', 'open',
        'file', 'input', 'raw_input', 'reload', 'vars', 'locals', 'globals',
        'dir', 'hasattr', 'getattr', 'setattr', 'delattr', 'callable',
        'socket', 'urllib', 'requests', 'http', 'ftplib', 'smtplib',
        'telnetlib', 'poplib', 'imaplib', 'nntplib', 'ssl', 'hashlib',
        'hmac', 'secrets', 'random', 'tempfile', 'threading', 'multiprocessing'
    }
    
    BLOCKED_FUNCTIONS = {
        'eval', 'exec', 'compile', '__import__', 'open', 'file', 'input',
        'raw_input', 'reload', 'vars', 'locals', 'globals', 'dir',
        'hasattr', 'getattr', 'setattr', 'delattr', 'callable'
    }
    
    BLOCKED_ATTRIBUTES = {
        '__class__', '__bases__', '__subclasses__', '__mro__', '__globals__',
        '__code__', '__func__', '__self__', '__module__', '__dict__',
        '__getattribute__', '__setattr__', '__delattr__', '__reduce__',
        '__reduce_ex__', '__getstate__', '__setstate__'
    }
    
    def validate_code(self, code: str) -> bool:
        """
        Validate Python code for security risks.
        
        Args:
            code: Python code string to validate
            
        Returns:
            True if code is safe, raises SecurityError if dangerous
        """
        try:
            # Parse the code into an AST
            tree = ast.parse(code)
            
            # Walk through all nodes in the AST
            for node in ast.walk(tree):
                self._check_node(node)
                
            return True
            
        except SyntaxError as e:
            raise SecurityError(f"Syntax error in code: {e}")
        except Exception as e:
            raise SecurityError(f"Code validation failed: {e}")
    
    def _check_node(self, node: ast.AST) -> None:
        """Check individual AST node for security risks."""
        
        # Check imports
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            self._check_import(node)
        
        # Check function calls
        elif isinstance(node, ast.Call):
            self._check_function_call(node)
        
        # Check attribute access
        elif isinstance(node, ast.Attribute):
            self._check_attribute_access(node)
        
        # Check name access
        elif isinstance(node, ast.Name):
            self._check_name_access(node)
    
    def _check_import(self, node: Union[ast.Import, ast.ImportFrom]) -> None:
        """Check import statements for dangerous modules."""
        if isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name in self.BLOCKED_IMPORTS:
                    raise SecurityError(f"Blocked import: {alias.name}")
        
        elif isinstance(node, ast.ImportFrom):
            if node.module and node.module in self.BLOCKED_IMPORTS:
                raise SecurityError(f"Blocked import from: {node.module}")
    
    def _check_function_call(self, node: ast.Call) -> None:
        """Check function calls for dangerous operations."""
        if isinstance(node.func, ast.Name):
            if node.func.id in self.BLOCKED_FUNCTIONS:
                raise SecurityError(f"Blocked function call: {node.func.id}")
    
    def _check_attribute_access(self, node: ast.Attribute) -> None:
        """Check attribute access for dangerous attributes."""
        if node.attr in self.BLOCKED_ATTRIBUTES:
            raise SecurityError(f"Blocked attribute access: {node.attr}")
    
    def _check_name_access(self, node: ast.Name) -> None:
        """Check name access for blocked identifiers."""
        if node.id in self.BLOCKED_FUNCTIONS:
            # Allow if it's being assigned to (not called)
            if not isinstance(node.ctx, ast.Store):
                raise SecurityError(f"Blocked name access: {node.id}")


class SecureCodeExecutor:
    """Secure Python code executor with mathematical capabilities."""
    
    def __init__(self, timeout: int = 30, memory_limit_mb: int = 512):
        """
        Initialize secure code executor.
        
        Args:
            timeout: Maximum execution time in seconds
            memory_limit_mb: Maximum memory usage in MB
        """
        self.timeout = timeout
        self.memory_limit_mb = memory_limit_mb
        self.validator = CodeSecurityValidator()
        
        # Available libraries status
        self.available_libraries = {
            'numpy': NUMPY_AVAILABLE,
            'pandas': PANDAS_AVAILABLE,
            'scipy': SCIPY_AVAILABLE,
            'sympy': SYMPY_AVAILABLE,
            'matplotlib': MATPLOTLIB_AVAILABLE
        }
        
        logger.info(f"SecureCodeExecutor initialized with {timeout}s timeout, {memory_limit_mb}MB limit")
        logger.info(f"Available libraries: {[lib for lib, avail in self.available_libraries.items() if avail]}")
    
    def execute_code(self, code: str, return_output: bool = True) -> Dict[str, Any]:
        """
        Execute Python code securely and return results.
        
        Args:
            code: Python code to execute
            return_output: Whether to capture and return output
            
        Returns:
            Dictionary with execution results
        """
        start_time = time.time()
        
        try:
            # Validate code security
            self.validator.validate_code(code)
            
            # Prepare execution environment
            execution_result = self._execute_in_subprocess(code, return_output)
            
            execution_time = time.time() - start_time
            
            return {
                'success': True,
                'result': execution_result.get('result'),
                'output': execution_result.get('output', ''),
                'error': None,
                'execution_time': execution_time,
                'libraries_used': self._detect_libraries_used(code)
            }
            
        except SecurityError as e:
            return {
                'success': False,
                'result': None,
                'output': '',
                'error': f"Security violation: {e}",
                'execution_time': time.time() - start_time,
                'libraries_used': []
            }
        
        except ExecutionTimeoutError as e:
            return {
                'success': False,
                'result': None,
                'output': '',
                'error': f"Execution timeout: {e}",
                'execution_time': self.timeout,
                'libraries_used': []
            }
        
        except Exception as e:
            return {
                'success': False,
                'result': None,
                'output': '',
                'error': f"Execution error: {e}",
                'execution_time': time.time() - start_time,
                'libraries_used': []
            }
    
    def _execute_in_subprocess(self, code: str, return_output: bool) -> Dict[str, Any]:
        """Execute code in a secure subprocess."""
        
        # Create temporary file for code execution
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            # Prepare safe execution environment
            safe_code = self._prepare_safe_code(code, return_output)
            f.write(safe_code)
            temp_file = f.name
        
        try:
            # Execute in subprocess with timeout and resource limits
            result = subprocess.run(
                [sys.executable, temp_file],
                capture_output=True,
                text=True,
                timeout=self.timeout,
                cwd=tempfile.gettempdir()  # Run in temp directory
            )
            
            if result.returncode == 0:
                # Parse output
                output_lines = result.stdout.strip().split('\n')
                if return_output and output_lines:
                    # Last line should be the result if we added result capture
                    if output_lines[-1].startswith('RESULT:'):
                        result_str = output_lines[-1][7:]  # Remove 'RESULT:' prefix
                        output = '\n'.join(output_lines[:-1])
                        try:
                            # Try to parse as JSON for complex types
                            parsed_result = json.loads(result_str)
                        except:
                            # Fall back to string result
                            parsed_result = result_str
                        
                        return {
                            'result': parsed_result,
                            'output': output
                        }
                    else:
                        return {
                            'result': None,
                            'output': result.stdout
                        }
                else:
                    return {
                        'result': None,
                        'output': result.stdout
                    }
            else:
                raise Exception(f"Code execution failed: {result.stderr}")
                
        except subprocess.TimeoutExpired:
            raise ExecutionTimeoutError(f"Code execution exceeded {self.timeout} seconds")
        
        finally:
            # Clean up temporary file
            try:
                os.unlink(temp_file)
            except:
                pass
    
    def _prepare_safe_code(self, code: str, capture_result: bool) -> str:
        """Prepare code for safe execution with necessary imports and result capture."""
        
        safe_imports = []
        
        # Add available mathematical libraries
        if NUMPY_AVAILABLE:
            safe_imports.append("import numpy as np")
        if PANDAS_AVAILABLE:
            safe_imports.append("import pandas as pd")
        if SCIPY_AVAILABLE:
            safe_imports.append("import scipy")
            safe_imports.append("from scipy import stats, optimize, integrate, linalg")
        if SYMPY_AVAILABLE:
            safe_imports.append("import sympy as sp")
            safe_imports.append("from sympy import symbols, solve, diff, integrate as sp_integrate, simplify, expand, factor")
        if MATPLOTLIB_AVAILABLE:
            safe_imports.append("import matplotlib")
            safe_imports.append("matplotlib.use('Agg')")
            safe_imports.append("import matplotlib.pyplot as plt")
        
        # Add basic math and other safe imports
        safe_imports.extend([
            "import math",
            "import cmath",
            "import decimal",
            "import fractions",
            "import statistics",
            "import itertools",
            "import functools",
            "import operator",
            "import json"
        ])
        
        # Prepare the complete code
        complete_code = '\n'.join(safe_imports) + '\n\n'
        
        if capture_result:
            # Wrap user code to capture the last expression result
            complete_code += '''
# User code execution
import sys
from io import StringIO

# Capture stdout
old_stdout = sys.stdout
sys.stdout = captured_output = StringIO()

try:
    # Execute user code and capture result
    user_code = """''' + code.replace('"""', '\\"\\"\\"') + '''"""
    
    # Execute the code
    exec(user_code)
    
    # Try to capture the result of the last expression
    import ast
    try:
        tree = ast.parse(user_code)
        if tree.body and isinstance(tree.body[-1], ast.Expr):
            # Last statement is an expression, evaluate it
            last_expr = ast.Expression(tree.body[-1].value)
            result = eval(compile(last_expr, '<string>', 'eval'))
            print(f"RESULT:{json.dumps(result) if isinstance(result, (int, float, str, list, dict, bool)) else str(result)}")
        else:
            print("RESULT:None")
    except:
        print("RESULT:None")

finally:
    # Restore stdout and print captured output
    sys.stdout = old_stdout
    output = captured_output.getvalue()
    if output:
        print(output, end='')
'''
        else:
            complete_code += code
        
        return complete_code
    
    def _detect_libraries_used(self, code: str) -> List[str]:
        """Detect which mathematical libraries are used in the code."""
        libraries_used = []
        
        # Simple detection based on import statements and usage
        if 'numpy' in code or 'np.' in code:
            libraries_used.append('numpy')
        if 'pandas' in code or 'pd.' in code:
            libraries_used.append('pandas')
        if 'scipy' in code:
            libraries_used.append('scipy')
        if 'sympy' in code or 'sp.' in code:
            libraries_used.append('sympy')
        if 'matplotlib' in code or 'plt.' in code:
            libraries_used.append('matplotlib')
        if 'math.' in code:
            libraries_used.append('math')
        
        return libraries_used


class CodeExecutionTool:
    """AGNO-compatible tool for secure Python code execution."""
    
    def __init__(self, timeout: int = 30, memory_limit_mb: int = 512):
        """Initialize the code execution tool."""
        self.executor = SecureCodeExecutor(timeout, memory_limit_mb)
        self.available = True
        
        logger.info("CodeExecutionTool initialized successfully")
    
    def execute_python_code(self, code: str) -> str:
        """
        Execute Python code and return the result.
        
        Args:
            code: Python code to execute
            
        Returns:
            Formatted result string
        """
        result = self.executor.execute_code(code, return_output=True)
        
        if result['success']:
            output_parts = []
            
            if result['output']:
                output_parts.append(f"Output:\n{result['output']}")
            
            if result['result'] is not None:
                output_parts.append(f"Result: {result['result']}")
            
            if result['libraries_used']:
                output_parts.append(f"Libraries used: {', '.join(result['libraries_used'])}")
            
            output_parts.append(f"Execution time: {result['execution_time']:.3f}s")
            
            return '\n'.join(output_parts)
        else:
            return f"Error: {result['error']}"
    
    def run_mathematical_computation(self, expression: str) -> str:
        """
        Run a mathematical computation using available libraries.
        
        Args:
            expression: Mathematical expression or computation
            
        Returns:
            Computation result
        """
        # Prepare code for mathematical computation
        code = f"""
# Mathematical computation
result = {expression}
print(f"Computation: {expression}")
print(f"Result: {{result}}")
result
"""
        
        return self.execute_python_code(code)
    
    def analyze_numerical_data(self, data: str, operation: str = "basic_stats") -> str:
        """
        Analyze numerical data using pandas and numpy.
        
        Args:
            data: Data as string (comma-separated values or JSON)
            operation: Type of analysis to perform
            
        Returns:
            Analysis results
        """
        code = f"""
import json

# Parse data
try:
    data = json.loads('{data}')
except:
    data = [float(x.strip()) for x in '{data}'.split(',') if x.strip()]

# Convert to numpy array for analysis
data_array = np.array(data)

# Perform analysis
if '{operation}' == 'basic_stats':
    result = {{
        'mean': float(np.mean(data_array)),
        'median': float(np.median(data_array)),
        'std': float(np.std(data_array)),
        'min': float(np.min(data_array)),
        'max': float(np.max(data_array)),
        'sum': float(np.sum(data_array)),
        'count': len(data_array)
    }}
elif '{operation}' == 'advanced_stats':
    result = {{
        'mean': float(np.mean(data_array)),
        'variance': float(np.var(data_array)),
        'skewness': float(stats.skew(data_array)) if 'stats' in globals() else 'N/A',
        'kurtosis': float(stats.kurtosis(data_array)) if 'stats' in globals() else 'N/A',
        'percentiles': {{
            '25th': float(np.percentile(data_array, 25)),
            '50th': float(np.percentile(data_array, 50)),
            '75th': float(np.percentile(data_array, 75))
        }}
    }}
else:
    result = 'Unknown operation'

print(f"Data analysis ({operation}):")
print(f"Data: {{data}}")
print(f"Results: {{result}}")
result
"""
        
        return self.execute_python_code(code)
    
    def get_status(self) -> Dict[str, Any]:
        """Get tool status and capabilities."""
        return {
            'available': self.available,
            'timeout': self.executor.timeout,
            'memory_limit_mb': self.executor.memory_limit_mb,
            'available_libraries': self.executor.available_libraries,
            'security_features': [
                'AST-based code validation',
                'Subprocess isolation',
                'Import restrictions',
                'Function call blocking',
                'Attribute access control',
                'Timeout protection',
                'Memory limits'
            ]
        }


# AGNO tool registration functions
def get_code_execution_tools():
    """Get code execution tools for AGNO registration."""
    tool = CodeExecutionTool()
    
    # Return tool methods that can be called by AGNO
    return [
        {
            'name': 'execute_python_code',
            'function': tool.execute_python_code,
            'description': 'Execute Python code securely with mathematical libraries'
        },
        {
            'name': 'run_mathematical_computation',
            'function': tool.run_mathematical_computation,
            'description': 'Perform mathematical computations using numpy, scipy, sympy'
        },
        {
            'name': 'analyze_numerical_data',
            'function': tool.analyze_numerical_data,
            'description': 'Analyze numerical data with statistical operations'
        }
    ]


if __name__ == "__main__":
    # Test the code execution tool
    tool = CodeExecutionTool()
    
    # Test basic mathematical computation
    test_code = """
import math
result = math.sqrt(2) * math.pi
print(f"Square root of 2 times pi: {result}")
result
"""
    
    print("Testing CodeExecutionTool:")
    print("=" * 50)
    result = tool.execute_python_code(test_code)
    print(result)
    print("=" * 50)
    
    # Test status
    status = tool.get_status()
    print("Tool Status:")
    print(json.dumps(status, indent=2))