chore(asm): add SSRF support for urllib.request

.github/CODEOWNERS

@@ -64,6 +64,7 @@ ddtrace/settings/asm.py             @DataDog/asm-python
 ddtrace/contrib/subprocess/         @DataDog/asm-python
 ddtrace/contrib/flask_login/        @DataDog/asm-python
 ddtrace/contrib/webbrowser          @DataDog/asm-python
+ddtrace/contrib/urllib              @DataDog/asm-python
 ddtrace/internal/_exceptions.py     @DataDog/asm-python
 tests/appsec/                       @DataDog/asm-python
 tests/contrib/dbapi/test_dbapi_appsec.py    @DataDog/asm-python

ddtrace/appsec/_iast/taint_sinks/ssrf.py

@@ -24,6 +24,7 @@ class SSRF(VulnerabilityBase):
 _FUNC_TO_URL_ARGUMENT = {
     "http.client.request": (1, "url"),
     "requests.sessions.request": (1, "url"),
+    "urllib.request.urlopen": (0, "url"),
     "urllib3._request_methods.request": (1, "url"),
     "urllib3.request.request": (1, "url"),
     "webbrowser.open": (0, "url"),

ddtrace/contrib/urllib/__init__.py

@@ -0,0 +1,12 @@
+"""
+Trace the standard library ``urllib.request`` library to trace
+HTTP requests and detect SSRF vulnerabilities. It is enabled by default
+if ``DD_IAST_ENABLED`` is set to ``True`` (for detecting sink points) and/or
+``DD_ASM_ENABLED`` is set to ``True`` (for exploit prevention).
+"""
+from .patch import get_version
+from .patch import patch
+from .patch import unpatch
+
+
+__all__ = ["patch", "unpatch", "get_version"]

ddtrace/contrib/urllib/patch.py

@@ -0,0 +1,34 @@
+import urllib.request
+
+from ddtrace.appsec._common_module_patches import wrapped_request_D8CB81E472AF98A2 as _wrap_open
+from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
+from ddtrace.appsec._iast.constants import VULN_SSRF
+from ddtrace.settings.asm import config as asm_config
+from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
+
+from ..trace_utils import unwrap as _u
+
+
+def get_version():
+    # type: () -> str
+    return ""
+
+
+def patch():
+    """patch the built-in urllib.request methods for tracing"""
+    if getattr(urllib.request, "__datadog_patch", False):
+        return
+    urllib.request.__datadog_patch = True
+
+    _w("urllib.request", "urlopen", _wrap_open)
+    if asm_config._iast_enabled:
+        _set_metric_iast_instrumented_sink(VULN_SSRF)
+
+
+def unpatch():
+    """unpatch any previously patched modules"""
+    if not getattr(urllib.request, "__datadog_patch", False):
+        return
+    urllib.request.__datadog_patch = False
+
+    _u(urllib.request, "urlopen")

releasenotes/notes/asm-ssrf-expanded-cc7d8abaa3f9c7dd.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Expand SSRF vulnerability support for Code Security and Exploit Prevention for the modules ``urllib3``, ``http.client``, 
+    ``webbrowser`` and ``urllib.request``.

tests/.suitespec.json

@@ -265,6 +265,9 @@
         "webbrowser": [
             "ddtrace/contrib/webbrowser/*"
         ],
+        "urllib": [
+            "ddtrace/contrib/urllib/*"
+        ],
         "rq": [
             "ddtrace/contrib/rq/*"
         ],
@@ -1298,6 +1301,14 @@
             "@webbrowser",
             "tests/appsec/iast/taint_sinks/test_ssrf.py"
         ],
+        "urllib": [
+            "@bootstrap",
+            "@core",
+            "@contrib",
+            "@tracing",
+            "@urllib",
+            "tests/appsec/iast/taint_sinks/test_ssrf.py"
+        ],
         "vertica": [
             "@bootstrap",
             "@core",

tests/appsec/iast/taint_sinks/test_ssrf.py

@@ -10,6 +10,8 @@
 from ddtrace.contrib.httplib.patch import unpatch as httplib_unpatch
 from ddtrace.contrib.requests.patch import patch as requests_patch
 from ddtrace.contrib.requests.patch import unpatch as requests_unpatch
+from ddtrace.contrib.urllib.patch import patch as urllib_patch
+from ddtrace.contrib.urllib.patch import unpatch as urllib_unpatch
 from ddtrace.contrib.urllib3.patch import patch as urllib3_patch
 from ddtrace.contrib.urllib3.patch import unpatch as urllib3_unpatch
 from ddtrace.contrib.webbrowser.patch import patch as webbrowser_patch
@@ -142,6 +144,26 @@ def test_ssrf_webbrowser(tracer, iast_span_defaults):
             webbrowser_unpatch()
 
 
+def test_urllib_request(tracer, iast_span_defaults):
+    with override_global_config(dict(_iast_enabled=True)):
+        urllib_patch()
+        try:
+            import urllib.request
+
+            tainted_url, tainted_path = _get_tainted_url()
+            try:
+                # label test_urllib_request
+                urllib.request.urlopen(tainted_url)
+            except urllib.error.URLError:
+                pass
+
+            span_report = core.get_item(IAST.CONTEXT_KEY, span=iast_span_defaults)
+            assert span_report
+            _check_report(span_report, tainted_path, "test_urllib_request")
+        finally:
+            urllib_unpatch()
+
+
 def _check_no_report_if_deduplicated(span_report, num_vuln_expected):
     if num_vuln_expected == 0:
         assert span_report is None
@@ -232,3 +254,23 @@ def test_ssrf_webbrowser_deduplication(num_vuln_expected, tracer, iast_span_dedu
         _check_no_report_if_deduplicated(span_report, num_vuln_expected)
     finally:
         webbrowser_unpatch()
+
+
+@pytest.mark.parametrize("num_vuln_expected", [1, 0, 0])
+def test_ssrf_urllib_deduplication(num_vuln_expected, tracer, iast_span_deduplication_enabled):
+    urllib_patch()
+    try:
+        import urllib.request
+
+        tainted_url, tainted_path = _get_tainted_url()
+        for _ in range(0, 5):
+            try:
+                # label test_urllib_request_deduplication
+                urllib.request.urlopen(tainted_url)
+            except urllib.error.URLError:
+                pass
+
+        span_report = core.get_item(IAST.CONTEXT_KEY, span=iast_span_deduplication_enabled)
+        _check_no_report_if_deduplicated(span_report, num_vuln_expected)
+    finally:
+        urllib_unpatch()