Patched code sample

import ast
import io
from contextlib import redirect_stdout

# This code represents a potential fix for a vulnerability like CVE-2025-5120.
# The vulnerability described involves bypassing inadequate checks in a Python sandbox.
# A robust fix involves parsing the code into an Abstract Syntax Tree (AST)
# and validating every node against a strict whitelist before execution.
# This is far more secure than simple string-based blacklisting.

class SafePythonExecutor:
    """
    A Python code executor that uses AST analysis to prevent sandbox escapes.
    It ensures that only whitelisted modules, functions, and language features are used.
    """

    def __init__(self, allowed_modules=None, max_output_chars=1024):
        """
        Initializes the safe executor.

        Args:
            allowed_modules (dict): A dictionary where keys are module names and
                                    values are lists of allowed function/attribute names.
                                    Example: {'math': ['sqrt', 'pi'], 'random': ['random']}
            max_output_chars (int): The maximum number of characters for stdout.
        """
        if allowed_modules is None:
            # By default, only allow a minimal set of safe math functions.
            self.allowed_modules = {
                'math': ['sqrt', 'pow', 'sin', 'cos', 'tan', 'pi', 'e']
            }
        else:
            self.allowed_modules = allowed_modules

        # A set of built-in functions that are generally safe to expose.
        self.allowed_builtins = {
            'print', 'len', 'str', 'int', 'float', 'round', 'range', 'abs',
            'min', 'max', 'sum', 'sorted'
        }
        self.max_output_chars = max_output_chars

    def _validate_node(self, node):
        """
        Recursively validates each node in the AST.
        This is the core of the security boundary.
        """
        # Disallow accessing private/dunder attributes to prevent introspection attacks.
        if isinstance(node, ast.Attribute) and node.attr.startswith('_'):
            raise PermissionError(f"Access to private/dunder attributes like '{node.attr}' is forbidden.")

        # Disallow dangerous function calls like exec, eval, open, getattr, etc.
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name) and node.func.id not in self.allowed_builtins:
                raise PermissionError(f"Calling the built-in function '{node.func.id}' is not allowed.")
            if isinstance(node.func, ast.Attribute):
                # e.g., in `math.sqrt()`, `node.func.value` is `math` (an ast.Name)
                # and `node.func.attr` is `sqrt` (a string).
                if isinstance(node.func.value, ast.Name):
                    module_name = node.func.value.id
                    func_name = node.func.attr
                    if module_name not in self.allowed_modules:
                        raise PermissionError(f"Module '{module_name}' is not whitelisted.")
                    if func_name not in self.allowed_modules[module_name]:
                        raise PermissionError(f"Function '{func_name}' from module '{module_name}' is not allowed.")

        # Disallow imports that are not on the whitelist.
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            for alias in node.names:
                if alias.name not in self.allowed_modules:
                    raise PermissionError(f"Importing module '{alias.name}' is not allowed.")

        # Recursively check all child nodes.
        for child_node in ast.iter_child_nodes(node):
            self._validate_node(child_node)

    def execute(self, code_string: str) -> str:
        """
        Validates and executes the given Python code in a restricted environment.

        Args:
            code_string (str): The untrusted Python code to execute.

        Returns:
            str: The captured stdout from the executed code.
        
        Raises:
            PermissionError: If the code uses disallowed features.
            Exception: Any other exception raised by the user's code.
        """
        try:
            # 1. Static Analysis: Parse the code into an AST.
            tree = ast.parse(code_string)

            # 2. Validation: Walk the AST and check every node against security rules.
            self._validate_node(tree)

            # 3. Restricted Execution: If validation passes, execute the code.
            # Prepare a safe global scope for the execution.
            safe_globals = {}
            for module_name, funcs in self.allowed_modules.items():
                try:
                    module = __import__(module_name)
                    # Only expose the whitelisted functions/attributes of the module.
                    safe_globals[module_name] = {
                        func: getattr(module, func) for func in funcs
                    }
                except ImportError:
                    # This should not happen if the environment is set up correctly.
                    raise RuntimeError(f"Could not import whitelisted module '{module_name}'.")
            
            # Add allowed built-ins to the scope.
            for builtin_name in self.allowed_builtins:
                safe_globals[builtin_name] = __builtins__.__dict__[builtin_name]
            
            # Capture stdout to return the result.
            stdout_capture = io.StringIO()
            with redirect_stdout(stdout_capture):
                # Execute the code with the restricted global scope and no local scope.
                # The '__builtins__' key is set to a restricted dict to prevent access to all builtins.
                exec(code_string, {'__builtins__': safe_globals})
            
            output = stdout_capture.getvalue()
            if len(output) > self.max_output_chars:
                return output[:self.max_output_chars] + "\n...[output truncated]..."
            return output

        except (PermissionError, SyntaxError, NameError, TypeError) as e:
            return f"Execution failed: {type(e).__name__}: {e}"
        except Exception as e:
            # Catch other potential runtime errors from the user's code.
            return f"An unexpected error occurred: {type(e).__name__}: {e}"