cmdi patch refactor + rasp support

ddtrace/appsec/_capabilities.py

@@ -31,6 +31,7 @@ class Flags(enum.IntFlag):
     ASM_SESSION_FINGERPRINT = 1 << 33
     ASM_NETWORK_FINGERPRINT = 1 << 34
     ASM_HEADER_FINGERPRINT = 1 << 35
+    ASM_RASP_CMDI = 1 << 37
 
 
 _ALL_ASM_BLOCKING = (
@@ -49,7 +50,7 @@ class Flags(enum.IntFlag):
     | Flags.ASM_HEADER_FINGERPRINT
 )
 
-_ALL_RASP = Flags.ASM_RASP_SQLI | Flags.ASM_RASP_LFI | Flags.ASM_RASP_SSRF | Flags.ASM_RASP_SHI
+_ALL_RASP = Flags.ASM_RASP_SQLI | Flags.ASM_RASP_LFI | Flags.ASM_RASP_SSRF | Flags.ASM_RASP_SHI | Flags.ASM_RASP_CMDI
 _FEATURE_REQUIRED = Flags.ASM_ACTIVATION | Flags.ASM_AUTO_USER
 
 

ddtrace/appsec/_common_module_patches.py

@@ -17,6 +17,7 @@
 from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
 from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
 from ddtrace.appsec._iast.constants import VULN_PATH_TRAVERSAL
+import ddtrace.contrib.internal.subprocess.patch as subprocess_patch
 from ddtrace.internal import core
 from ddtrace.internal._exceptions import BlockingException
 from ddtrace.internal._unpatched import _gc as gc
@@ -39,7 +40,10 @@ def patch_common_modules():
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
     try_wrap_function_wrapper("_io", "BytesIO.read", wrapped_read_F3E51D71B4EC16EF)
     try_wrap_function_wrapper("_io", "StringIO.read", wrapped_read_F3E51D71B4EC16EF)
-    try_wrap_function_wrapper("os", "system", wrapped_system_5542593D237084A7)
+    # ensure that the subprocess patch is applied even after one click activation
+    subprocess_patch.patch()
+    subprocess_patch.add_str_callback("rasp os.system", wrapped_system_5542593D237084A7)
+    subprocess_patch.add_lst_callback("rasp Popen", popen_FD233052260D8B4D)
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
     if asm_config._iast_enabled:
         _set_metric_iast_instrumented_sink(VULN_PATH_TRAVERSAL)
@@ -54,6 +58,7 @@ def unpatch_common_modules():
     try_unwrap("urllib.request", "OpenerDirector.open")
     try_unwrap("_io", "BytesIO.read")
     try_unwrap("_io", "StringIO.read")
+    subprocess_patch.del_str_callback("rasp os.system")
     _is_patched = False
 
 
@@ -211,45 +216,58 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
     return original_request_callable(*args, **kwargs)
 
 
-def wrapped_system_5542593D237084A7(original_command_callable, instance, args, kwargs):
+def wrapped_system_5542593D237084A7(command: str) -> None:
     """
     wrapper for os.system function
     """
