feat(iast): security controls

ddtrace/appsec/_iast/_patch_modules.py

@@ -10,8 +10,9 @@
 The module uses wrapt's function wrapping capabilities to intercept calls to security-sensitive
 functions and enable taint tracking and vulnerability detection.
 """
-
+import functools
 from typing import Callable
+from typing import Optional
 from typing import Set
 from typing import Text
 
@@ -21,6 +22,12 @@
 from ddtrace.appsec._common_module_patches import try_wrap_function_wrapper
 from ddtrace.appsec._common_module_patches import wrap_object
 from ddtrace.appsec._iast._logs import iast_instrumentation_wrapt_debug_log
+from ddtrace.appsec._iast.secure_marks import SecurityControl
+from ddtrace.appsec._iast.secure_marks import get_security_controls_from_env
+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.sanitizers import create_sanitizer
+from ddtrace.appsec._iast.secure_marks.validators import create_validator
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
@@ -181,3 +188,55 @@ def _testing_unpatch_iast():
     """
     iast_funcs = WrapFunctonsForIAST()
     iast_funcs.testing_unpatch()
+
+
+def _apply_custom_security_controls(iast_funcs: Optional[WrapFunctonsForIAST] = None):
+    """Apply custom security controls from DD_IAST_SECURITY_CONTROLS_CONFIGURATION environment variable."""
+    try:
+        if iast_funcs is None:
+            iast_funcs = WrapFunctonsForIAST()
+        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(iast_funcs, control)
+            except Exception:
+                log.warning("Failed to apply security control %s", control, exc_info=True)
+        return iast_funcs
+    except Exception:
+        log.warning("Failed to load custom security controls", exc_info=True)
+
+
+def _apply_security_control(iast_funcs: WrapFunctonsForIAST, 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)
+        return
+
+    iast_funcs.wrap_function(
+        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],
+    )

ddtrace/appsec/_iast/main.py

@@ -23,6 +23,7 @@
 from wrapt import when_imported
 
 from ddtrace.appsec._iast._patch_modules import WrapFunctonsForIAST
+from ddtrace.appsec._iast._patch_modules import _apply_custom_security_controls
 from ddtrace.appsec._iast.secure_marks import cmdi_sanitizer
 from ddtrace.appsec._iast.secure_marks import path_traversal_sanitizer
 from ddtrace.appsec._iast.secure_marks import sqli_sanitizer
@@ -72,6 +73,9 @@ def patch_iast(patch_modules=IAST_PATCH):
         when_imported("hashlib")(_on_import_factory(module, "ddtrace.appsec._iast.taint_sinks.%s", raise_errors=False))
 
     iast_funcs = WrapFunctonsForIAST()
+
+    _apply_custom_security_controls(iast_funcs)
+
     # CMDI sanitizers
     iast_funcs.wrap_function("shlex", "quote", cmdi_sanitizer)
 

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",
 ]