🔍 Code Extractor

A sandboxed Python script executor that safely runs user-provided Python code with timeout controls, security restrictions, and isolated execution environments for data analysis tasks.

Purpose

This class provides a secure execution environment for running untrusted Python scripts, primarily for data analysis workflows. It validates scripts for dangerous operations, executes them in restricted namespaces with limited builtins, captures outputs and plots, and manages session-based file storage. The executor is designed to prevent malicious code execution while allowing legitimate data science operations using pandas, numpy, matplotlib, and other scientific libraries.

Source Code

class ScriptExecutor:
    """Sandboxed Python script executor with timeout and restrictions"""
    
    def __init__(self, config):
        self.config = config
        self.sandbox_dir = getattr(config, 'SANDBOX_FOLDER', Path(config.OUTPUT_DIR) / 'sandbox')
        self.timeout = getattr(config, 'SCRIPT_TIMEOUT', 60)
        self.allowed_imports = getattr(config, 'ALLOWED_IMPORTS', [
            'pandas', 'numpy', 'scipy', 'statsmodels', 'matplotlib', 'seaborn', 
            'sklearn', 'math', 'statistics', 'datetime', 'json'
        ])
        
        # Ensure sandbox directory exists
        os.makedirs(self.sandbox_dir, exist_ok=True)
        
    def execute_script(self, script: str, data: pd.DataFrame, 
                      session_id: str, step_id: str) -> Dict[str, Any]:
        """
        Execute Python script in sandboxed environment
        Returns execution results, outputs, and any errors
        """
        execution_result = {
            'success': False,
            'output': '',
            'error': '',
            'results': {},
            'plots': [],
            'execution_time': 0,
            'warnings': []
        }
        
        start_time = time.time()
        
        try:
            # Validate script safety
            validation_result = self._validate_script(script)
            if not validation_result['safe']:
                execution_result['error'] = f"Script validation failed: {validation_result['reason']}"
                return execution_result
            
            # Create isolated execution environment
            execution_context = self._create_execution_context(data, session_id, step_id)
            
            # Execute script with timeout
            output, error, results, plots = self._execute_with_timeout(script, execution_context)
            
            execution_result.update({
                'success': error == '',
                'output': output,
                'error': error,
                'results': results,
                'plots': plots,
                'execution_time': time.time() - start_time
            })
            
        except Exception as e:
            try:
                error_msg = f"Execution failed: {str(e)}"
                # Try to add traceback if possible
                exc_type, exc_value, exc_traceback = sys.exc_info()
                if exc_traceback:
                    try:
                        tb_lines = traceback.format_tb(exc_traceback)
                        error_msg += "\n" + "".join(tb_lines)
                    except Exception:
                        error_msg += "\n(Traceback unavailable)"
            except Exception:
                error_msg = f"Execution failed: {str(e)}"
            
            execution_result['error'] = error_msg
            execution_result['execution_time'] = time.time() - start_time
            logger.error(f"Script execution error: {error_msg}")
            
        return execution_result
    
    def _validate_script(self, script: str) -> Dict[str, Any]:
        """Validate script for security and safety"""
        validation = {'safe': True, 'reason': '', 'warnings': []}
        
        # Check for dangerous imports
        dangerous_imports = ['os', 'sys', 'subprocess', 'eval', 'exec', 'compile', 
                           'open', '__import__', 'vars']
        
        for dangerous in dangerous_imports:
            if dangerous in script:
                # More precise checking
                import re
                patterns = [
                    rf'\bimport\s+{dangerous}\b',
                    rf'\bfrom\s+{dangerous}\s+import\b',
                    rf'\b{dangerous}\s*\(',
                ]
                
                for pattern in patterns:
                    if re.search(pattern, script):
                        validation['safe'] = False
                        validation['reason'] = f"Dangerous import/function detected: {dangerous}"
                        return validation
        
        # Check for file operations (except matplotlib save)
        file_ops = ['open(', 'file(', 'read(', 'write(']
        for op in file_ops:
            if op in script and 'plt.savefig' not in script:
                validation['warnings'].append(f"File operation detected: {op}")
        
        # Check script length
        if len(script) > 10000:  # 10KB limit
            validation['safe'] = False
            validation['reason'] = "Script too long"
            return validation
        
        # Check for infinite loops (basic detection)
        if 'while True:' in script:
            validation['warnings'].append("Potential infinite loop detected")
        
        return validation
    
    def _create_execution_context(self, data: pd.DataFrame, session_id: str, step_id: str) -> Dict[str, Any]:
        """Create isolated execution context"""
        
        # Create session-specific directory
        session_dir = self.sandbox_dir / session_id
        os.makedirs(session_dir, exist_ok=True)
        
        # Create step-specific directory for outputs
        step_dir = session_dir / step_id
        os.makedirs(step_dir, exist_ok=True)
        
        context = {
            'df': data.copy(),  # Provide data as 'df'
            'session_dir': str(session_dir),
            'step_dir': str(step_dir),
            'plots_dir': str(step_dir / 'plots'),
            'np': np,
            'pd': pd,
            'plt': plt,
            'sns': sns,
            'stats': stats,
            'sm': sm,
            'warnings': warnings
        }
        
        # Create plots directory
        os.makedirs(context['plots_dir'], exist_ok=True)
        
        return context
    
    def _execute_with_timeout(self, script: str, context: Dict[str, Any]) -> Tuple[str, str, Dict, List[str]]:
        """Execute script with timeout and capture outputs"""
        
        # Prepare script with context setup
        full_script = self._prepare_script(script, context)
        
        # Capture stdout and stderr
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        results = {}
        plots = []
        
        try:
            # Redirect output
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture
            
            # Execute script in restricted environment (no timeout - handled by subprocess in production)
            # Signal-based timeout doesn't work in Flask worker threads
            exec_globals = self._create_restricted_globals(context)
            exec_locals = {}
            
            exec(full_script, exec_globals, exec_locals)
            
            # Extract results if available (check both locals and globals)
            if 'results' in exec_locals:
                results = exec_locals['results']
            elif 'results' in exec_globals:
                results = exec_globals['results']
            
            # Find generated plots
            plots = self._collect_plots(context['plots_dir'])
            
            output = stdout_capture.getvalue()
            error = stderr_capture.getvalue()
            
        except TimeoutError:
            error = f"Script execution timed out after {self.timeout} seconds"
            output = stdout_capture.getvalue()
            
        except Exception as e:
            # Get exception info without using format_exc which might have compatibility issues
            exc_type, exc_value, exc_traceback = sys.exc_info()
            error_lines = [f"Execution error: {str(e)}"]
            
            # Try to get traceback info safely
            try:
                if exc_traceback:
                    tb_lines = traceback.format_tb(exc_traceback)
                    error_lines.extend(tb_lines)
            except Exception:
                error_lines.append("(Traceback unavailable due to compatibility issue)")
            
            error = "\n".join(error_lines)
            output = stdout_capture.getvalue()
            
        finally:
            # Restore stdout/stderr
            sys.stdout = old_stdout
            sys.stderr = old_stderr
        
        return output, error, results, plots
    
    def _prepare_script(self, script: str, context: Dict[str, Any]) -> str:
        """Prepare script with necessary setup and context"""
        
        setup_code = f"""
# Setup code - automatically added
# Modules are pre-imported and available as: pd, np, plt, sns, stats, sm, warnings

# Suppress warnings
warnings.filterwarnings('ignore')

# Set up matplotlib for non-interactive use
plt.ioff()
plt.style.use('default')

# Set up directories
plots_dir = r"{context['plots_dir']}"

# Data is available as 'df' and 'data'
data = df  # Alias for compatibility

# Initialize results dictionary at global scope
results = {{
    'summary_statistics': {{}},
    'test_results': {{}},
    'plots': [],
    'assumptions': {{}},
    'interpretation': ''
}}

# Override plt.show to auto-save instead
original_show = plt.show
def show_and_save():
    global results
    try:
        # Get current figure and save it directly
        fig = plt.gcf()
        if fig.get_axes():  # Only save if figure has content
            filename = f"plot_{{len(results.get('plots', []))}}.png"
            filepath = plots_dir + "/" + filename
            fig.savefig(filepath, dpi=150, bbox_inches='tight')
            if 'plots' not in results:
                results['plots'] = []
            results['plots'].append(filename)
            plt.close(fig)
    except Exception as e:
        print(f"Warning: Could not save plot: {{e}}")
    
plt.show = show_and_save

"""
        
        return setup_code + "\n" + script
    
    def _create_restricted_globals(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Create restricted global namespace for execution"""
        
        # Start with minimal safe builtins
        safe_builtins = {
            'abs', 'all', 'any', 'bool', 'dict', 'enumerate', 'filter',
            'float', 'int', 'len', 'list', 'map', 'max', 'min', 'range',
            'round', 'set', 'sorted', 'str', 'sum', 'tuple', 'type', 'zip',
            'print', 'isinstance', 'hasattr', 'getattr', 'setattr', 'locals', 'globals',
            # Add exception types for error handling
            'Exception', 'ValueError', 'TypeError', 'KeyError', 'IndexError', 'AttributeError'
        }
        
        # Create restricted builtins
        restricted_builtins = {name: __builtins__[name] for name in safe_builtins if name in __builtins__}
        
        # Add mathematical functions
        import math
        restricted_builtins['math'] = math
        
        globals_dict = {
            '__builtins__': restricted_builtins,
            # Add allowed modules
            'pd': pd,
            'np': np,
            'plt': plt,
            'sns': sns,
            'stats': stats,
            'sm': sm,
            'warnings': warnings,
            # Add data
            'df': context['df']
        }
        
        return globals_dict
    
    def _collect_plots(self, plots_dir: str) -> List[str]:
        """Collect generated plot files"""
        plots = []
        
        if os.path.exists(plots_dir):
            for file in os.listdir(plots_dir):
                if file.endswith(('.png', '.jpg', '.jpeg', '.svg', '.pdf')):
                    plots.append(os.path.join(plots_dir, file))
        
        return plots
    
    def extract_results_summary(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Extract and format key results for display"""
        summary = {
            'statistics': {},
            'tests': {},
            'plots_count': 0,
            'key_findings': []
        }
        
        try:
            # Extract summary statistics
            if 'summary_statistics' in results:
                summary['statistics'] = results['summary_statistics']
            
            # Extract test results
            if 'test_results' in results:
                summary['tests'] = results['test_results']
            
            # Count plots
            if 'plots' in results:
                summary['plots_count'] = len(results['plots'])
            
            # Extract interpretation
            if 'interpretation' in results:
                summary['interpretation'] = results['interpretation']
            
            # Try to extract key numerical findings
            for key, value in results.items():
                if isinstance(value, (int, float)) and not np.isnan(value):
                    summary['key_findings'].append(f"{key}: {value}")
                elif isinstance(value, dict):
                    for subkey, subvalue in value.items():
                        if isinstance(subvalue, (int, float)) and not np.isnan(subvalue):
                            summary['key_findings'].append(f"{key}.{subkey}: {subvalue}")
                            
        except Exception as e:
            logger.warning(f"Error extracting results summary: {str(e)}")
        
        return summary
    
    def cleanup_session(self, session_id: str, keep_recent: int = 5):
        """Clean up old session files"""
        try:
            session_dir = self.sandbox_dir / session_id
            if session_dir.exists():
                # Get all step directories
                step_dirs = [d for d in session_dir.iterdir() if d.is_dir()]
                step_dirs.sort(key=lambda x: x.stat().st_mtime, reverse=True)
                
                # Remove old step directories
                for step_dir in step_dirs[keep_recent:]:
                    import shutil
                    shutil.rmtree(step_dir)
                    
        except Exception as e:
            logger.warning(f"Error cleaning up session {session_id}: {str(e)}")
    
    def get_plot_data(self, plot_path: str) -> Optional[str]:
        """Get plot as base64 encoded string for web display"""
        try:
            if os.path.exists(plot_path):
                with open(plot_path, 'rb') as f:
                    plot_data = base64.b64encode(f.read()).decode('utf-8')
                return plot_data
        except Exception as e:
            logger.error(f"Error reading plot {plot_path}: {str(e)}")
        return None
    
    def validate_data_access(self, script: str, available_columns: List[str]) -> Dict[str, Any]:
        """Validate that script only accesses available columns"""
        validation = {
            'valid': True,
            'warnings': [],
            'accessed_columns': []
        }
        
        import re
        
        # Find column references in script
        column_patterns = [
            r"df\[['\"](.*?)['\"]\]",  # df['column']
            r"df\.(\w+)",              # df.column
        ]
        
        accessed_columns = set()
        for pattern in column_patterns:
            matches = re.findall(pattern, script)
            accessed_columns.update(matches)
        
        validation['accessed_columns'] = list(accessed_columns)
        
        # Check if columns exist
        missing_columns = accessed_columns - set(available_columns)
        if missing_columns:
            validation['warnings'].append(f"Columns not found in dataset: {list(missing_columns)}")
        
        return validation

Parameters

Parameter Details

config: Configuration object containing settings for the executor. Expected attributes include: SANDBOX_FOLDER (Path for isolated execution directories), SCRIPT_TIMEOUT (execution timeout in seconds, default 60), ALLOWED_IMPORTS (list of permitted module names), and OUTPUT_DIR (base directory for outputs). The config object should have these as attributes accessible via getattr.

Return Value

Instantiation returns a ScriptExecutor object. The main execute_script method returns a dictionary with keys: 'success' (bool), 'output' (captured stdout string), 'error' (error messages string), 'results' (dict of script results), 'plots' (list of plot file paths), 'execution_time' (float in seconds), and 'warnings' (list of warning messages). Other methods return validation dictionaries, plot data as base64 strings, or None.

Class Interface

Methods

Attributes

Dependencies

• os
• sys
• subprocess
• tempfile
• json
• logging
• traceback
• io
• contextlib
• time
• signal
• threading
• queue
• typing
• pathlib
• pickle
• base64
• pandas
• numpy
• matplotlib
• seaborn
• scipy
• statsmodels
• warnings
• math
• re
• shutil

Required Imports

import os
import sys
import subprocess
import tempfile
import json
import logging
import traceback
import io
import contextlib
import time
import signal
import threading
import queue
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path
import pickle
import base64
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
import statsmodels.api as sm
import warnings
import math
import re

Conditional/Optional Imports

These imports are only needed under specific conditions:

import shutil

Usage Example

import pandas as pd
from pathlib import Path

# Create a simple config object
class Config:
    OUTPUT_DIR = Path('./output')
    SANDBOX_FOLDER = Path('./output/sandbox')
    SCRIPT_TIMEOUT = 60
    ALLOWED_IMPORTS = ['pandas', 'numpy', 'matplotlib', 'seaborn']

config = Config()

# Instantiate the executor
executor = ScriptExecutor(config)

# Prepare sample data
data = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})

