feat(iast): security controls

ddtrace/appsec/_iast/_patch_modules.py

@@ -1,15 +1,27 @@
+import functools
+
 from wrapt.importer import when_imported
 
+from ddtrace.appsec._common_module_patches import try_unwrap
 from ddtrace.appsec._common_module_patches import try_wrap_function_wrapper
+from ddtrace.appsec._iast.secure_marks.configuration import SC_SANITIZER
+from ddtrace.appsec._iast.secure_marks.configuration import SC_VALIDATOR
+from ddtrace.appsec._iast.secure_marks.configuration import SecurityControl
+from ddtrace.appsec._iast.secure_marks.configuration import get_security_controls_from_env
 from ddtrace.appsec._iast.secure_marks.sanitizers import cmdi_sanitizer
+from ddtrace.appsec._iast.secure_marks.sanitizers import create_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import header_injection_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import path_traversal_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import sqli_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import xss_sanitizer
+from ddtrace.appsec._iast.secure_marks.validators import create_validator
 from ddtrace.appsec._iast.secure_marks.validators import header_injection_validator
 from ddtrace.appsec._iast.secure_marks.validators import ssrf_validator
 from ddtrace.appsec._iast.secure_marks.validators import unvalidated_redirect_validator
+from ddtrace.internal.logger import get_logger
+
 
