chore(iast): fix IAST imports when under pytest [Backport 2.21] (#12762)

Backport # 12323 to 2.21


## Description

PR #12198 had the unintended consequence of not honoring
`DD_IAST_ENABLED` if set after the `_common_module_patches.py` was
evaluated. This make some tests (`ssrf` and probably others) to not run.
This fixes the problem by moving `is_iast_request_enabled` and
`_IAST_CONTEXT` to `asm_config`.


## Checklist
- [X] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Co-authored-by: Juanjo Alvarez Martinez <[email protected]>

Author: avara1986

ddtrace/appsec/_common_module_patches.py

@@ -25,14 +25,6 @@
 from ddtrace.settings.asm import config as asm_config
 
 
-if asm_config._iast_enabled:
-    from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
-else:
-
-    def is_iast_request_enabled() -> bool:
-        return False
-
-
 log = get_logger(__name__)
 _DD_ORIGINAL_ATTRIBUTES: Dict[Any, Any] = {}
 
@@ -50,15 +42,6 @@ def patch_common_modules():
     subprocess_patch.add_lst_callback(_RASP_POPEN, popen_FD233052260D8B4D)
     if _is_patched:
         return
-    # for testing purposes, we need to update is_iast_request_enabled
-    if asm_config._iast_enabled:
-        global is_iast_request_enabled
-        from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
-    else:
-        global is_iast_request_enabled
-
-        def is_iast_request_enabled() -> bool:
-            return False
 
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
@@ -91,7 +74,7 @@ def wrapped_read_F3E51D71B4EC16EF(original_read_callable, instance, args, kwargs
     wrapper for _io.BytesIO and _io.StringIO read function
     """
     result = original_read_callable(*args, **kwargs)
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         from ddtrace.appsec._iast._taint_tracking import OriginType
         from ddtrace.appsec._iast._taint_tracking import Source
         from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
@@ -117,7 +100,7 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
     """
     wrapper for open file function
     """
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         try:
             from ddtrace.appsec._iast.taint_sinks.path_traversal import check_and_report_path_traversal
 
@@ -208,7 +191,7 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
     wrapper for third party requests.request function
     https://requests.readthedocs.io
     """
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         from ddtrace.appsec._iast.taint_sinks.ssrf import _iast_report_ssrf
 
         _iast_report_ssrf(original_request_callable, *args, **kwargs)

ddtrace/appsec/_constants.py

@@ -161,6 +161,7 @@ class IAST(metaclass=Constant_Class):
 
     TEXT_TYPES = (str, bytes, bytearray)
     TAINTEABLE_TYPES = (str, bytes, bytearray, Match, BytesIO, StringIO)
+    REQUEST_CONTEXT_KEY: Literal["_iast_env"] = "_iast_env"
 
 
 class IAST_SPAN_TAGS(metaclass=Constant_Class):

ddtrace/appsec/_iast/_handlers.py

@@ -17,7 +17,6 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
-from ._iast_request_context import is_iast_request_enabled
 from ._logs import iast_instrumentation_wrapt_debug_log
 from ._logs import iast_propagation_listener_log_log
 from ._taint_tracking._taint_objects import taint_pyobject
@@ -57,7 +56,7 @@ def _on_set_http_meta_iast(
 
 def _on_request_init(wrapped, instance, args, kwargs):
     wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         try:
             instance.query_string = taint_pyobject(
                 pyobject=instance.query_string,
@@ -154,7 +153,7 @@ def _iast_on_wrapped_view(kwargs):
 
 
 def _on_wsgi_environ(wrapped, _instance, args, kwargs):
-    if asm_config._iast_enabled and args and is_iast_request_enabled():
+    if asm_config._iast_enabled and args and asm_config.is_iast_request_enabled:
         return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
 
     return wrapped(*args, **kwargs)
@@ -189,7 +188,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
     # If IAST is enabled, and we're wrapping a Django view call, taint the kwargs (view's
     # path parameters)
     if asm_config._iast_enabled and fn_args and isinstance(fn_args[0], first_arg_expected_type):
-        if not is_iast_request_enabled():
+        if not asm_config.is_iast_request_enabled:
             return
 
         http_req = fn_args[0]
@@ -304,7 +303,7 @@ def _on_grpc_response(message):
 
 
 def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         try:
             for key, value in wrapped(*args, **kwargs):
                 new_key = taint_pyobject(pyobject=key, source_name=key, source_value=key, source_origin=origins[0])
@@ -321,7 +320,7 @@ def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
 
 def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
     value = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         try:
             if not is_pyobject_tainted(value):
                 name = str(args[0]) if len(args) else "http.request.body"
@@ -335,7 +334,7 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
 
 def if_iast_taint_starlette_datastructures(origin, wrapped, instance, args, kwargs):
     value = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         try:
             res = []
             for element in value:
@@ -441,7 +440,7 @@ def _on_pre_tracedrequest_iast(ctx):
 
 
 def _on_set_request_tags_iast(request, span, flask_config):
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         request.cookies = taint_structure(
             request.cookies,
             OriginType.COOKIE_NAME,

ddtrace/appsec/_iast/_iast_request_context.py

@@ -29,8 +29,6 @@
 else:
     from typing_extensions import Literal  # noqa:F401
 
-_IAST_CONTEXT: Literal["_iast_env"] = "_iast_env"
-
 
 class IASTEnvironment:
     """
@@ -49,17 +47,17 @@ def __init__(self, span: Optional[Span] = None):
 
 
 def _get_iast_context() -> Optional[IASTEnvironment]:
-    return core.get_item(_IAST_CONTEXT)
+    return core.get_item(IAST.REQUEST_CONTEXT_KEY)
 
 
 def in_iast_context() -> bool:
-    return core.get_item(_IAST_CONTEXT) is not None
+    return core.get_item(IAST.REQUEST_CONTEXT_KEY) is not None
 
 
 def start_iast_context():
     if asm_config._iast_enabled:
         create_propagation_context()
-        core.set_item(_IAST_CONTEXT, IASTEnvironment())
+        core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment())
 
 
 def end_iast_context(span: Optional[Span] = None):
@@ -70,7 +68,7 @@ def end_iast_context(span: Optional[Span] = None):
 
 
 def finalize_iast_env(env: IASTEnvironment) -> None:
-    core.discard_item(_IAST_CONTEXT)
+    core.discard_item(IAST.REQUEST_CONTEXT_KEY)
 
 
 def set_iast_reporter(iast_reporter: IastSpanReporter) -> None:
@@ -104,13 +102,6 @@ def set_iast_request_enabled(request_enabled) -> None:
         log.debug("iast::propagation::context::Trying to set IAST reporter but no context is present")
 
 
-def is_iast_request_enabled() -> bool:
-    env = _get_iast_context()
-    if env:
-        return env.request_enabled
-    return False
-
-
 def _move_iast_data_to_root_span():
     return asbool(os.getenv("_DD_IAST_USE_ROOT_SPAN"))
 
@@ -156,7 +147,7 @@ def _iast_end_request(ctx=None, span=None, *args, **kwargs):
             existing_data = req_span.get_tag(IAST.JSON)
             if existing_data is None:
                 if req_span.get_metric(IAST.ENABLED) is None:
-                    if not is_iast_request_enabled():
+                    if not asm_config.is_iast_request_enabled:
                         req_span.set_metric(IAST.ENABLED, 0.0)
                         end_iast_context(req_span)
                         oce.release_request()

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -4,7 +4,6 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
-from .._iast_request_context import is_iast_request_enabled
 from .._patch import set_and_check_module_is_patched
 from .._patch import set_module_unpatched
 from .._patch import try_wrap_function_wrapper
@@ -42,7 +41,7 @@ def wrapped_loads(wrapped, instance, args, kwargs):
     from .._taint_utils import taint_structure
 
     obj = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and is_iast_request_enabled():
+    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
         from .._taint_tracking._taint_objects import get_tainted_ranges
         from .._taint_tracking._taint_objects import taint_pyobject
 

ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py

@@ -5,7 +5,6 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
-from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
 from ddtrace.appsec._iast._logs import iast_propagation_debug_log
 from ddtrace.appsec._iast._logs import iast_propagation_error_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_source
@@ -18,13 +17,14 @@
 from ddtrace.appsec._iast._taint_tracking import set_ranges
 from ddtrace.appsec._iast._taint_tracking import set_ranges_from_values
 from ddtrace.internal.logger import get_logger
+from ddtrace.settings.asm import config as asm_config
 
 
 log = get_logger(__name__)
 
 
 def is_pyobject_tainted(pyobject: Any) -> bool:
-    if not is_iast_request_enabled():
+    if not asm_config.is_iast_request_enabled:
         return False
     if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):  # type: ignore[misc]
         return False
@@ -37,7 +37,7 @@ def is_pyobject_tainted(pyobject: Any) -> bool:
 
 
 def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
-    if not is_iast_request_enabled():
+    if not asm_config.is_iast_request_enabled:
         return pyobject
 
     if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):  # type: ignore[misc]
@@ -69,7 +69,7 @@ def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, sou
 
 
 def taint_pyobject_with_ranges(pyobject: Any, ranges: Tuple) -> bool:
-    if not is_iast_request_enabled():
+    if not asm_config.is_iast_request_enabled:
         return False
     if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):  # type: ignore[misc]
         return False
@@ -82,7 +82,7 @@ def taint_pyobject_with_ranges(pyobject: Any, ranges: Tuple) -> bool:
 
 
 def get_tainted_ranges(pyobject: Any) -> Tuple:
-    if not is_iast_request_enabled():
+    if not asm_config.is_iast_request_enabled:
         return tuple()
     if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):  # type: ignore[misc]
         return tuple()

ddtrace/appsec/_iast/reporter.py

@@ -19,6 +19,7 @@
 from ddtrace.appsec._iast.constants import VULN_WEAK_CIPHER_TYPE
 from ddtrace.appsec._iast.constants import VULN_WEAK_RANDOMNESS
 from ddtrace.internal.logger import get_logger
+from ddtrace.settings.asm import config as asm_config
 
 
 log = get_logger(__name__)
@@ -237,9 +238,7 @@ def taint_ranges_as_evidence_info(pyobject: Any) -> Tuple[List[Source], List[Dic
         return sources, tainted_ranges_to_dict
 
     def add_ranges_to_evidence_and_extract_sources(self, vuln):
-        from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
-
-        if not is_iast_request_enabled():
+        if not asm_config.is_iast_request_enabled:
             log.debug(
                 "iast::propagation::context::add_ranges_to_evidence_and_extract_sources. "
                 "No request quota or this vulnerability is outside the context"
@@ -258,9 +257,7 @@ def build_and_scrub_value_parts(self) -> Dict[str, Any]:
         Returns:
         - Dict[str, Any]: Dictionary representation of the IAST span reporter.
         """
-        from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
-
-        if not is_iast_request_enabled():
+        if not asm_config.is_iast_request_enabled:
             log.debug(
                 "iast::propagation::context::build_and_scrub_value_parts. "
                 "No request quota or this vulnerability is outside the context"

ddtrace/appsec/_iast/taint_sinks/_base.py

@@ -11,7 +11,6 @@
 from ddtrace.trace import tracer
 
 from .._iast_request_context import get_iast_reporter
-from .._iast_request_context import is_iast_request_enabled
 from .._iast_request_context import set_iast_reporter
 from .._overhead_control_engine import Operation
 from .._stacktrace import get_info_frame
@@ -57,7 +56,7 @@ def wrapper(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
             """Get the current root Span and attach it to the wrapped function. We need the span to report the
             vulnerability and update the context with the report information.
             """
-            if not is_iast_request_enabled():
+            if not asm_config.is_iast_request_enabled:
                 if _is_iast_debug_enabled():
                     log.debug(
                         "iast::propagation::context::VulnerabilityBase.wrapper. No request quota or this vulnerability "
@@ -74,7 +73,7 @@ def wrapper(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     @classmethod
     @taint_sink_deduplication
     def _prepare_report(cls, vulnerability_type, evidence, file_name, line_number):
-        if not is_iast_request_enabled():
+        if not asm_config.is_iast_request_enabled:
             if _is_iast_debug_enabled():
                 log.debug(
                     "iast::propagation::context::VulnerabilityBase._prepare_report. "

ddtrace/appsec/_iast/taint_sinks/code_injection.py

@@ -4,7 +4,6 @@
 from ddtrace.appsec._common_module_patches import try_unwrap
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast import oce
-from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_sink
 from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
 from ddtrace.appsec._iast._metrics import increment_iast_span_metric
@@ -82,6 +81,6 @@ class CodeInjection(VulnerabilityBase):
 def _iast_report_code_injection(code_string: Text):
     increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, CodeInjection.vulnerability_type)
     _set_metric_iast_executed_sink(CodeInjection.vulnerability_type)
-    if is_iast_request_enabled():
+    if asm_config.is_iast_request_enabled:
         if is_pyobject_tainted(code_string):
             CodeInjection.report(evidence_value=code_string)

ddtrace/appsec/_iast/taint_sinks/command_injection.py

@@ -3,7 +3,6 @@
 
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast import oce
-from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_sink
 from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
 from ddtrace.appsec._iast._metrics import increment_iast_span_metric
@@ -50,7 +49,7 @@ def _iast_report_cmdi(shell_args: Union[str, List[str]]) -> None:
     increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, CommandInjection.vulnerability_type)
     _set_metric_iast_executed_sink(CommandInjection.vulnerability_type)
 
-    if is_iast_request_enabled() and CommandInjection.has_quota():
+    if asm_config.is_iast_request_enabled and CommandInjection.has_quota():
         from .._taint_tracking.aspects import join_aspect
 
         if isinstance(shell_args, (list, tuple)):