# Define a script to execute
script = '''
import pandas as pd
import numpy as np

# Perform analysis
results['mean_A'] = df['A'].mean()
results['sum_B'] = df['B'].sum()
results['correlation'] = df['A'].corr(df['B'])

print(f"Mean of A: {results['mean_A']}")
print(f"Sum of B: {results['sum_B']}")
'''

# Execute the script
result = executor.execute_script(
    script=script,
    data=data,
    session_id='session_123',
    step_id='step_001'
)

# Check results
if result['success']:
    print(f"Output: {result['output']}")
    print(f"Results: {result['results']}")
    print(f"Execution time: {result['execution_time']}s")
else:
    print(f"Error: {result['error']}")

# Extract summary
summary = executor.extract_results_summary(result['results'])
print(f"Key findings: {summary['key_findings']}")

# Clean up old sessions
executor.cleanup_session('session_123', keep_recent=3)

Best Practices

• Always validate scripts before execution using the built-in _validate_script method (called automatically by execute_script)
• Use unique session_id and step_id for each execution to prevent file conflicts and enable proper cleanup
• Regularly call cleanup_session to prevent disk space issues from accumulated plot files and execution artifacts
• Check the 'success' field in the returned dictionary before accessing 'results' or 'plots'
• The executor modifies sys.stdout and sys.stderr during execution but restores them in finally blocks - avoid nested executions
• Scripts execute with a restricted global namespace - only safe builtins and specified modules are available
• The 'df' variable is automatically provided to scripts as a copy of the input data to prevent mutations
• Plots are automatically saved when plt.show() is called in the script - the method is overridden to save instead of display
• Results must be stored in a global 'results' dictionary within the script to be captured
• File operations are restricted - scripts cannot directly open files except through matplotlib's savefig
• The timeout mechanism may not work in all environments (e.g., Flask worker threads) - consider external process isolation for production
• Dangerous imports (os, sys, subprocess, eval, exec) are blocked by validation - scripts using these will fail
• The sandbox directory structure is: sandbox_dir/session_id/step_id/plots/ - plan storage accordingly
• Use extract_results_summary to get a formatted view of execution results for display purposes
• The executor is not thread-safe for the same session_id - use different session IDs for concurrent executions

Tags

Similar Components

AI-powered semantic similarity - components with related functionality:

• class ScriptExecutor_v1 96.6% similar

  A sandboxed Python script executor that safely runs user-provided Python code with timeout controls, security restrictions, and isolated execution environments for data analysis tasks.

  From: /tf/active/vicechatdev/full_smartstat/script_executor.py

• class AgentExecutor 60.0% similar

  Agent-based script executor that generates standalone Python files, manages dependencies, and provides iterative debugging capabilities

  From: /tf/active/vicechatdev/vice_ai/agent_executor.py

• class AgentExecutor_v1 59.8% similar

  Agent-based script executor that generates standalone Python files, manages dependencies, and provides iterative debugging capabilities

  From: /tf/active/vicechatdev/full_smartstat/agent_executor.py

• class AgentExecutor_v2 59.8% similar

  Agent-based script executor that generates standalone Python files, manages dependencies, and provides iterative debugging capabilities

  From: /tf/active/vicechatdev/smartstat/agent_executor.py

• function test_agent_executor 49.3% similar

  Integration test function that validates the AgentExecutor's ability to generate and execute data analysis projects using synthetic test data.

  From: /tf/active/vicechatdev/full_smartstat/debug_agent.py