-    command = args[0] if args else kwargs.get("command", None)
-    if command is not None:
-        if asm_config._iast_enabled and is_iast_request_enabled():
-            from ddtrace.appsec._iast.taint_sinks.command_injection import _iast_report_cmdi
-
-            _iast_report_cmdi(command)
-
-        if (
-            asm_config._asm_enabled
-            and asm_config._ep_enabled
-            and ddtrace.tracer._appsec_processor is not None
-            and ddtrace.tracer._appsec_processor.rasp_cmdi_enabled
-        ):
-            try:
-                from ddtrace.appsec._asm_request_context import call_waf_callback
-                from ddtrace.appsec._asm_request_context import in_asm_context
-                from ddtrace.appsec._constants import EXPLOIT_PREVENTION
-            except ImportError:
-                return original_command_callable(*args, **kwargs)
-
-            if in_asm_context():
-                res = call_waf_callback(
-                    {EXPLOIT_PREVENTION.ADDRESS.CMDI: command},
-                    crop_trace="wrapped_system_5542593D237084A7",
-                    rule_type=EXPLOIT_PREVENTION.TYPE.CMDI,
-                )
-                if res and _must_block(res.actions):
-                    raise BlockingException(get_blocked(), "exploit_prevention", "cmdi", command)
-    try:
-        return original_command_callable(*args, **kwargs)
-    except Exception as e:
-        previous_frame = e.__traceback__.tb_frame.f_back
-        raise e.with_traceback(
-            e.__traceback__.__class__(None, previous_frame, previous_frame.f_lasti, previous_frame.f_lineno)
-        )
+    if (
+        asm_config._asm_enabled
+        and asm_config._ep_enabled
+        and ddtrace.tracer._appsec_processor is not None
+        and ddtrace.tracer._appsec_processor.rasp_shi_enabled
+    ):
+        try:
+            from ddtrace.appsec._asm_request_context import call_waf_callback
+            from ddtrace.appsec._asm_request_context import in_asm_context
+            from ddtrace.appsec._constants import EXPLOIT_PREVENTION
+        except ImportError:
+            return
+
+        if in_asm_context():
+            res = call_waf_callback(
+                {EXPLOIT_PREVENTION.ADDRESS.SHI: command},
+                crop_trace="wrapped_system_5542593D237084A7",
+                rule_type=EXPLOIT_PREVENTION.TYPE.SHI,
+            )
+            if res and _must_block(res.actions):
+                raise BlockingException(get_blocked(), "exploit_prevention", "shi", command)
+
+
+def popen_FD233052260D8B4D(arg_list) -> None:
+    """
+    listener for subprocess.Popen class
+    """
+    if (
+        asm_config._asm_enabled
+        and asm_config._ep_enabled
+        and ddtrace.tracer._appsec_processor is not None
+        and ddtrace.tracer._appsec_processor.rasp_cmdi_enabled
+    ):
+        try:
+            from ddtrace.appsec._asm_request_context import call_waf_callback
+            from ddtrace.appsec._asm_request_context import in_asm_context
+            from ddtrace.appsec._constants import EXPLOIT_PREVENTION
+        except ImportError:
+            return
+
+        if in_asm_context():
+            res = call_waf_callback(
+                {EXPLOIT_PREVENTION.ADDRESS.CMDI: arg_list},
+                crop_trace="popen_FD233052260D8B4D",
+                rule_type=EXPLOIT_PREVENTION.TYPE.CMDI,
+            )
+            if res and _must_block(res.actions):
+                raise BlockingException(get_blocked(), "exploit_prevention", "cmdi", arg_list)
 
 
 _DB_DIALECTS = {

ddtrace/appsec/_constants.py

@@ -202,7 +202,8 @@ class WAF_DATA_NAMES(metaclass=Constant_Class):
 
     # EPHEMERAL ADDRESSES
     PROCESSOR_SETTINGS: Literal["waf.context.processor"] = "waf.context.processor"
-    CMDI_ADDRESS: Literal["server.sys.shell.cmd"] = "server.sys.shell.cmd"
+    CMDI_ADDRESS: Literal["server.sys.exec.cmd"] = "server.sys.exec.cmd"
+    SHI_ADDRESS: Literal["server.sys.shell.cmd"] = "server.sys.shell.cmd"
     LFI_ADDRESS: Literal["server.io.fs.file"] = "server.io.fs.file"
     SSRF_ADDRESS: Literal["server.io.net.url"] = "server.io.net.url"
     SQLI_ADDRESS: Literal["server.db.statement"] = "server.db.statement"
@@ -339,13 +340,15 @@ class EXPLOIT_PREVENTION(metaclass=Constant_Class):
 
     class TYPE(metaclass=Constant_Class):
         CMDI: Literal["command_injection"] = "command_injection"
+        SHI: Literal["command_injection"] = "command_injection"
         LFI: Literal["lfi"] = "lfi"
         SSRF: Literal["ssrf"] = "ssrf"
         SQLI: Literal["sql_injection"] = "sql_injection"
 
     class ADDRESS(metaclass=Constant_Class):
         CMDI: Literal["CMDI_ADDRESS"] = "CMDI_ADDRESS"
         LFI: Literal["LFI_ADDRESS"] = "LFI_ADDRESS"
+        SHI: Literal["SHI_ADDRESS"] = "SHI_ADDRESS"
         SSRF: Literal["SSRF_ADDRESS"] = "SSRF_ADDRESS"
         SQLI: Literal["SQLI_ADDRESS"] = "SQLI_ADDRESS"
         SQLI_TYPE: Literal["SQLI_SYSTEM_ADDRESS"] = "SQLI_SYSTEM_ADDRESS"

ddtrace/appsec/_iast/_pytest_plugin.py

@@ -27,6 +27,8 @@ def ddtrace_iast(request, ddspan):
         Optionally output the test as failed if vulnerabilities are found.
         """
         yield
+        if ddspan is None:
+            return
         data = ddspan.get_tag(IAST.JSON)
         if not data:
             return

ddtrace/appsec/_iast/taint_sinks/command_injection.py

@@ -1,18 +1,15 @@
-import os
-import subprocess  # nosec
 from typing import List
 from typing import Union
 
-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
-from ddtrace.appsec._iast._patch import try_wrap_function_wrapper
 from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast.constants import VULN_CMDI
+import ddtrace.contrib.internal.subprocess.patch as subprocess_patch
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
@@ -29,45 +26,14 @@ def get_version() -> str:
 def patch():
     if not asm_config._iast_enabled:
         return
-
-    if not getattr(os, "_datadog_cmdi_patch", False):
-        # all os.spawn* variants eventually use this one:
-        try_wrap_function_wrapper("os", "_spawnvef", _iast_cmdi_osspawn)
-
-    if not getattr(subprocess, "_datadog_cmdi_patch", False):
-        try_wrap_function_wrapper("subprocess", "Popen.__init__", _iast_cmdi_subprocess_init)
-
-        os._datadog_cmdi_patch = True
-        subprocess._datadog_cmdi_patch = True
-
+    subprocess_patch.add_str_callback("iast cmdi", _iast_report_cmdi)
+    subprocess_patch.add_lst_callback("iast cmdi", _iast_report_cmdi)
     _set_metric_iast_instrumented_sink(VULN_CMDI)
 
 
 def unpatch() -> None:
-    try_unwrap("os", "system")
-    try_unwrap("os", "_spawnvef")
-    try_unwrap("subprocess", "Popen.__init__")
-
-    os._datadog_cmdi_patch = False  # type: ignore[attr-defined]
-    subprocess._datadog_cmdi_patch = False  # type: ignore[attr-defined]
-
-
-def _iast_cmdi_osspawn(wrapped, instance, args, kwargs):
-    mode, file, func_args, _, _ = args
-    _iast_report_cmdi(func_args)
-
-    if hasattr(wrapped, "__func__"):
-        return wrapped.__func__(instance, *args, **kwargs)
-    return wrapped(*args, **kwargs)
-
-
-def _iast_cmdi_subprocess_init(wrapped, instance, args, kwargs):
-    cmd_args = args[0] if len(args) else kwargs["args"]
-    _iast_report_cmdi(cmd_args)
-
-    if hasattr(wrapped, "__func__"):
-        return wrapped.__func__(instance, *args, **kwargs)
-    return wrapped(*args, **kwargs)
+    subprocess_patch.del_str_callback("iast cmdi")
+    subprocess_patch.del_lst_callback("iast cmdi")
 
 
 @oce.register

ddtrace/appsec/_processor.py

@@ -202,6 +202,10 @@ def _update_rules(self, new_rules: Dict[str, Any]) -> bool:
     def rasp_lfi_enabled(self) -> bool:
         return WAF_DATA_NAMES.LFI_ADDRESS in self._addresses_to_keep
 
+    @property
+    def rasp_shi_enabled(self) -> bool:
+        return WAF_DATA_NAMES.SHI_ADDRESS in self._addresses_to_keep
+
     @property
     def rasp_cmdi_enabled(self) -> bool:
         return WAF_DATA_NAMES.CMDI_ADDRESS in self._addresses_to_keep

ddtrace/contrib/internal/subprocess/patch.py

@@ -6,6 +6,7 @@
 import shlex
 import subprocess  # nosec
 from threading import RLock
+from typing import Callable  # noqa:F401
 from typing import Deque  # noqa:F401
 from typing import Dict  # noqa:F401
 from typing import List  # noqa:F401
@@ -38,11 +39,30 @@ def get_version():
     return ""
 
 
-def patch():
-    # type: () -> List[str]
-    patched = []  # type: List[str]
-    if not asm_config._asm_enabled:
-        return patched
+_STR_CALLBACKS: Dict[str, Callable[[str], None]] = {}
+_LST_CALLBACKS: Dict[str, Callable[[List[str]], None]] = {}
+
+
+def add_str_callback(name: str, callback: Callable[[str], None]):
+    _STR_CALLBACKS[name] = callback
+
+
+def del_str_callback(name: str):
+    _STR_CALLBACKS.pop(name, None)
+
+
+def add_lst_callback(name: str, callback: Callable[[List[str]], None]):
+    _LST_CALLBACKS[name] = callback
+
+
+def del_lst_callback(name: str):
+    _LST_CALLBACKS.pop(name, None)
+
+
+def patch() -> List[str]:
+    if not (asm_config._asm_enabled or asm_config._iast_enabled):
+        return []
+    patched: List[str] = []
 
     import os
 
@@ -71,7 +91,7 @@ def patch():
 
 
 @dataclass(eq=False)
-class SubprocessCmdLineCacheEntry(object):
+class SubprocessCmdLineCacheEntry:
     binary: Optional[str] = None
     arguments: Optional[List] = None
     truncated: bool = False
@@ -80,7 +100,7 @@ class SubprocessCmdLineCacheEntry(object):
     as_string: Optional[str] = None
 
 
-class SubprocessCmdLine(object):
+class SubprocessCmdLine:
     # This catches the computed values into a SubprocessCmdLineCacheEntry object
     _CACHE = {}  # type: Dict[str, SubprocessCmdLineCacheEntry]
     _CACHE_DEQUE = collections.deque()  # type: Deque[str]
@@ -299,13 +319,16 @@ def unpatch():
 
     SubprocessCmdLine._clear_cache()
 
-    os._datadog_patch = False
-    subprocess._datadog_patch = False
+    os.__dict__.pop("_datadog_patch", None)
+    subprocess.__dict__.pop("_datadog_patch", None)
 
 
 @trace_utils.with_traced_module
 def _traced_ossystem(module, pin, wrapped, instance, args, kwargs):
     try:
+        if isinstance(args[0], str):
+            for callback in _STR_CALLBACKS.values():
+                callback(args[0])
         shellcmd = SubprocessCmdLine(args[0], shell=True)  # nosec
 
         with pin.tracer.trace(COMMANDS.SPAN_NAME, resource=shellcmd.binary, span_type=SpanTypes.SYSTEM) as span:
@@ -342,6 +365,10 @@ def _traced_fork(module, pin, wrapped, instance, args, kwargs):
 def _traced_osspawn(module, pin, wrapped, instance, args, kwargs):
     try:
         mode, file, func_args, _, _ = args
+        if isinstance(func_args, (list, tuple)):
+            commands = [file] + list(func_args)
+            for callback in _LST_CALLBACKS.values():
+                callback(commands)
         shellcmd = SubprocessCmdLine(func_args, shell=False)
 
         with pin.tracer.trace(COMMANDS.SPAN_NAME, resource=shellcmd.binary, span_type=SpanTypes.SYSTEM) as span:
@@ -366,6 +393,12 @@ def _traced_osspawn(module, pin, wrapped, instance, args, kwargs):
 def _traced_subprocess_init(module, pin, wrapped, instance, args, kwargs):
     try:
         cmd_args = args[0] if len(args) else kwargs["args"]
+        if isinstance(cmd_args, str):
+            for callback in _STR_CALLBACKS.values():
+                callback(cmd_args)
+        elif isinstance(cmd_args, list):
+            for callback in _LST_CALLBACKS.values():
+                callback(cmd_args)
         cmd_args_list = shlex.split(cmd_args) if isinstance(cmd_args, str) else cmd_args
         is_shell = kwargs.get("shell", False)
         shellcmd = SubprocessCmdLine(cmd_args_list, shell=is_shell)  # nosec