Patched code sample

import ast
import sys

# A set of dangerous attributes that should not be accessed in a sandbox.
# Access to these can lead to sandbox escapes and arbitrary code execution.
BLOCKED_ATTRIBUTES = {
    '__subclasses__',
    '__bases__',
    '__globals__',
    '__code__',
    '__closure__',
    '__func__',
    '__self__',
    '__mro__',
    '__init__',
    '__dict__',
    '__getattribute__',
    'getattr',
    'setattr'
}

# A set of functions/methods known to be used for dynamic attribute access.
# The fix involves monitoring calls to these functions.
DYNAMIC_ACCESS_FUNCTIONS = {
    '__getattribute__',
    'getattr'
}

class PatchedAstValidator(ast.NodeVisitor):
    """
    An AST visitor that enforces security constraints on Python code.
    
    This class represents the fix for CVE-2026-40158 by inspecting not only
    direct attribute access (`ast.Attribute`) but also dynamic attribute access
    via function calls (`ast.Call`) where dangerous attribute names are passed
    as string constants (`ast.Constant`).
    """

    def visit_Attribute(self, node: ast.Attribute):
        """
        Blocks direct access to forbidden attributes.
        Example: blocks code like `some_object.__globals__`.
        """
        if node.attr in BLOCKED_ATTRIBUTES:
            raise PermissionError(
                f"Direct access to the dangerous attribute '{node.attr}' is forbidden."
            )
        # Continue traversing the rest of the node's children.
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call):
        """
        **This method contains the core of the vulnerability fix.**
        
        It blocks dynamic access to forbidden attributes passed as string literals.
        The original PraisonAI sandbox only checked `ast.Attribute` nodes, but the
        exploit bypassed this by using an `ast.Constant` (the string '__subclasses__')
        as an argument to a function like `type.__getattribute__`.

        This patched method inspects calls to sensitive functions and checks their
        string arguments against the blocklist, closing the loophole.
        """
        # Check if the function being called is an attribute itself (e.g., `type.getattribute`).
        if isinstance(node.func, ast.Attribute):
            function_name = node.func.attr
            
            # Check if it's a function we are monitoring for dynamic attribute access.
            if function_name in DYNAMIC_ACCESS_FUNCTIONS:
                # The sensitive attribute name is typically the second argument.
                if len(node.args) > 1:
                    arg = node.args[1]
                    attr_name = None

                    # In Python 3.8+, string literals are `ast.Constant`.
                    if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                        attr_name = arg.value
                    # In older Python versions (<3.8), they were `ast.Str`.
                    elif sys.version_info < (3, 8) and isinstance(arg, ast.Str):
                        attr_name = arg.s

                    # If the extracted string argument is a blocked attribute, raise an error.
                    if attr_name and attr_name in BLOCKED_ATTRIBUTES:
                        raise PermissionError(
                            f"Dynamic access to '{attr_name}' via the function "
                            f"'{function_name}' is forbidden."
                        )
        
        # Continue traversing the rest of the node's children.
        self.generic_visit(node)


def validate_code_with_fix(untrusted_code: str):
    """
    A wrapper function demonstrating how to use the PatchedAstValidator to
    check if a string of code is safe. In a real application, this would be
    called before executing any untrusted agent code.
    
    Args:
        untrusted_code: The string of Python code from an untrusted source.

    Raises:
        PermissionError: If the code contains forbidden patterns.
        SyntaxError: If the code is not valid Python.
    """
    tree = ast.parse(untrusted_code)
    validator = PatchedAstValidator()
    validator.visit(tree)
    # If no exception is raised, the code is considered safe by the validator.
    # In a real application, you would now proceed to execute the code
    # in a sandboxed environment.