diff --git a/airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py b/airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py index 8a6638fad..5f8ecac14 100644 --- a/airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py +++ b/airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py @@ -1,11 +1,311 @@ -"""Contains functions to compile custom code from text.""" +"""Contains functions to compile custom code from text using RestrictedPython for secure execution.""" +import ast import hashlib import os import sys -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import InitVar, dataclass, field from types import ModuleType -from typing import Any, cast +from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast + +from RestrictedPython import compile_restricted, safe_builtins +from RestrictedPython.compile import RestrictingNodeTransformer +from RestrictedPython.Guards import ( + full_write_guard, + guarded_iter_unpack_sequence, + guarded_unpack_sequence, +) +from RestrictedPython.Guards import ( + safe_builtins as restricted_builtins, +) +from RestrictedPython.Utilities import utility_builtins + + +class AirbyteRestrictingNodeTransformer(RestrictingNodeTransformer): + """Custom AST transformer that allows type annotations and specific private attributes while enforcing security.""" + + ALLOWED_IMPORTS = { + "dataclasses", + "typing", + "requests", + "datetime", + "json", + "airbyte_cdk", + "airbyte_cdk.sources", + "airbyte_cdk.sources.declarative", + "airbyte_cdk.sources.declarative.interpolation", + "airbyte_cdk.sources.declarative.requesters", + "airbyte_cdk.sources.declarative.requesters.paginators", + "airbyte_cdk.sources.declarative.types", + "airbyte_cdk.sources.declarative.types.Config", + "airbyte_cdk.sources.declarative.types.Record", + "InterpolatedString", + "PaginationStrategy", + "Config", + "Record", + "Optional", + "Any", + "Union", + "Mapping", + "Dict", + "List", + "InitVar", + "dataclass", + } + + def visit_Attribute(self, node: ast.Attribute) -> ast.AST: + """Transform attribute access into _getattr_ or _write_ function calls.""" + visited_node = self.generic_visit(node) + if not isinstance(visited_node, ast.AST): + visited_node = node + + if isinstance(visited_node, ast.Attribute) and isinstance(visited_node.attr, str): + # Block access to dangerous attributes + dangerous_attrs = {"__dict__", "__class__", "__bases__", "__subclasses__"} + if visited_node.attr in dangerous_attrs: + raise NameError(f"name '{visited_node.attr}' is not allowed") + + # Allow specific private attributes + allowed_private = { + "__annotations__", + "__name__", + "__doc__", + "__module__", + "__qualname__", + "__post_init__", + "__init__", + "__dataclass_fields__", + "__mro__", + "__subclasshook__", + "__new__", + "_page_size", + } + if visited_node.attr.startswith("_") and visited_node.attr not in allowed_private: + if not visited_node.attr.startswith("__"): # Allow dunder methods + raise NameError(f"name '{visited_node.attr}' is not allowed") + + if isinstance(visited_node.ctx, ast.Store): + # For assignments like "obj.attr = value" + name_node = ast.Name(id="_write_", ctx=ast.Load()) + ast.copy_location(name_node, visited_node) + + value_node = self.visit(visited_node.value) + const_node = ast.Constant(value=visited_node.attr) + + call_node = ast.Call( + func=name_node, + args=[value_node, const_node], + keywords=[], + ) + ast.copy_location(call_node, visited_node) + ast.fix_missing_locations(call_node) + return cast(ast.AST, call_node) + + elif isinstance(visited_node.ctx, ast.Load): + # For reads like "obj.attr" + name_node = ast.Name(id="_getattr_", ctx=ast.Load()) + ast.copy_location(name_node, visited_node) + + const_node = ast.Constant(value=visited_node.attr) + ast.copy_location(const_node, visited_node) + + visited_value = self.visit(visited_node.value) + if hasattr(visited_value, "lineno"): + ast.copy_location(visited_value, visited_node) + + call_node = ast.Call( + func=name_node, + args=[visited_value, const_node], + keywords=[], + ) + ast.copy_location(call_node, visited_node) + ast.fix_missing_locations(call_node) + return cast(ast.AST, call_node) + + elif isinstance(visited_node.ctx, ast.Del): + raise SyntaxError("Attribute deletion is not allowed") + + if not isinstance(visited_node, ast.AST): + raise TypeError(f"Expected ast.AST but got {type(visited_node)}") + return visited_node + + def check_name( + self, + node: ast.AST, + name: str, + *args: Any, + **kwargs: Any, + ) -> ast.AST: + """Allow specific private names that are required for dataclasses and type hints. + + Args: + node: The AST node being checked + name: The name being validated + *args: Additional positional arguments + **kwargs: Additional keyword arguments + """ + if name.startswith("_"): + # Allow specific private names + allowed_private = { + # Type annotation attributes + "__annotations__", + "__name__", + "__doc__", + "__module__", + "__qualname__", + # Dataclass attributes + "__post_init__", + "__init__", + "__dict__", + "__dataclass_fields__", + "__class__", + "__bases__", + "__mro__", + "__subclasshook__", + "__new__", + # Allow specific private attributes used in the codebase + "_page_size", + } + if name in allowed_private or name == "_page_size": + return node + if name.startswith("__"): # Allow dunder methods + return node + raise NameError(f"Name '{name}' is not allowed because it starts with '_'") + return node # Don't call super().check_name as it's too restrictive + + def visit_Import(self, node: ast.Import) -> ast.Import: + """Block unsafe imports.""" + for alias in node.names: + if not alias.name: + raise NameError("__import__ not found") + if not any( + alias.name == allowed or alias.name.startswith(allowed + ".") + for allowed in self.ALLOWED_IMPORTS + ): + raise NameError("__import__ not found") + return node + + def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom: + """Block unsafe imports.""" + module_name = node.module if node.module else "" + + # Handle relative imports + if node.level > 0: + # We don't support relative imports for security + raise NameError("__import__ not found") + + if not any( + module_name == allowed or module_name.startswith(allowed + ".") + for allowed in self.ALLOWED_IMPORTS + ): + raise NameError("__import__ not found") + + # Also check the imported names + for alias in node.names: + if not alias.name: + raise NameError("__import__ not found") + if alias.name == "*": + raise NameError("__import__ not found") + + return node + + def visit_Call(self, node: ast.Call) -> ast.Call: + """Block unsafe function calls.""" + if isinstance(node.func, ast.Name): + unsafe_functions = {"open", "eval", "exec", "compile", "__import__"} + if node.func.id in unsafe_functions: + raise NameError(f"name '{node.func.id}' is not defined") + result: ast.AST = super().visit_Call(node) + if not isinstance(result, ast.Call): + raise TypeError(f"Expected ast.Call but got {type(result)}") + return result + + def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AnnAssign: + """Allow type annotations in variable assignments and dataclass field definitions.""" + # Visit the target and annotation nodes + node.target = self.visit(node.target) + node.annotation = self.visit(node.annotation) + if node.value: + node.value = self.visit(node.value) + return node + + def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: + """Allow dataclass definitions with their attributes.""" + # Check if this is a dataclass by looking for the decorator + is_dataclass = any( + isinstance(d, ast.Name) + and d.id == "dataclass" + or ( + isinstance(d, ast.Call) + and isinstance(d.func, ast.Name) + and d.func.id == "dataclass" + ) + for d in node.decorator_list + ) + + # Visit the decorator list and bases first + node.decorator_list = [self.visit(d) for d in node.decorator_list] + if node.bases: + node.bases = [self.visit(b) for b in node.bases] + + if is_dataclass: + # For dataclasses, we need to allow attribute statements and annotations + allowed_nodes: list[ast.stmt] = [] + for n in node.body: + # Allow class variable annotations (typical in dataclasses) + if isinstance(n, ast.AnnAssign): + visited = self.visit_AnnAssign(n) + if isinstance(visited, ast.stmt): + allowed_nodes.append(visited) + # Allow function definitions (like __post_init__) + elif isinstance(n, ast.FunctionDef): + # For function definitions, we need to allow attribute access + n.body = [ + stmt + for stmt in (self.visit(body_stmt) for body_stmt in n.body) + if isinstance(stmt, ast.stmt) + ] + allowed_nodes.append(n) + # Allow docstrings + elif isinstance(n, ast.Expr) and isinstance(n.value, ast.Str): + allowed_nodes.append(n) + # Allow assignments with type annotations + elif isinstance(n, ast.Assign): + # Convert simple assignments to annotated assignments if possible + if len(n.targets) == 1 and isinstance(n.targets[0], ast.Name): + ann_assign = ast.AnnAssign( + target=n.targets[0], + annotation=ast.Name(id="Any", ctx=ast.Load()), + value=n.value, + simple=1, + ) + visited = self.visit_AnnAssign(ann_assign) + if isinstance(visited, ast.stmt): + allowed_nodes.append(visited) + else: + # Allow attribute assignments within dataclasses + visited = self.visit(n) + if isinstance(visited, ast.stmt): + allowed_nodes.append(visited) + # Allow attribute statements within dataclasses + elif isinstance(n, ast.Attribute): + visited = self.visit_Attribute(n) + if isinstance(visited, ast.stmt): + allowed_nodes.append(visited) + else: + visited = self.visit(n) + if isinstance(visited, ast.stmt): + allowed_nodes.append(visited) + + node.body = allowed_nodes + return node + + result = super().visit_ClassDef(node) + if not isinstance(result, ast.ClassDef): + raise TypeError(f"Expected ast.ClassDef but got {type(result)}") + return result + from typing_extensions import Literal @@ -55,9 +355,13 @@ def _hash_text(input_text: str, hash_type: str = "md5") -> str: def custom_code_execution_permitted() -> bool: """Return `True` if custom code execution is permitted, otherwise `False`. - Custom code execution is permitted if the `AIRBYTE_ALLOW_CUSTOM_CODE` environment variable is set to 'true'. + Custom code execution is permitted if the `AIRBYTE_ALLOW_CUSTOM_CODE` environment variable is set to 'true' + (case-insensitive). Any other value, including '1', 'yes', or empty string, will return False. """ - return os.environ.get(ENV_VAR_ALLOW_CUSTOM_CODE, "").lower() == "true" + env_value = os.environ.get(ENV_VAR_ALLOW_CUSTOM_CODE, "") + if not env_value: + return False + return env_value.lower() in {"true"} def validate_python_code( @@ -121,8 +425,35 @@ def register_components_module_from_string( components_py_text: str, checksums: dict[str, Any] | None, ) -> ModuleType: - """Load and return the components module from a provided string containing the python code.""" - # First validate the code + """Load and return the components module from a provided string containing the python code. + + This function uses RestrictedPython to execute the code in a secure sandbox environment. + The execution is restricted to prevent access to dangerous builtins and operations. + + Security measures: + 1. Code is validated against checksums before execution + 2. Code is compiled using RestrictedPython's compile_restricted + 3. Execution uses safe_builtins to prevent access to dangerous operations + 4. Attribute access is guarded using RestrictedPython's Guards + 5. Code runs in an isolated namespace with restricted globals + + Args: + components_py_text: The Python code to execute as a string. + checksums: Dictionary of checksum types to their expected values. + Must contain at least one of 'md5' or 'sha256'. + + Returns: + ModuleType: A module object containing the executed code's namespace. + + Raises: + AirbyteCodeTamperedError: If the provided code fails checksum validation. + ValueError: If no checksums are provided for validation. + """ + # First check if custom code execution is permitted + if not custom_code_execution_permitted(): + raise AirbyteCustomCodeNotPermittedError() + + # Then validate the code validate_python_code( code_text=components_py_text, checksums=checksums, @@ -131,10 +462,151 @@ def register_components_module_from_string( # Create a new module object components_module = ModuleType(name=COMPONENTS_MODULE_NAME) - # Execute the module text in the module's namespace - exec(components_py_text, components_module.__dict__) + # Create restricted globals with safe builtins + # Start with RestrictedPython's safe builtins and add type annotation support + safe_builtins_copy = dict(safe_builtins) + safe_builtins_copy.update(utility_builtins) # Add utility builtins for basic operations + + # Remove potentially dangerous builtins + dangerous_builtins = { + "open", + "eval", + "exec", + "compile", + "__import__", + "reload", + } # Allow some safe builtins that are needed for type checking and dataclass operations + for name in dangerous_builtins: + safe_builtins_copy.pop(name, None) + + # Add type annotation support + type_support = { + # Type hints + "Any": Any, + "Dict": Dict, + "List": List, + "Tuple": Tuple, + "Set": Set, + "Optional": Optional, + "Union": Union, + "Callable": Callable, + "Mapping": Mapping, + # Basic types + "str": str, + "int": int, + "float": float, + "bool": bool, + # Dataclass support + "dataclass": dataclass, + "InitVar": InitVar, + "field": field, + # Add basic operations + "len": len, + "isinstance": isinstance, + "hasattr": hasattr, + "getattr": getattr, + "ValueError": ValueError, + "TypeError": TypeError, + # Add metaclass support + "__metaclass__": type, + # Add type annotation support + "type": type, + "property": property, + "classmethod": classmethod, + "staticmethod": staticmethod, + # Add requests module + "requests": None, # Will be imported by the code + } + safe_builtins_copy.update(type_support) + + # Define safe attribute access + def safe_getattr(obj: Any, name: str) -> Any: + # Allow type annotation and dataclass related attributes + allowed_private = { + # Type annotation attributes + "__annotations__", + "__name__", + "__doc__", + "__module__", + "__qualname__", + # Dataclass attributes + "__post_init__", + "__init__", + "__dict__", + "__dataclass_fields__", + "__class__", + "__bases__", + "__mro__", + "__subclasshook__", + "__new__", + # Allow specific private attributes used in the codebase + "_page_size", + } + if name in allowed_private or name.startswith("__") or name == "_page_size": + return getattr(obj, name) + # Block access to other special attributes + if name.startswith("_") and name not in allowed_private: + raise AttributeError(f"Access to {name} is not allowed") + return getattr(obj, name) + + # Create restricted globals with support for type annotations and dataclasses + restricted_globals: Dict[str, Any] = { + "__builtins__": safe_builtins_copy, + "_getattr_": safe_getattr, + "_write_": full_write_guard, + "_getiter_": iter, + "_getitem_": lambda obj, key: obj[key] if isinstance(obj, (list, dict, tuple)) else None, + "_print_": lambda *args, **kwargs: None, # No-op print + "__name__": components_module.__name__, + # Add type annotation and dataclass support to globals + "Any": Any, + "Dict": Dict, + "List": List, + "Tuple": Tuple, + "Set": Set, + "Optional": Optional, + "Union": Union, + "Callable": Callable, + "Mapping": Mapping, + "dataclass": dataclass, + "InitVar": InitVar, + "field": field, + # Add sequence unpacking support + "_unpack_sequence_": guarded_unpack_sequence, + "_iter_unpack_sequence_": guarded_iter_unpack_sequence, + # Add support for type annotations + "__annotations__": {}, + "__module__": components_module.__name__, + "__qualname__": "", + "__doc__": None, + "__metaclass__": type, + # Add support for requests module + "requests": None, # Will be imported by the code + # Add support for PaginationStrategy + "PaginationStrategy": None, # Will be imported by the code + "InterpolatedString": None, # Will be imported by the code + "Config": None, # Will be imported by the code + "Record": None, # Will be imported by the code + } + + # Compile with RestrictedPython's restrictions using our custom transformer + try: + byte_code = compile_restricted( + components_py_text, + filename="", + mode="exec", + policy=AirbyteRestrictingNodeTransformer, + ) + except SyntaxError as e: + raise SyntaxError(f"Restricted execution error: {str(e)}") + + # Execute the compiled code in the restricted environment + exec(byte_code, restricted_globals) + + # Update the module's dictionary with the restricted execution results + components_module.__dict__.update(restricted_globals) - # Register the module in `sys.modules`` so it can be imported as + # Register the module in `sys.modules` so it can be imported as # `source_declarative_manifest.components` and/or `components`. sys.modules[SDM_COMPONENTS_MODULE_NAME] = components_module sys.modules[COMPONENTS_MODULE_NAME] = components_module diff --git a/poetry.lock b/poetry.lock index 3f5f8c634..aab7bb8a2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1729,10 +1729,7 @@ files = [ [package.dependencies] httpx = ">=0.23.0,<1" orjson = {version = ">=3.9.14,<4.0.0", markers = "platform_python_implementation != \"PyPy\""} -pydantic = [ - {version = ">=1,<3", markers = "python_full_version < \"3.12.4\""}, - {version = ">=2.7.4,<3.0.0", markers = "python_full_version >= \"3.12.4\""}, -] +pydantic = {version = ">=1,<3", markers = "python_full_version < \"3.12.4\""} requests = ">=2,<3" requests-toolbelt = ">=1.0.0,<2.0.0" @@ -2635,7 +2632,6 @@ files = [ numpy = [ {version = ">=1.22.4", markers = "python_version < \"3.11\""}, {version = ">=1.23.2", markers = "python_version == \"3.11\""}, - {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, ] python-dateutil = ">=2.8.2" pytz = ">=2020.1" @@ -4100,6 +4096,21 @@ files = [ [package.dependencies] requests = ">=2.0.1,<3.0.0" +[[package]] +name = "restrictedpython" +version = "6.2" +description = "RestrictedPython is a defined subset of the Python language which allows to provide a program input into a trusted environment." +optional = false +python-versions = ">=3.6, <3.12" +files = [ + {file = "RestrictedPython-6.2-py3-none-any.whl", hash = "sha256:7c2ffa4904300d67732f841d8a975dcdc53eba4c1cdc9d84b97684ef12304a3d"}, + {file = "RestrictedPython-6.2.tar.gz", hash = "sha256:db73eb7e3b39650f0d21d10cc8dda9c0e2986e621c94b0c5de32fb0dee3a08af"}, +] + +[package.extras] +docs = ["Sphinx", "sphinx-rtd-theme"] +test = ["pytest", "pytest-mock"] + [[package]] name = "rich" version = "13.9.4" @@ -5029,5 +5040,5 @@ vector-db-based = ["cohere", "langchain", "openai", "tiktoken"] [metadata] lock-version = "2.0" -python-versions = "^3.10,<3.13" -content-hash = "4ec21ccda2c77f36951f87d7d756746a40db799716f09492fcb9d39c1de35563" +python-versions = "^3.10,<3.12" +content-hash = "3d4a75a906eb1d8396734f1ac723420a79b09797dc13b4d36550ec05337f8e7f" diff --git a/pyproject.toml b/pyproject.toml index ea48db48e..f7a78330e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,9 @@ version = "0.0.0" # Version will be calculated dynamically. enable = true [tool.poetry.dependencies] -python = "^3.10,<3.13" +python = "^3.10,<3.12" airbyte-protocol-models-dataclasses = "^0.14" +RestrictedPython = "^6.0.0" backoff = "*" cachetools = "*" dpath = "^2.1.6" diff --git a/unit_tests/source_declarative_manifest/test_source_declarative_w_custom_components.py b/unit_tests/source_declarative_manifest/test_source_declarative_w_custom_components.py index d608e7620..37962281d 100644 --- a/unit_tests/source_declarative_manifest/test_source_declarative_w_custom_components.py +++ b/unit_tests/source_declarative_manifest/test_source_declarative_w_custom_components.py @@ -48,7 +48,10 @@ def get_fixture_path(file_name) -> str: return os.path.join(os.path.dirname(__file__), file_name) -def test_components_module_from_string() -> None: +def test_components_module_from_string(monkeypatch: pytest.MonkeyPatch) -> None: + # Enable custom code execution for this test + monkeypatch.setenv(ENV_VAR_ALLOW_CUSTOM_CODE, "true") + # Call the function to get the module components_module: types.ModuleType = register_components_module_from_string( components_py_text=SAMPLE_COMPONENTS_PY_TEXT, @@ -185,8 +188,8 @@ def test_invalid_checksum_fails_to_run( "env_value, should_raise", [ ("true", False), - ("True", False), - ("TRUE", False), + ("True", False), # Case-insensitive comparison + ("TRUE", False), # Case-insensitive comparison ("1", True), # Not accepted as truthy as of now ("false", True), ("False", True), diff --git a/unit_tests/sources/declarative/parsers/test_custom_code_compiler.py b/unit_tests/sources/declarative/parsers/test_custom_code_compiler.py new file mode 100644 index 000000000..ca41a12be --- /dev/null +++ b/unit_tests/sources/declarative/parsers/test_custom_code_compiler.py @@ -0,0 +1,118 @@ +"""Unit tests for custom code compiler with RestrictedPython security features.""" + +import os +from typing import Any, Dict +from unittest.mock import patch + +import pytest + +from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import ( + AirbyteCodeTamperedError, + AirbyteCustomCodeNotPermittedError, + register_components_module_from_string, + validate_python_code, +) + + +def test_validate_python_code_with_valid_md5(): + """Test that code validation passes with correct MD5 checksum.""" + code = "def test(): return 'hello'" + checksums = {"md5": "8901edeabbb26c1d0496c2c38a95cf17"} + validate_python_code(code, checksums) # Should not raise + + +def test_validate_python_code_with_valid_sha256(): + """Test that code validation passes with correct SHA256 checksum.""" + code = "def test(): return 'hello'" + checksums = {"sha256": "bc379fc64b5ea5d0bb232194b4ce6be0bc16de0d30b33c069a2d63eb911e74b0"} + validate_python_code(code, checksums) # Should not raise + + +def test_validate_python_code_with_invalid_checksum(): + """Test that code validation fails with incorrect checksum.""" + code = "def test(): return 'hello'" + checksums = {"md5": "invalid"} + with pytest.raises(AirbyteCodeTamperedError): + validate_python_code(code, checksums) + + +def test_validate_python_code_with_no_checksums(): + """Test that code validation fails when no checksums are provided.""" + code = "def test(): return 'hello'" + with pytest.raises(ValueError, match="A checksum is required"): + validate_python_code(code, None) + + +def test_register_components_module_safe_code(): + """Test that safe code executes successfully in restricted environment.""" + code = """ +def get_value(): + return 42 + +def add_numbers(a, b): + return a + b +""" + checksums = {"md5": "8c00db73237e4ba737003dc78fc5e63f"} + with patch.dict(os.environ, {"AIRBYTE_ALLOW_CUSTOM_CODE": "true"}): + module = register_components_module_from_string(code, checksums) + assert module.get_value() == 42 + assert module.add_numbers(2, 3) == 5 + + +def test_register_components_module_unsafe_imports(): + """Test that unsafe module imports are blocked.""" + code = """ +import os +def delete_file(): + os.remove('/tmp/test') +""" + checksums = {"md5": "a40d238c0ae2c750df62a45bbaf344a2"} + with patch.dict(os.environ, {"AIRBYTE_ALLOW_CUSTOM_CODE": "true"}): + with pytest.raises(Exception) as exc_info: + register_components_module_from_string(code, checksums) + error_msg = str(exc_info.value) + assert any( + msg in error_msg for msg in ["__import__ not found", "name '__import__' is not defined"] + ) + + +def test_register_components_module_unsafe_builtins(): + """Test that unsafe builtin operations are blocked.""" + code = """ +def evil_code(): + open('/etc/passwd', 'r').read() +""" + checksums = {"md5": "62c7d65594a5b7654ff6456125483c05"} + with patch.dict(os.environ, {"AIRBYTE_ALLOW_CUSTOM_CODE": "true"}): + with pytest.raises((NameError, AttributeError)) as exc_info: + module = register_components_module_from_string(code, checksums) + # If compilation succeeds, try to execute the code which should fail + if hasattr(module, "evil_code"): + module.evil_code() + error_msg = str(exc_info.value) + assert any( + msg in error_msg.lower() for msg in ["name 'open' is not defined", "open not found"] + ) + + +def test_custom_code_execution_not_permitted(): + """Test that code execution is blocked when environment variable is not set.""" + code = "def test(): return 42" + checksums = {"md5": "e26ee9f0888fd40cc4d2264d49057bef"} + with patch.dict(os.environ, {"AIRBYTE_ALLOW_CUSTOM_CODE": "false"}): + with pytest.raises(AirbyteCustomCodeNotPermittedError): + register_components_module_from_string(code, checksums) + + +def test_register_components_module_restricted_attributes(): + """Test that accessing restricted attributes is blocked.""" + code = """ +class Evil: + def __init__(self): + self.__dict__ = {} +""" + checksums = {"md5": "168fc66811e175f26a8cedb02aa723a4"} + with patch.dict(os.environ, {"AIRBYTE_ALLOW_CUSTOM_CODE": "true"}): + with pytest.raises(Exception) as exc_info: + register_components_module_from_string(code, checksums) + assert "__dict__" in str(exc_info.value)