Patched code sample

import ast
import sys

# A whitelist of dunder attributes that are considered safe to access.
# In a real-world scenario, this list would be carefully curated.
SAFE_DUNDERS = {'__name__', '__class__'}

class DunderValidator(ast.NodeVisitor):
    """
    AST visitor that checks for access to potentially dangerous dunder attributes.:
    """
    def visit_Attribute(self, node):
        """Secure function that fixes the vulnerability."""
        if isinstance(node.ctx, ast.Load):
            attr_name = node.attr
            if attr_name.startswith('__') and attr_name.endswith('__'):
                if attr_name not in SAFE_DUNDERS:
                    raise PermissionError(
                    f"Sandbox escape attempt: Access to dunder attribute '{attr_name}' is forbidden."
                    )
                    self.generic_visit(node)

                    def secure_execute_code(code_string: str, allowed_globals: dict = None):
                        """
                        Executes Python code after validating its AST to prevent dunder attribute abuse.

                        This function represents the "fixed" version of the execution environment.
                        It parses the code into an Abstract Syntax Tree (AST) and uses a visitor
                        to traverse the tree, ensuring no forbidden dunder attributes are accessed.
                        """
                        if allowed_globals is None:
        # Provide a very limited global scope.
                            allowed_globals = {"__builtins__": {"print": print, "range": range, "list": list}}

                            try:
        # 1. Parse the code into an AST.
                                tree = ast.parse(code_string)

        # 2. Validate the AST for forbidden dunder attribute access.:
                                validator = DunderValidator()
                                validator.visit(tree)

        # 3. If validation passes, compile and execute the code.
                                print(f"--- AST Validation Passed for: ---\n{code_string}\n----------------------------------")
                                code_obj = compile(tree, '<string>', 'exec')
                                exec(code_obj, allowed_globals)

                                except (PermissionError, SyntaxError, Exception) as e:
                                    print(f"--- Execution Blocked: {type(e).__name__} ---")
                                    print(f"Reason: {e}")
                                    print("-----------------------------------------")


                                    def vulnerable_execute_code(code_string: str, allowed_globals: dict = None):
                                        """
                                        Executes Python code with a naive blocklist, which is easily bypassed.

                                        This function represents the "vulnerable" environment. It only checks for
                                        obvious dangerous keywords in the string, but fails to prevent sophisticated:
                                        attacks that use dunder attributes to access restricted objects.
                                        """
                                        if allowed_globals is None:
                                            allowed_globals = {"__builtins__": {"print": print}}

    # Naive, incomplete validation (The Vulnerability)
                                            dangerous_keywords = ['import', 'os', 'subprocess', 'eval', 'exec', 'open']
                                            if any(keyword in code_string for keyword in dangerous_keywords):
                                                print("--- Execution Blocked by Naive Check ---")
                                                return

                                                print(f"--- Naive Check Passed, Executing Vulnerable Code: ---\n{code_string}\n------------------------------------------------------")
                                                try:
                                                    exec(code_string, allowed_globals)
                                                    except Exception as e:
                                                        print(f"Execution failed: {e}")


                                                        if __name__ == '__main__':
    # This payload uses a chain of dunder attributes to escape the sandbox.
    # It finds the `object` class, lists all its subclasses, finds a class that
    # can import modules (`warnings.catch_warnings`), and uses it to import `os`
    # and run a command.
                                                            malicious_payload = (
                                                            "print('Attempting sandbox escape...'); "
                                                            "importer = [c for c in ().__class__.__base__.__subclasses__() ":
                                                            "if c.__name__ == 'catch_warnings'][0]()._module.__builtins__['__import__']; ":
                                                            "os_module = importer('os'); "
                                                            "os_module.system('echo PAYLOAD EXECUTED: Sandbox escaped successfully!')"
                                                            )
    
    # A benign piece of code that should be allowed to run
                                                            safe_payload = "print('This is a safe operation.')"

                                                            print("=========================================================")
                                                            print("DEMONSTRATING THE VULNERABLE EXECUTION ENVIRONMENT")
                                                            print("=========================================================\n")
    
                                                            print("--- Testing vulnerable environment with safe code ---")
                                                            vulnerable_execute_code(safe_payload)
                                                            print("\n--- Testing vulnerable environment with malicious payload ---")
                                                            vulnerable_execute_code(malicious_payload)

                                                            print("\n\n=========================================================")
                                                            print("DEMONSTRATING THE FIXED (SECURE) EXECUTION ENVIRONMENT")
                                                            print("=========================================================\n")

                                                            print("--- Testing fixed environment with safe code ---")
                                                            secure_execute_code(safe_payload)
                                                            print("\n--- Testing fixed environment with malicious payload ---")
    # This call will be blocked by the DunderValidator in the secure_execute_code function
                                                            secure_execute_code(malicious_payload)