chore(iast): refactor import modules hook (#13665)

This PR fixes several logic issues and removes code that was duplicated
dozens of times.
Move _iast_instrument_starlette_url to _handlers

Since `iast/_patch.py` and `iast/_patch_modules.py` had overlapping
responsibilities and it was unclear what each file was for, the
functions that serve as the entry point for starting IAST have been
moved to main.py.

The new class that handles all the logic for wrapping functions and
modules now lives in patch_modules.

On one hand, we had the function `try_wrap_function_wrapper` defined
twice — once for AppSec and once for IAST:

The **AppSec** version used `ModuleWatchdog`:

```python
# ddtrace/appsec/_common_module_patches.py
def try_wrap_function_wrapper(module_name: str, name: str, wrapper: Callable) -> None:
    @ModuleWatchdog.after_module_imported(module_name)
    def _(module):
        try:
            wrap_object(module, name, FunctionWrapper, (wrapper,))
        except (ImportError, AttributeError):
            log.debug("ASM patching. Module %s.%s does not exist", module_name, name)
```

Whereas the **IAST** version used `wrapt` directly:

```python
# ddtrace/appsec/_iast/_patch.py
def try_wrap_function_wrapper(module: Text, name: Text, wrapper: Callable):
    try:
        wrap_object(module, name, FunctionWrapper, (wrapper,))
    except (ImportError, AttributeError):
        iast_instrumentation_wrapt_debug_log(f"Module {module}.{name} not exists")
```

Additionally, the IAST version was sometimes called like this:

```python
try_wrap_function_wrapper(("_%s" % MD5_DEF), "MD5Type.digest", wrapped_md5_function)
```

And other times using `wrapt.when_imported`:

```python
@when_imported("django.shortcuts")
def _(m):
    try_wrap_function_wrapper(m, "redirect", _unvalidated_redirect_for_django)
```

These two patterns were repeated dozens of times, and now all of this
logic has been centralized in the class
`ddtrace.appsec._iast._patch_modules.WrapModulesForIAST`.



## 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)

Author: avara1986

ddtrace/_monkey.py

@@ -352,7 +352,7 @@ def _patch_all(**patch_modules: bool) -> None:
 
     patch(raise_errors=False, **modules)
     if asm_config._iast_enabled:
-        from ddtrace.appsec._iast._patch_modules import patch_iast
+        from ddtrace.appsec._iast.main import patch_iast
         from ddtrace.appsec.iast import enable_iast_propagation
 
         patch_iast()

ddtrace/appsec/_common_module_patches.py

@@ -42,26 +42,32 @@ def _(module):
         subprocess_patch.patch()
         subprocess_patch.add_str_callback(_RASP_SYSTEM, wrapped_system_5542593D237084A7)
         subprocess_patch.add_lst_callback(_RASP_POPEN, popen_FD233052260D8B4D)
+        log.debug("Patching common modules: subprocess_patch")
 
     if _is_patched:
         return
 
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
+    log.debug("Patching common modules: builtins and urllib.request")
     _is_patched = True
 
 
 def unpatch_common_modules():
     global _is_patched
     if not _is_patched:
         return
+
     try_unwrap("builtins", "open")
     try_unwrap("urllib.request", "OpenerDirector.open")
     try_unwrap("_io", "BytesIO.read")
     try_unwrap("_io", "StringIO.read")
+    subprocess_patch.unpatch()
     subprocess_patch.del_str_callback(_RASP_SYSTEM)
     subprocess_patch.del_lst_callback(_RASP_POPEN)
+
+    log.debug("Unpatching common modules subprocess, builtins and urllib.request")
     _is_patched = False
 
 
@@ -323,7 +329,7 @@ def try_unwrap(module, name):
             apply_patch(parent, attribute, original)
             del _DD_ORIGINAL_ATTRIBUTES[(parent, attribute)]
     except ModuleNotFoundError:
-        pass
+        log.debug("ERROR unwrapping %s.%s ", module, name)
 
 
 def try_wrap_function_wrapper(module_name: str, name: str, wrapper: Callable) -> None:
@@ -332,7 +338,7 @@ def _(module):
         try:
             wrap_object(module, name, FunctionWrapper, (wrapper,))
         except (ImportError, AttributeError):
-            log.debug("ASM patching. Module %s.%s does not exist", module_name, name)
+            log.debug("Module %s.%s does not exist", module_name, name)
 
 
 def wrap_object(module, name, factory, args=(), kwargs=None):

ddtrace/appsec/_iast/_handlers.py

@@ -1,18 +1,14 @@
 from collections.abc import MutableMapping
 import functools
 
-from wrapt import when_imported
-from wrapt import wrap_function_wrapper as _w
-
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._iast_request_context_base import get_iast_stacktrace_reported
 from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_endpoint
 from ddtrace.appsec._iast._iast_request_context_base import set_iast_stacktrace_reported
 from ddtrace.appsec._iast._logs import iast_instrumentation_wrapt_debug_log
 from ddtrace.appsec._iast._logs import iast_propagation_listener_log_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
-from ddtrace.appsec._iast._patch import _iast_instrument_starlette_url
-from ddtrace.appsec._iast._patch import try_wrap_function_wrapper
+from ddtrace.appsec._iast._patch_modules import WrapFunctonsForIAST
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
@@ -62,27 +58,29 @@ def _on_flask_patch(flask_version):
     """
     if asm_config._iast_enabled:
         try:
-            try_wrap_function_wrapper(
+            iast_funcs = WrapFunctonsForIAST()
+
+            iast_funcs.wrap_function(
                 "werkzeug.datastructures",
                 "Headers.items",
                 functools.partial(if_iast_taint_yield_tuple_for, (OriginType.HEADER_NAME, OriginType.HEADER)),
             )
 
-            try_wrap_function_wrapper(
+            iast_funcs.wrap_function(
                 "werkzeug.datastructures",
                 "EnvironHeaders.__getitem__",
                 functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
             )
             # Since werkzeug 3.1.0 get doesn't call to __getitem__
-            try_wrap_function_wrapper(
+            iast_funcs.wrap_function(
                 "werkzeug.datastructures",
                 "EnvironHeaders.get",
                 functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
             )
             _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
             _set_metric_iast_instrumented_source(OriginType.HEADER)
 
-            try_wrap_function_wrapper(
+            iast_funcs.wrap_function(
                 "werkzeug.datastructures",
                 "ImmutableMultiDict.__getitem__",
                 functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
@@ -91,17 +89,17 @@ def _on_flask_patch(flask_version):
 
             if flask_version >= (2, 0, 0):
                 # instance.query_string: raising an error on werkzeug/_internal.py "AttributeError: read only property"
-                try_wrap_function_wrapper("werkzeug.wrappers.request", "Request.__init__", _on_request_init)
+                iast_funcs.wrap_function("werkzeug.wrappers.request", "Request.__init__", _on_request_init)
 
             _set_metric_iast_instrumented_source(OriginType.PATH)
             _set_metric_iast_instrumented_source(OriginType.QUERY)
 
-            try_wrap_function_wrapper(
+            iast_funcs.wrap_function(
                 "werkzeug.wrappers.request",
                 "Request.get_data",
                 functools.partial(taint_dictionary, OriginType.BODY, OriginType.BODY),
             )
-            try_wrap_function_wrapper(
+            iast_funcs.wrap_function(
                 "werkzeug.wrappers.request",
                 "Request.get_json",
                 functools.partial(taint_dictionary, OriginType.BODY, OriginType.BODY),
@@ -110,13 +108,15 @@ def _on_flask_patch(flask_version):
             _set_metric_iast_instrumented_source(OriginType.BODY)
 
             if flask_version < (2, 0, 0):
-                _w(
+                iast_funcs.wrap_function(
                     "werkzeug._internal",
                     "_DictAccessorProperty.__get__",
                     functools.partial(if_iast_taint_returned_object_for, OriginType.QUERY),
                 )
                 _set_metric_iast_instrumented_source(OriginType.QUERY)
 
+            iast_funcs.patch()
+
             # Instrumented on _ddtrace.appsec._asm_request_context._on_wrapped_view
             _set_metric_iast_instrumented_source(OriginType.PATH_PARAMETER)
 
@@ -156,14 +156,17 @@ def _on_django_patch():
     """Handle Django framework patch event."""
     if asm_config._iast_enabled:
         try:
-            when_imported("django.http.request")(
-                lambda m: try_wrap_function_wrapper(
-                    m,
-                    "QueryDict.__getitem__",
-                    functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
-                )
+            iast_funcs = WrapFunctonsForIAST()
+
+            iast_funcs.wrap_function(
+                "django.http.request",
+                "QueryDict.__getitem__",
+                functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
             )
-            try_wrap_function_wrapper("django.utils.shlex", "quote", cmdi_sanitizer)
+            iast_funcs.wrap_function("django.utils.shlex", "quote", cmdi_sanitizer)
+
+            iast_funcs.patch()
+
             # we instrument those sources on _on_django_func_wrapped
             _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
             _set_metric_iast_instrumented_source(OriginType.HEADER)
@@ -354,82 +357,88 @@ def if_iast_taint_starlette_datastructures(origin, wrapped, instance, args, kwar
 
 
 def _on_iast_fastapi_patch():
-    # Cookies sources
-    try_wrap_function_wrapper(
-        "starlette.requests",
-        "cookie_parser",
-        functools.partial(taint_dictionary, OriginType.COOKIE_NAME, OriginType.COOKIE),
-    )
-    _set_metric_iast_instrumented_source(OriginType.COOKIE)
-    _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
-
-    # Parameter sources
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "QueryParams.__getitem__",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
-    )
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "QueryParams.get",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
-    )
-    _set_metric_iast_instrumented_source(OriginType.PARAMETER)
+    try:
+        iast_funcs = WrapFunctonsForIAST()
+        # Cookies sources
+        iast_funcs.wrap_function(
+            "starlette.requests",
+            "cookie_parser",
+            functools.partial(taint_dictionary, OriginType.COOKIE_NAME, OriginType.COOKIE),
+        )
+        _set_metric_iast_instrumented_source(OriginType.COOKIE)
+        _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
+
+        # Parameter sources
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "QueryParams.__getitem__",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
+        )
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "QueryParams.get",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
+        )
+        _set_metric_iast_instrumented_source(OriginType.PARAMETER)
 
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "QueryParams.keys",
-        functools.partial(if_iast_taint_starlette_datastructures, OriginType.PARAMETER_NAME),
-    )
-    _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "QueryParams.keys",
+            functools.partial(if_iast_taint_starlette_datastructures, OriginType.PARAMETER_NAME),
+        )
+        _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
 
-    # Header sources
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "Headers.__getitem__",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
-    )
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "Headers.get",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
-    )
-    _set_metric_iast_instrumented_source(OriginType.HEADER)
+        # Header sources
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "Headers.__getitem__",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
+        )
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "Headers.get",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.HEADER),
+        )
+        _set_metric_iast_instrumented_source(OriginType.HEADER)
 
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "Headers.keys",
-        functools.partial(if_iast_taint_starlette_datastructures, OriginType.HEADER_NAME),
-    )
-    _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
-
-    # Path source
-    try_wrap_function_wrapper("starlette.datastructures", "URL.__init__", _iast_instrument_starlette_url)
-    _set_metric_iast_instrumented_source(OriginType.PATH)
-
-    # Body source
-    try_wrap_function_wrapper("starlette.requests", "Request.__init__", _iast_instrument_starlette_request)
-    try_wrap_function_wrapper("starlette.requests", "Request.body", _iast_instrument_starlette_request_body)
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "FormData.__getitem__",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.BODY),
-    )
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "FormData.get",
-        functools.partial(if_iast_taint_returned_object_for, OriginType.BODY),
-    )
-    try_wrap_function_wrapper(
-        "starlette.datastructures",
-        "FormData.keys",
-        functools.partial(if_iast_taint_starlette_datastructures, OriginType.PARAMETER_NAME),
-    )
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "Headers.keys",
+            functools.partial(if_iast_taint_starlette_datastructures, OriginType.HEADER_NAME),
+        )
+        _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
+
+        # Path source
+        iast_funcs.wrap_function("starlette.datastructures", "URL.__init__", _iast_instrument_starlette_url)
+        _set_metric_iast_instrumented_source(OriginType.PATH)
+
+        # Body source
+        iast_funcs.wrap_function("starlette.requests", "Request.__init__", _iast_instrument_starlette_request)
+        iast_funcs.wrap_function("starlette.requests", "Request.body", _iast_instrument_starlette_request_body)
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "FormData.__getitem__",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.BODY),
+        )
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "FormData.get",
+            functools.partial(if_iast_taint_returned_object_for, OriginType.BODY),
+        )
+        iast_funcs.wrap_function(
+            "starlette.datastructures",
+            "FormData.keys",
+            functools.partial(if_iast_taint_starlette_datastructures, OriginType.PARAMETER_NAME),
+        )
 
-    _set_metric_iast_instrumented_source(OriginType.BODY)
+        iast_funcs.patch()
 
-    # Instrumented on _iast_starlette_scope_taint
-    _set_metric_iast_instrumented_source(OriginType.PATH_PARAMETER)
+        _set_metric_iast_instrumented_source(OriginType.BODY)
+
+        # Instrumented on _iast_starlette_scope_taint
+        _set_metric_iast_instrumented_source(OriginType.PATH_PARAMETER)
+    except Exception:
+        iast_propagation_listener_log_log("Unexpected exception while tainting pyobject", exc_info=True)
 
 
 def _on_pre_tracedrequest_iast(ctx):
@@ -570,3 +579,16 @@ def _iast_instrument_starlette_scope(scope, route):
                     )
     except Exception:
         iast_propagation_listener_log_log("Unexpected exception while tainting path parameters", exc_info=True)
+
+
+def _iast_instrument_starlette_url(wrapped, instance, args, kwargs):
+    def path(self) -> str:
+        return taint_pyobject(
+            self.components.path,
+            source_name=origin_to_str(OriginType.PATH),
+            source_value=self.components.path,
+            source_origin=OriginType.PATH,
+        )
+
+    instance.__class__.path = property(path)
+    wrapped(*args, **kwargs)