+log = get_logger(__name__)
 
 IAST_PATCH = {
     "code_injection": True,
@@ -34,6 +46,9 @@ def patch_iast(patch_modules=IAST_PATCH):
     for module in (m for m, e in patch_modules.items() if e):
         when_imported("hashlib")(_on_import_factory(module, "ddtrace.appsec._iast.taint_sinks.%s", raise_errors=False))
 
+    # Apply custom security controls from environment variable
+    _apply_custom_security_controls()
+
     # CMDI sanitizers
     when_imported("shlex")(
         lambda _: try_wrap_function_wrapper(
@@ -137,3 +152,66 @@ def patch_iast(patch_modules=IAST_PATCH):
     #     )
     # )
     when_imported("json")(_on_import_factory("json_tainting", "ddtrace.appsec._iast._patches.%s", raise_errors=False))
+
+
+def _apply_custom_security_controls():
+    """Apply custom security controls from DD_IAST_SECURITY_CONTROLS_CONFIGURATION environment variable."""
+    try:
+        security_controls = get_security_controls_from_env()
+
+        if not security_controls:
+            log.debug("No custom security controls configured")
+            return
+
+        log.debug("Applying %s custom security controls", len(security_controls))
+
+        for control in security_controls:
+            try:
+                _apply_security_control(control)
+            except Exception:
+                log.warning("Failed to apply security control %s", control, exc_info=True)
+    except Exception:
+        log.warning("Failed to load custom security controls", exc_info=True)
+
+
+def _apply_security_control(control: SecurityControl):
+    """Apply a single security control configuration.
+
+    Args:
+        control: SecurityControl object containing the configuration
+    """
+    # Create the appropriate wrapper function
+    if control.control_type == SC_SANITIZER:
+        wrapper_func = functools.partial(create_sanitizer, control.vulnerability_types)
+    elif control.control_type == SC_VALIDATOR:
+        wrapper_func = functools.partial(create_validator, control.vulnerability_types, control.parameters)
+    else:
+        log.warning("Unknown control type: %s", control.control_type)
+
+    # Split module path to get the base module for when_imported
+    base_module = control.module_path.split(".")[0]
+    when_imported(base_module)(
+        lambda _: try_wrap_function_wrapper(
+            control.module_path,
+            control.method_name,
+            wrapper_func,
+        )
+    )
+    log.debug(
+        "Configured %s for %s.%s (vulnerabilities: %s)",
+        control.control_type,
+        control.module_path,
+        control.method_name,
+        [v.name for v in control.vulnerability_types],
+    )
+
+
+def _unapply_security_control():
+    """Remove security control configuration for testing proposes
+
+    Args:
+        control: SecurityControl object containing the configuration
+    """
+    security_controls = get_security_controls_from_env()
+    for control in security_controls:
+        try_unwrap(control.module_path, control.method_name)

ddtrace/appsec/_iast/secure_marks/README_CONFIGURATION.md

@@ -0,0 +1,118 @@
+# IAST Security Controls Configuration
+
+This document explains how to configure custom security controls for IAST using the `DD_IAST_SECURITY_CONTROLS_CONFIGURATION` environment variable.
+
+## Overview
+
+The `DD_IAST_SECURITY_CONTROLS_CONFIGURATION` environment variable allows you to specify custom sanitizers and validators that IAST should recognize when analyzing your application for security vulnerabilities.
+
+## Format
+
+The configuration uses the following format:
+
+```
+CONTROL_TYPE:VULNERABILITY_TYPES:MODULE:METHOD[:PARAMETER_POSITIONS]
+```
+
+Multiple security controls are separated by semicolons (`;`).
+
+### Fields
+
+1. **CONTROL_TYPE**: Either `INPUT_VALIDATOR` or `SANITIZER`
+2. **VULNERABILITY_TYPES**: Comma-separated list of vulnerability types or `*` for all types
+3. **MODULE**: Python module path (e.g., `shlex`, `django.utils.http`)
+4. **METHOD**: Method name to instrument
+5. **PARAMETER_POSITIONS** (Optional): Zero-based parameter positions to validate (INPUT_VALIDATOR only)
+
+### Vulnerability Types
+
+Supported vulnerability types:
+- `COMMAND_INJECTION` / `CMDI`
+- `CODE_INJECTION`
+- `SQL_INJECTION` / `SQLI`
+- `XSS`
+- `HEADER_INJECTION`
+- `PATH_TRAVERSAL`
+- `SSRF`
+- `UNVALIDATED_REDIRECT`
+- `INSECURE_COOKIE`
+- `NO_HTTPONLY_COOKIE`
+- `NO_SAMESITE_COOKIE`
+- `WEAK_CIPHER`
+- `WEAK_HASH`
+- `WEAK_RANDOMNESS`
+- `STACKTRACE_LEAK`
+
+Use `*` to apply to all vulnerability types.
+
+## Examples
+
+### Basic Examples
+
+#### Input Validator for Command Injection
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="INPUT_VALIDATOR:COMMAND_INJECTION:shlex:quote"
+```
+
+#### Sanitizer for XSS
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="SANITIZER:XSS:html:escape"
+```
+
+#### Multiple Vulnerability Types
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="INPUT_VALIDATOR:COMMAND_INJECTION,XSS:custom.validator:validate_input"
+```
+
+#### All Vulnerability Types
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="SANITIZER:*:custom.sanitizer:sanitize_all"
+```
+
+### Advanced Examples
+
+#### Multiple Security Controls
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="INPUT_VALIDATOR:COMMAND_INJECTION:shlex:quote;SANITIZER:XSS:html:escape;SANITIZER:SQLI:custom.db:escape_sql"
+```
+
+#### Validator with Specific Parameter Positions
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="INPUT_VALIDATOR:COMMAND_INJECTION:custom.validator:validate:0,2"
+```
+This validates only the 1st and 3rd parameters (0-based indexing).
+
+#### Complex Configuration
+```bash
+export DD_IAST_SECURITY_CONTROLS_CONFIGURATION="INPUT_VALIDATOR:COMMAND_INJECTION,XSS:security.validators:validate_user_input:0,1;SANITIZER:SQLI:database.utils:escape_sql_string;SANITIZER:*:security.sanitizers:clean_all_inputs"
+```
+
+## How It Works
+
+### Input Validators
+- **Purpose**: Mark input parameters as safe after validation
+- **When to use**: When your function validates input and returns a boolean or throws an exception
+- **Effect**: Parameters are marked as secure for the specified vulnerability types
+- **Parameter positions**: Optionally specify which parameters to mark (0-based index)
+
+### Sanitizers
+- **Purpose**: Mark return values as safe after sanitization
+- **When to use**: When your function cleans/escapes input and returns the sanitized value
+- **Effect**: Return value is marked as secure for the specified vulnerability types
+
+## Integration with Existing Controls
+
+Your custom security controls work alongside the built-in IAST security controls:
+
+- `shlex.quote` (Command injection sanitizer)
+- `html.escape` (XSS sanitizer)
+- Database escape functions (SQL injection sanitizers)
+- Django validators (Various validators)
+- And more...
+
+## Error Handling
+
+If there are errors in the configuration:
+- Invalid configurations are logged and skipped
+- The application continues to run with built-in security controls
+- Check application logs for configuration warnings/errors

ddtrace/appsec/_iast/secure_marks/__init__.py

@@ -3,8 +3,15 @@
 This package provides functions to mark values as secure from specific vulnerabilities.
 It includes both sanitizers (which transform and secure values) and validators (which
 verify values are secure).
+
+It also provides configuration capabilities for custom security controls via the
+DD_IAST_SECURITY_CONTROLS_CONFIGURATION environment variable.
 """
 
+from .configuration import VULNERABILITY_TYPE_MAPPING
+from .configuration import SecurityControl
+from .configuration import get_security_controls_from_env
+from .configuration import parse_security_controls_config
 from .sanitizers import cmdi_sanitizer
 from .sanitizers import path_traversal_sanitizer
 from .sanitizers import sqli_sanitizer
@@ -22,4 +29,9 @@
     "path_traversal_validator",
     "sqli_validator",
     "cmdi_validator",
+    # Configuration
+    "get_security_controls_from_env",
+    "parse_security_controls_config",
+    "SecurityControl",
+    "VULNERABILITY_TYPE_MAPPING",
 ]