test_poc_vulnerability.py

"""
Test suite for CodeSafe security vulnerabilities.
Tests the fix for the sandbox escape vulnerability and other potential attack vectors.
"""
import unittest
import os
from pathlib import Path
from codesafe import safe_eval, UnsafeExpressionError


# Top-level helper function for multiprocessing pickling compatibility on Windows
def custom_func(x):
    return x * 2


class TestVulnerabilityFixes(unittest.TestCase):
    """Test cases for security vulnerabilities and their fixes."""

    def setUp(self):
        """Set up test fixtures."""
        self.marker_file = Path(__file__).parent / "codesafe_rce_marker.txt"
        if self.marker_file.exists():
            self.marker_file.unlink()

    def tearDown(self):
        """Clean up test files."""
        if self.marker_file.exists():
            self.marker_file.unlink()

    # ==================== MAIN VULNERABILITY TESTS ====================

    def test_poc_direct_exec_blocked(self):
        """Test that direct exec() calls are blocked (negative control)."""
        inner_code = (
            "from pathlib import Path\n"
            f"Path({str(self.marker_file)!r}).write_text("
            "'codesafe-rce', encoding='utf-8')"
        )
        payload = f"exec({inner_code!r})"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        self.assertFalse(self.marker_file.exists())
        print("[PASS] Direct exec is correctly blocked")

    def test_poc_subscript_bypass_fixed(self):
        """Test the main PoC: __builtins__['exec']() bypass is now fixed."""
        inner_code = (
            "from pathlib import Path\n"
            f"Path({str(self.marker_file)!r}).write_text("
            "'codesafe-rce', encoding='utf-8')"
        )
        payload = f"__builtins__['exec']({inner_code!r})"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        self.assertFalse(self.marker_file.exists())
        print("[PASS] Subscript bypass (__builtins__['exec']) is now blocked")

    def test_poc_getattr_exec_blocked(self):
        """Test getattr(__builtins__, 'exec') bypass is blocked."""
        inner_code = (
            "from pathlib import Path\n"
            f"Path({str(self.marker_file)!r}).write_text('getattr-rce')"
        )
        payload = f"getattr(__builtins__, 'exec')({inner_code!r})"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        self.assertFalse(self.marker_file.exists())
        print("[PASS] getattr exec bypass is blocked")

    def test_poc_eval_via_builtins_blocked(self):
        """Test __builtins__['eval']() is blocked."""
        payload = "__builtins__['eval']('1 + 1')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] __builtins__['eval'] is blocked")

    def test_poc_compile_via_builtins_blocked(self):
        """Test __builtins__['compile']() is blocked."""
        payload = "__builtins__['compile']('1+1', '<string>', 'eval')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] __builtins__['compile'] is blocked")

    def test_poc_import_via_builtins_blocked(self):
        """Test __builtins__['__import__']() is blocked."""
        payload = "__builtins__['__import__']('os')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] __builtins__['__import__'] is blocked")

    # ==================== INDIRECT CALL BYPASS TESTS ====================

    def test_lambda_call_bypass_blocked(self):
        """Test lambda-based call bypass is blocked."""
        payload = "(lambda: __import__('os'))()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Lambda call bypass is blocked")

    def test_nested_subscript_bypass_blocked(self):
        """Test nested subscript access is blocked."""
        payload = "__builtins__['dict']['keys']()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Nested subscript bypass is blocked")

    def test_attribute_then_call_blocked(self):
        """Test attribute access followed by call is blocked."""
        payload = "__builtins__.get('exec')('1+1')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Attribute-then-call bypass is blocked")

    def test_list_comprehension_call_blocked(self):
        """Test list comprehension with function call is blocked."""
        payload = "[__builtins__['exec']('1') for x in [1]]"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] List comprehension call bypass is blocked")

    def test_generator_expression_call_blocked(self):
        """Test generator expression with function call is blocked."""
        payload = "list(__builtins__['exec']('1') for x in [1])"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Generator expression call bypass is blocked")

    # ==================== DUNDER METHOD ATTACKS ====================

    def test_class_dunder_access_blocked(self):
        """Test __class__ attribute access is blocked."""
        payload = "''.__class__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __class__ access is blocked")

    def test_bases_dunder_access_blocked(self):
        """Test __bases__ attribute access is blocked."""
        payload = "().__class__.__bases__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __bases__ access is blocked")

    def test_subclasses_dunder_access_blocked(self):
        """Test __subclasses__ attribute access is blocked."""
        payload = "''.__class__.__mro__[1].__subclasses__()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __subclasses__ access is blocked")

    def test_globals_dunder_access_blocked(self):
        """Test __globals__ attribute access is blocked."""
        payload = "(lambda: None).__globals__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __globals__ access is blocked")

    def test_closure_dunder_access_blocked(self):
        """Test __closure__ attribute access is blocked."""
        payload = "(lambda: None).__closure__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __closure__ access is blocked")

    def test_code_dunder_access_blocked(self):
        """Test __code__ attribute access is blocked."""
        payload = "(lambda: None).__code__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __code__ access is blocked")

    # ==================== IMPORT AND MODULE ACCESS TESTS ====================

    def test_direct_import_blocked(self):
        """Test direct import statements are blocked."""
        payload = "import os"
        
        try:
            safe_eval(payload)
            self.fail("Should have raised SyntaxError")
        except SyntaxError:
            pass
        
        print("[PASS] Direct import is blocked (syntax error)")

    def test_from_import_blocked(self):
        """Test from-import statements are blocked."""
        payload = "from os import system"
        
        try:
            safe_eval(payload)
            self.fail("Should have raised SyntaxError")
        except SyntaxError:
            pass
        
        print("[PASS] From-import is blocked (syntax error)")

    def test_os_system_via_builtins_blocked(self):
        """Test accessing os.system via builtins is blocked."""
        payload = "__builtins__['__import__']('os').system('echo pwned')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] os.system via __import__ is blocked")

    def test_subprocess_call_blocked(self):
        """Test subprocess module access is blocked."""
        payload = "__builtins__['__import__']('subprocess').call(['ls'])"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] subprocess.call is blocked")

    # ==================== FILE ACCESS TESTS ====================

    def test_open_function_blocked_by_default(self):
        """Test open() is blocked when file_access=False."""
        payload = "open('/etc/passwd').read()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, file_access=False)
        
        print("[PASS] open() is blocked by default")

    def test_pathlib_write_blocked(self):
        """Test pathlib.Path write operations are blocked."""
        marker = Path(__file__).parent / "pathlib_test.txt"
        if marker.exists():
            marker.unlink()
        
        payload = f"from pathlib import Path; Path({str(marker)!r}).write_text('test')"
        
        try:
            safe_eval(payload)
            self.fail("Should have raised SyntaxError")
        except SyntaxError:
            pass
        
        print("[PASS] pathlib write is blocked (syntax error)")

    def test_io_module_blocked(self):
        """Test io module access is blocked."""
        payload = "__builtins__['__import__']('io').open('/tmp/test')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] io.open is blocked")

    # ==================== NETWORK ACCESS TESTS ====================

    def test_socket_module_blocked(self):
        """Test socket module access is blocked."""
        payload = "__builtins__['__import__']('socket').socket()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] socket module is blocked")

    def test_requests_module_blocked(self):
        """Test requests module access is blocked."""
        payload = "__builtins__['__import__']('requests').get('http://example.com')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] requests module is blocked")

    def test_urllib_module_blocked(self):
        """Test urllib module access is blocked."""
        payload = "__builtins__['__import__']('urllib.request').urlopen('http://example.com')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] urllib module is blocked")

    # ==================== CODE EXECUTION PRIMITIVES ====================

    def test_exec_function_blocked(self):
        """Test exec() function call is blocked."""
        payload = "exec('print(1)')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] exec() is blocked")

    def test_eval_function_blocked(self):
        """Test eval() function call is blocked."""
        payload = "eval('1+1')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] eval() is blocked")

    def test_compile_function_blocked(self):
        """Test compile() function call is blocked."""
        payload = "compile('1+1', '<string>', 'eval')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] compile() is blocked")

    def test_input_function_behavior(self):
        """Test input() behavior (should be blocked as it's not in safe builtins)."""
        payload = "input()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] input() is blocked")

    # ==================== METACLASS AND TYPE ATTACKS ====================

    def test_type_subclasses_blocked(self):
        """Test type().subclasses() attack is blocked."""
        payload = "type('')[0].__subclasses__()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] type subclasses attack is blocked")

    def test_object_subclasses_blocked(self):
        """Test object.__subclasses__() attack is blocked."""
        payload = "object().__subclasses__()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] object.__subclasses__ is blocked")

    def test_mro_access_blocked(self):
        """Test __mro__ access is blocked."""
        payload = "''.__class__.__mro__"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload, allow_attributes=True)
        
        print("[PASS] __mro__ access is blocked")

    # ==================== BUILTIN MODIFICATION TESTS ====================

    def test_builtins_dict_modification_blocked(self):
        """Test modifying __builtins__ dict is blocked."""
        payload = "__builtins__['malicious'] = lambda: 'pwned'"
        
        try:
            safe_eval(payload)
            self.fail("Should have raised SyntaxError")
        except SyntaxError:
            pass
        
        print("[PASS] __builtins__ modification is blocked (syntax error)")

    def test_builtins_update_blocked(self):
        """Test __builtins__.update() is blocked."""
        payload = "__builtins__.update({'evil': eval})"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] __builtins__.update() is blocked")

    # ==================== SAFE OPERATIONS STILL WORK ====================

    def test_safe_arithmetic_works(self):
        """Test that safe arithmetic operations still work."""
        result = safe_eval("2 + 2 * 3")
        self.assertEqual(result, 8)
        print("[PASS] Safe arithmetic works")

    def test_safe_string_operations_work(self):
        """Test that safe string operations still work."""
        result = safe_eval("'hello'.upper()", allow_attributes=True)
        self.assertEqual(result, "HELLO")
        print("[PASS] Safe string operations work")

    def test_safe_list_operations_work(self):
        """Test that safe list operations still work."""
        result = safe_eval("[1, 2, 3][1]")
        self.assertEqual(result, 2)
        print("[PASS] Safe list operations work")

    def test_safe_dict_operations_work(self):
        """Test that safe dict operations still work."""
        result = safe_eval("{'a': 1}['a']")
        self.assertEqual(result, 1)
        print("[PASS] Safe dict operations work")

    def test_safe_builtin_functions_work(self):
        """Test that safe builtin functions still work."""
        result = safe_eval("max([1, 5, 3])")
        self.assertEqual(result, 5)
        
        result = safe_eval("abs(-5)")
        self.assertEqual(result, 5)
        
        result = safe_eval("len([1, 2, 3])")
        self.assertEqual(result, 3)
        
        print("[PASS] Safe builtin functions work")

    def test_allowed_function_calls_work(self):
        """Test that explicitly allowed function calls work."""
        result = safe_eval(
            "custom_func(5)",
            allowed_vars={"custom_func": custom_func},
            allowed_function_calls=["custom_func"]
        )
        self.assertEqual(result, 10)
        print("[PASS] Allowed function calls work")


