chore(iast): refactor import modules hook

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)