Patched code sample

import ast
import operator

# The fictitious CVE-2026-4963 describes a code injection vulnerability in a
# component responsible for executing Python code from an Abstract Syntax Tree (AST).
# The vulnerability stems from an incomplete fix for a prior issue, where
# an attacker can bypass checks on what functions or operations are allowed.
#
# A secure fix involves abandoning blocklists or naive string-based checks.
# Instead, a strict allowlist-based approach must be used, where every part of
# a code expression (e.g., 'os.system') is validated against a predefined set of
# safe, approved objects. The code below demonstrates this principle.

class FixedLocalPythonExecutor:
    """
    A representation of a local Python executor that safely evaluates
    AST nodes, fixing the described code injection vulnerability.
    """

    def __init__(self, allowed_modules=None, allowed_functions=None):
        """
        Initializes the executor with a strict allowlist of what is safe to run.
        """
        # --- FIX: Use a strict, object-based allowlist, not string names. ---
        # This maps safe-to-call names to their actual Python objects.
        self.allowed_scope = {
            'print': print,
            'min': min,
            'max': max,
            # Represent safe "modules" as nested dictionaries.
            'math': {
                'sqrt': operator.pow,  # Example: map to a safe implementation
            }
        }
        # For demonstration, we can merge custom allowed items.
        if allowed_modules:
            self.allowed_scope.update(allowed_modules)
        if allowed_functions:
            self.allowed_scope.update(allowed_functions)

    def _safely_resolve_node(self, node, current_scope):
        """
        Recursively resolves an AST node (e.g., a.b.c) against the allowed scope.
        This is the core of the security fix. It ensures that every part of an
        attribute access chain is explicitly permitted.
        """
        if isinstance(node, ast.Name):
            # Base case: an identifier like 'print' or 'math'.
            if node.id in current_scope:
                return current_scope[node.id]
        elif isinstance(node, ast.Attribute):
            # Recursive case: an access like 'math.sqrt'.
            # First, resolve the base ('math').
            base_object = self._safely_resolve_node(node.value, current_scope)
            # The resolved base must be a dictionary for further attribute access.
            if isinstance(base_object, dict) and node.attr in base_object:
                # Then, resolve the attribute ('sqrt') within the base's scope.
                return base_object[node.attr]
        elif isinstance(node, ast.Constant):
            # Allow literal values like numbers and strings.
            return node.value

        # --- FIX: Deny by default. ---
        # Any construct not explicitly matched and validated is forbidden.
        # This prevents bypassing checks with constructs like `__import__` or `getattr`.
        raise PermissionError(f"Access to '{ast.unparse(node)}' is not allowed.")

    def evaluate_call(self, node: ast.Call):
        """
        Safely evaluates an `ast.Call` node. This function was specified as
        vulnerable in the CVE.
        """
        # 1. Safely resolve the function being called. This will traverse the
        #    AST for the callable (e.g., `math.sqrt`) and ensure each part is allowed.
        callable_func = self._safely_resolve_node(node.func, self.allowed_scope)

        # 2. Verify that the resolved object is actually a function we can call.
        if not callable(callable_func):
            raise TypeError(f"'{ast.unparse(node.func)}' did not resolve to a callable function.")

        # 3. Safely resolve all arguments.
        #    A real implementation would recursively call a master evaluation function.
        #    For this example, we'll use the same resolver.
        args = [self._safely_resolve_node(arg, self.allowed_scope) for arg in node.args]
        kwargs = {kw.arg: self._safely_resolve_node(kw.value, self.allowed_scope) for kw in node.keywords}

        # 4. Execute the call with the validated function and arguments.
        return callable_func(*args, **kwargs)

    def evaluate_augassign(self, node: ast.AugAssign):
        """
        Safely evaluates an `ast.AugAssign` node (e.g., `x += 1`).
        This function was specified as vulnerable in the CVE.
        The fix is to deny this operation entirely if it's not part of the
        required feature set, which is the most secure approach.
        """
        # --- FIX: Deny complex state-modifying operations by default. ---
        raise PermissionError("In-place assignments (e.g., +=) are not permitted.")

    def evaluate_with(self, node: ast.With):
        """
        Safely evaluates an `ast.With` node.
        This function was specified as vulnerable in the CVE.
        The `with` statement can be dangerous (e.g., `with open(...)`), so the
        safest fix is to deny it unless a specific, safe context manager is allowed.
        """
        # --- FIX: Deny `with` statements by default. ---
        raise PermissionError("The 'with' statement is not permitted.")

    def execute_safely(self, code_string: str):
        """
        Main entry point to parse and execute code safely.
        """
        try:
            # We parse the code into an AST.
            tree = ast.parse(code_string, mode='exec')
            # A secure executor would walk the entire tree and only allow
            # a small subset of node types (Call, Constant, Name, etc.).
            # For this demonstration, we focus only on a single call expression.
            if len(tree.body) == 1 and isinstance(tree.body[0], ast.Expr):
                node = tree.body[0].value
                if isinstance(node, ast.Call):
                    return self.evaluate_call(node)

            raise SyntaxError("Execution is restricted to a single, allowed function call.")
        except (PermissionError, SyntaxError, TypeError) as e:
            # Catch our specific security exceptions.
            print(f"Execution blocked: {e}")
            return None
        except Exception as e:
            # Catch other potential parsing/evaluation errors.
            print(f"An unexpected error occurred: {e}")
            return None