import ast import os import sys import textwrap from types import ModuleType from typing import Optional from typing import Text from ddtrace.appsec._constants import IAST from ddtrace.appsec._iast._ast import iastpatch from ddtrace.appsec._iast._logs import iast_ast_debug_log from ddtrace.appsec._iast._logs import iast_compiling_debug_log from ddtrace.appsec._iast._logs import iast_instrumentation_ast_patching_debug_log from ddtrace.internal.logger import get_logger from ddtrace.internal.module import origin from ddtrace.internal.settings.asm import config as asm_config from ddtrace.internal.utils.formats import asbool from .visitor import AstVisitor _VISITOR = AstVisitor() _PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX IAST_PATCHING_LAZY_LOADED = True log = get_logger(__name__) def initialize_iast_lists(): """Initialize IAST module lists safely from Python. This function initializes the user allowlist and denylist for IAST module patching. It is critical that this initialization happens from Python rather than from C code during module initialization for several reasons: 1. Python GIL (Global Interpreter Lock) Management: - During C module initialization, GIL handling can be problematic - Python operations from C during initialization may not be fully thread-safe - The interpreter state might not be fully ready for certain Python API calls 2. Module State: - When called from Python, we ensure the module is fully initialized - All required Python objects and state are properly set up - Memory management is handled by Python's garbage collector 3. Error Handling: - Python-level initialization provides better error handling - Exceptions can be properly caught and managed - Prevents potential segmentation faults or undefined behavior 4. Thread Safety: - Ensures thread-safe initialization of global lists - Avoids race conditions during module loading - Provides consistent state across all threads The function specifically: 1. Builds the user allowlist from _DD_IAST_PATCH_MODULES environment variable 2. Builds the user denylist from _DD_IAST_DENY_MODULES environment variable 3. Imports and sets the packages_distributions function for first-party package detection This approach is safer than C-level initialization in init_globals() which can lead to inconsistent state or crashes due to GIL-related issues. """ # Import and set the packages_distributions function for the C extension try: if sys.version_info < (3, 10): import importlib_metadata as metadata else: import importlib.metadata as metadata result = set(metadata.packages_distributions()) iastpatch.set_packages_distributions(result) except (ImportError, AttributeError): # If metadata module is not available, the C extension will handle # first-party detection gracefully by returning False log.debug("Could not import metadata module for first-party detection") except Exception: log.debug("Failed to set packages in C extension", exc_info=True) iastpatch.build_list_from_env(IAST.PATCH_MODULES) iastpatch.build_list_from_env(IAST.DENY_MODULES) def _should_iast_patch(module_name: str) -> bool: """Determines whether a module should be patched by IAST instrumentation. This function checks if a given module should be instrumented by the IAST (Interactive Application Security Testing) system. It uses a series of rules to make this determination, including checking against allowlists and denylists. Performance characteristics: - Without debug: 0.005-0.020s per call - With debug: 0.017-0.024s per call (1.14x-3.08x slower) - Standard library modules are checked fastest - Third-party modules take longest to check Args: module_name (str): The fully qualified name of the module to check (e.g., "os.path", "requests") Returns: bool: True if the module should be patched, False otherwise. Returns True when: - Module is in the user allowlist - Module is in the static allowlist - Module is a first-party module Returns False when: - Module is in Python's stdlib - Module is in the user denylist - Module is in the static denylist - Module is not found in any list Note: When asm_config._iast_debug is True, the function will log detailed information about why a module was allowed or denied patching. """ global IAST_PATCHING_LAZY_LOADED if IAST_PATCHING_LAZY_LOADED: initialize_iast_lists() IAST_PATCHING_LAZY_LOADED = False result = False try: result = iastpatch.should_iast_patch(module_name) if asm_config._iast_debug: if result == iastpatch.DENIED_BUILTINS_DENYLIST: iast_ast_debug_log(f"denying {module_name}. it's in the python_stdlib") elif result == iastpatch.ALLOWED_USER_ALLOWLIST: iast_ast_debug_log(f"allowing {module_name}. it's in the USER_ALLOWLIST") elif result == iastpatch.ALLOWED_STATIC_ALLOWLIST: iast_ast_debug_log(f"allowing {module_name}. it's in the ALLOWLIST") elif result == iastpatch.DENIED_USER_DENYLIST: iast_ast_debug_log(f"denying {module_name}. it's in the USER_DENYLIST") elif result == iastpatch.DENIED_STATIC_DENYLIST: iast_ast_debug_log(f"denying {module_name}. it's in the DENYLIST") elif result == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST: iast_ast_debug_log(f"allowing {module_name}. it's a first party module") elif result == iastpatch.DENIED_NOT_FOUND: iast_ast_debug_log(f"denying {module_name}. it's NOT in the ALLOWLIST") except Exception as e: iast_instrumentation_ast_patching_debug_log( f"An error occurred while attempting to patch the {module_name} module. Error: {e}" ) return result >= iastpatch.ALLOWED_USER_ALLOWLIST def visit_ast( source_text: bytes, module_path: Text, module_name: Text = "", ) -> Optional[ast.Module]: """Visits and modifies a module's AST for IAST instrumentation. This function parses source code into an AST (Abstract Syntax Tree), visits each node using the IAST visitor pattern, and applies necessary modifications for security instrumentation. The visitor pattern is implemented by the AstVisitor class. Args: source_text (bytes): The raw source code of the module to be parsed and modified. module_path (Text): The file system path to the module. Used for error reporting and AST node location information. module_name (Text, optional): The fully qualified name of the module. Used by the visitor for context. Defaults to empty string. Returns: Optional[ast.Module]: The modified AST if changes were made, None if: - No modifications were needed (ast_modified flag is False) - The source couldn't be parsed - The visitor didn't make any changes Note: - Uses the global _VISITOR instance of AstVisitor - Automatically fixes source locations in the modified AST - The visitor's ast_modified flag determines if any changes were made - Returns None instead of unmodified AST to optimize memory usage """ parsed_ast = ast.parse(source_text, module_path) _VISITOR.update_location(filename=module_path, module_name=module_name) modified_ast = _VISITOR.visit(parsed_ast) if not _VISITOR.ast_modified: return None ast.fix_missing_locations(modified_ast) return modified_ast _DIR_WRAPPER = textwrap.dedent( f""" def {_PREFIX}dir(): orig_dir = globals().get("{_PREFIX}orig_dir__") if orig_dir: # Use the original __dir__ method and filter the results results = [name for name in orig_dir() if not name.startswith("{_PREFIX}")] else: # List names from the module's __dict__ and filter out the unwanted names results = [ name for name in globals() if not (name.startswith("{_PREFIX}") or name == "__dir__") ] return results def {_PREFIX}set_dir_filter(): if "__dir__" in globals(): # Store the original __dir__ method globals()["{_PREFIX}orig_dir__"] = __dir__ # Replace the module's __dir__ with the custom one globals()["__dir__"] = {_PREFIX}dir {_PREFIX}set_dir_filter() """ ).encode() def astpatch_module(module: ModuleType) -> tuple[str, Optional[ast.Module]]: """Patches a Python module's AST for IAST instrumentation. This function processes a Python module for IAST (Interactive Application Security Testing) instrumentation by modifying its Abstract Syntax Tree (AST). It handles various edge cases and module types while ensuring proper logging of the patching process. The function performs the following steps: 1. Resolves the module's file path 2. Validates the file (size, extension, accessibility) 3. Reads and processes the source code 4. Optionally adds a __dir__ wrapper to hide IAST internals 5. Generates and returns the modified AST Args: module (ModuleType): The Python module to patch. Must be an imported module object. Returns: tuple[str, Optional[ast.Module]]: A tuple containing: - str: The module's file path if successful, empty string if failed - Optional[ast.Module]: The modified AST if successful, None if: - Module file cannot be found - File is empty (e.g., __init__.py) - File extension not supported (.dll, .so, etc.) - File cannot be read or decoded - AST modification was not needed Note: - Debug logging only occurs when asm_config._iast_debug is True - Handles various file types (.py, .pyc, .pyo, .pyw) - Skips binary/native modules - Can be controlled via IAST.ENV_NO_DIR_PATCH environment variable to disable __dir__ wrapping """ module_name = module.__name__ module_origin = origin(module) if module_origin is None: iast_compiling_debug_log(f"could not find the module: {module_name}") return "", None module_path = str(module_origin) try: if module_origin.stat().st_size == 0: # Don't patch empty files like __init__.py iast_compiling_debug_log(f"empty file: {module_path}") return "", None except OSError: iast_compiling_debug_log(f"could not find the file: {module_path}", exc_info=True) return "", None # Get the file extension, if it's dll, os, pyd, dyn, dynlib: return # If its pyc or pyo, change to .py and check that the file exists. If not, # return with warning. _, module_ext = os.path.splitext(module_path) if module_ext.lower() not in {".pyo", ".pyc", ".pyw", ".py"}: # Probably native or built-in module iast_compiling_debug_log(f"Extension not supported: {module_ext} for: {module_path}") return "", None with open(module_path, "rb") as source_file: try: source_text = source_file.read() except UnicodeDecodeError: iast_compiling_debug_log(f"Encode decode error for file: {module_path}", exc_info=True) return "", None except Exception: iast_compiling_debug_log(f"Unexpected read error: {module_path}", exc_info=True) return "", None if len(source_text.strip()) == 0: # Don't patch empty files like __init__.py iast_compiling_debug_log(f"Empty file: {module_path}") return "", None if not asbool(os.environ.get(IAST.ENV_NO_DIR_PATCH, "false")): # Add the dir filter so __ddtrace stuff is not returned by dir(module) source_text += _DIR_WRAPPER new_ast = visit_ast( source_text, module_path, module_name=module_name, ) if new_ast is None: iast_compiling_debug_log(f"file not ast patched: {module_path}") return "", None return module_path, new_ast