class TestEdgeCases(unittest.TestCase):
    """Test edge cases and complex attack vectors."""

    def test_multiple_indirect_lookups(self):
        """Test multiple levels of indirect lookup."""
        payload = "__builtins__['dict']()"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Multiple indirect lookups blocked")

    def test_tuple_unpacking_call_blocked(self):
        """Test tuple unpacking with call is blocked."""
        payload = "(*[__builtins__['exec']('1')],)"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Tuple unpacking call blocked")

    def test_dict_comprehension_call_blocked(self):
        """Test dict comprehension with call is blocked."""
        payload = "{x: __builtins__['exec']('1') for x in [1]}"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Dict comprehension call blocked")

    def test_set_comprehension_call_blocked(self):
        """Test set comprehension with call is blocked."""
        payload = "{__builtins__['exec']('1') for x in [1]}"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Set comprehension call blocked")

    def test_ternary_with_call_blocked(self):
        """Test ternary operator with call is blocked."""
        payload = "__builtins__['exec']('1') if True else None"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Ternary with call blocked")

    def test_and_or_shortcircuit_call_blocked(self):
        """Test short-circuit evaluation with call is blocked."""
        payload = "True and __builtins__['exec']('1')"
        
        with self.assertRaises(UnsafeExpressionError):
            safe_eval(payload)
        
        print("[PASS] Short-circuit call blocked")

    def test_walrus_operator_call_blocked(self):
        """Test walrus operator with call is blocked."""
        try:
            payload = "(x := __builtins__['exec']('1'))"
            with self.assertRaises(UnsafeExpressionError):
                safe_eval(payload)
            print("[PASS] Walrus operator call blocked")
        except SyntaxError:
            print("[SKIP] Walrus operator not available in this Python version")


if __name__ == "__main__":
    unittest.main(verbosity=2)