chore(iast): removing redundant operations

ddtrace/appsec/_deduplications.py

@@ -32,7 +32,7 @@ def _check_deduplication(self):
         return asm_config._asm_deduplication_enabled
 
     def __call__(self, *args, **kwargs):
-        result = None
+        result = False
         if self._check_deduplication():
             raw_log_hash = hash("".join([str(arg) for arg in self._extract(args)]))
             last_reported_timestamp = self.reported_logs.get(raw_log_hash, M_INF) + self._time_lapse

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -28,10 +28,10 @@ def _set_span_tag_iast_request_tainted(span):
         span.set_tag(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED, total_objects_tainted)
 
 
-def start_iast_context():
+def start_iast_context(span: Optional["Span"] = None):
     if asm_config._iast_enabled:
         create_propagation_context()
-        core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment())
+        core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment(span))
 
 
 def end_iast_context(span: Optional["Span"] = None):
@@ -81,7 +81,7 @@ def _move_iast_data_to_root_span():
 def _iast_start_request(span=None, *args, **kwargs):
     try:
         if asm_config._iast_enabled:
-            start_iast_context()
+            start_iast_context(span)
             request_iast_enabled = False
             if oce.acquire_request(span):
                 request_iast_enabled = True

ddtrace/appsec/_iast/constants.py

@@ -1,6 +1,4 @@
 import re
-from typing import Any
-from typing import Dict
 
 
 VULN_INSECURE_HASHING_TYPE = "WEAK_HASH"
@@ -18,8 +16,6 @@
 VULN_SSRF = "SSRF"
 VULN_STACKTRACE_LEAK = "STACKTRACE_LEAK"
 
-VULNERABILITY_TOKEN_TYPE = Dict[int, Dict[str, Any]]
-
 HEADER_NAME_VALUE_SEPARATOR = ": "
 
 MD5_DEF = "md5"

ddtrace/appsec/_iast/taint_sinks/_base.py

@@ -1,7 +1,5 @@
 import os
 import sysconfig
-from typing import Any
-from typing import Callable
 from typing import Optional
 from typing import Tuple
 from typing import Union
@@ -18,7 +16,6 @@
 from .._iast_request_context import set_iast_reporter
 from .._overhead_control_engine import Operation
 from .._stacktrace import get_info_frame
-from .._utils import _is_iast_debug_enabled
 from ..reporter import Evidence
 from ..reporter import IastSpanReporter
 from ..reporter import Location
@@ -60,26 +57,6 @@ class VulnerabilityBase(Operation):
     vulnerability_type = ""
     secure_mark = 0
 
-    @classmethod
-    def wrap(cls, func: Callable) -> Callable:
-        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 asm_config.is_iast_request_enabled:
-                if _is_iast_debug_enabled():
-                    log.debug(
-                        "iast::propagation::context::VulnerabilityBase.wrapper. No request quota or this vulnerability "
-                        "is outside the context"
-                    )
-                return wrapped(*args, **kwargs)
-            elif cls.has_quota():
-                return func(wrapped, instance, args, kwargs)
-            else:
-                return wrapped(*args, **kwargs)
-
-        return wrapper
-
     @classmethod
     @taint_sink_deduplication
     def _prepare_report(
@@ -93,13 +70,6 @@ def _prepare_report(
         *args,
         **kwargs,
     ) -> bool:
-        if not asm_config.is_iast_request_enabled:
-            if _is_iast_debug_enabled():
-                log.debug(
-                    "iast::propagation::context::VulnerabilityBase._prepare_report. "
-                    "No request quota or this vulnerability is outside the context"
-                )
-            return False
         if line_number is not None and (line_number == 0 or line_number < -1):
             line_number = -1
 

ddtrace/appsec/_iast/taint_sinks/code_injection.py

@@ -49,7 +49,7 @@ def unpatch():
 
 
 def _iast_coi(wrapped, instance, args, kwargs):
-    if asm_config._iast_enabled and len(args) >= 1:
+    if len(args) >= 1 and asm_config.is_iast_request_enabled:
         _iast_report_code_injection(args[0])
 
     caller_frame = None
@@ -85,7 +85,7 @@ def _iast_report_code_injection(code_string: Text):
     reported = False
     try:
         if asm_config.is_iast_request_enabled:
-            if isinstance(code_string, IAST.TEXT_TYPES) and CodeInjection.has_quota():
+            if code_string and isinstance(code_string, IAST.TEXT_TYPES) and CodeInjection.has_quota():
                 if CodeInjection.is_tainted_pyobject(code_string):
                     CodeInjection.report(evidence_value=code_string)
 

ddtrace/appsec/_iast/taint_sinks/header_injection.py

@@ -101,7 +101,7 @@ def unpatch():
 
 
 def _iast_h(wrapped, instance, args, kwargs):
-    if asm_config._iast_enabled and args:
+    if asm_config.is_iast_request_enabled:
         _iast_report_header_injection(args)
     if hasattr(wrapped, "__func__"):
         return wrapped.__func__(instance, *args, **kwargs)
@@ -129,17 +129,16 @@ def _process_header(headers_args):
             if header_name_lower == header_to_exclude or header_name_lower.startswith(header_to_exclude):
                 return
 
-        if asm_config.is_iast_request_enabled:
-            if HeaderInjection.has_quota() and (
-                HeaderInjection.is_tainted_pyobject(header_name) or HeaderInjection.is_tainted_pyobject(header_value)
-            ):
-                header_evidence = add_aspect(add_aspect(header_name, HEADER_NAME_VALUE_SEPARATOR), header_value)
-                HeaderInjection.report(evidence_value=header_evidence)
-
-            # Reports Span Metrics
-            increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, HeaderInjection.vulnerability_type)
-            # Report Telemetry Metrics
-            _set_metric_iast_executed_sink(HeaderInjection.vulnerability_type)
+        if HeaderInjection.has_quota() and (
+            HeaderInjection.is_tainted_pyobject(header_name) or HeaderInjection.is_tainted_pyobject(header_value)
+        ):
+            header_evidence = add_aspect(add_aspect(header_name, HEADER_NAME_VALUE_SEPARATOR), header_value)
+            HeaderInjection.report(evidence_value=header_evidence)
+
+        # Reports Span Metrics
+        increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, HeaderInjection.vulnerability_type)
+        # Report Telemetry Metrics
+        _set_metric_iast_executed_sink(HeaderInjection.vulnerability_type)
     except Exception as e:
         iast_error(f"propagation::sink_point::Error in _iast_report_header_injection. {e}")
 

ddtrace/appsec/_iast/taint_sinks/insecure_cookie.py

@@ -150,7 +150,7 @@ def _iast_response_cookies(wrapped, instance, args, kwargs):
             cookie_value = kwargs.get("value")
 
         if cookie_value and cookie_key:
-            if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+            if asm_config.is_iast_request_enabled and CookiesVulnerability.has_quota():
                 report_samesite = False
                 samesite = kwargs.get("samesite", "")
                 if samesite:

ddtrace/appsec/_iast/taint_sinks/sql_injection.py

@@ -25,7 +25,7 @@ class SqlInjection(VulnerabilityBase):
 
 
 def check_and_report_sqli(
-    args: Tuple[Text], kwargs: Dict[str, Any], integration_name: Text, method: Callable[..., Any]
+    args: Tuple[Text, ...], kwargs: Dict[str, Any], integration_name: Text, method: Callable[..., Any]
 ) -> bool:
     """Check for SQL injection vulnerabilities in database operations and report them.
 
@@ -41,14 +41,8 @@ def check_and_report_sqli(
     reported = False
     try:
         if supported_dbapi_integration(integration_name) and method.__name__ == "execute":
-            if (
-                len(args)
-                and args[0]
-                and isinstance(args[0], IAST.TEXT_TYPES)
-                and asm_config.is_iast_request_enabled
-                and SqlInjection.has_quota()
-            ):
-                if SqlInjection.is_tainted_pyobject(args[0]):
+            if len(args) and args[0] and isinstance(args[0], IAST.TEXT_TYPES) and asm_config.is_iast_request_enabled:
+                if SqlInjection.has_quota() and SqlInjection.is_tainted_pyobject(args[0]):
                     SqlInjection.report(evidence_value=args[0], dialect=integration_name)
                     reported = True
 

ddtrace/appsec/_iast/taint_sinks/weak_cipher.py

@@ -130,12 +130,12 @@ def wrapped_aux_blowfish_function(wrapped, instance, args, kwargs):
     return result
 
 
-@WeakCipher.wrap
 def wrapped_rc4_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
-        WeakCipher.report(
-            evidence_value="RC4",
-        )
+        if WeakCipher.has_quota():
+            WeakCipher.report(
+                evidence_value="RC4",
+            )
         # Reports Span Metrics
         increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)
         # Report Telemetry Metrics
@@ -146,12 +146,12 @@ def wrapped_rc4_function(wrapped: Callable, instance: Any, args: Any, kwargs: An
     return wrapped(*args, **kwargs)
 
 
-@WeakCipher.wrap
 def wrapped_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
         if hasattr(instance, "_dd_weakcipher_algorithm"):
-            evidence = instance._dd_weakcipher_algorithm + "_" + str(instance.__class__.__name__)
-            WeakCipher.report(evidence_value=evidence)
+            if WeakCipher.has_quota():
+                evidence = instance._dd_weakcipher_algorithm + "_" + str(instance.__class__.__name__)
+                WeakCipher.report(evidence_value=evidence)
 
             # Reports Span Metrics
             increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)
@@ -163,14 +163,14 @@ def wrapped_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -
     return wrapped(*args, **kwargs)
 
 
-@WeakCipher.wrap
 def wrapped_cryptography_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
         algorithm_name = instance.algorithm.name.lower()
         if algorithm_name in get_weak_cipher_algorithms():
-            WeakCipher.report(
-                evidence_value=algorithm_name,
-            )
+            if WeakCipher.has_quota():
+                WeakCipher.report(
+                    evidence_value=algorithm_name,
+                )
 
             # Reports Span Metrics
             increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)

ddtrace/appsec/_iast/taint_sinks/weak_hash.py

@@ -1,13 +1,10 @@
 import os
 import sys
-from typing import TYPE_CHECKING  # noqa:F401
 from typing import Any
 from typing import Callable
 from typing import Set
-from typing import Text  # noqa:F401
 
 from ddtrace.appsec._common_module_patches import try_unwrap
-from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
 from ..._constants import IAST_SPAN_TAGS
@@ -25,9 +22,6 @@
 from ._base import VulnerabilityBase
 
 
-log = get_logger(__name__)
-
-
 def get_weak_hash_algorithms() -> Set:
     CONFIGURED_WEAK_HASH_ALGORITHMS = None
     DD_IAST_WEAK_HASH_ALGORITHMS = os.getenv("DD_IAST_WEAK_HASH_ALGORITHMS")
@@ -65,7 +59,7 @@ def unpatch_iast():
     try_unwrap("Crypto.Hash.SHA1", "SHA1Hash.hexdigest")
 
 
-def get_version() -> Text:
+def get_version() -> str:
     return ""
 
 
@@ -119,35 +113,31 @@ def patch():
         _set_metric_iast_instrumented_sink(VULN_INSECURE_HASHING_TYPE, num_instrumented_sinks)
 
 
-@WeakHash.wrap
 def wrapped_digest_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
         if WeakHash.has_quota() and instance.name.lower() in get_weak_hash_algorithms():
             WeakHash.report(
                 evidence_value=instance.name,
             )
 
-            # Reports Span Metrics
-            increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
-            # Report Telemetry Metrics
-            _set_metric_iast_executed_sink(WeakHash.vulnerability_type)
+        # Reports Span Metrics
+        increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
+        # Report Telemetry Metrics
+        _set_metric_iast_executed_sink(WeakHash.vulnerability_type)
 
     if hasattr(wrapped, "__func__"):
         return wrapped.__func__(instance, *args, **kwargs)
     return wrapped(*args, **kwargs)
 
 
-@WeakHash.wrap
 def wrapped_md5_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     return wrapped_function(wrapped, MD5_DEF, instance, args, kwargs)
 
 
-@WeakHash.wrap
 def wrapped_sha1_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     return wrapped_function(wrapped, SHA1_DEF, instance, args, kwargs)
 
 
-@WeakHash.wrap
 def wrapped_new_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
         if WeakHash.has_quota() and args[0].lower() in get_weak_hash_algorithms():
@@ -160,11 +150,12 @@ def wrapped_new_function(wrapped: Callable, instance: Any, args: Any, kwargs: An
     return wrapped(*args, **kwargs)
 
 
-def wrapped_function(wrapped: Callable, evidence: Text, instance: Any, args: Any, kwargs: Any) -> Any:
+def wrapped_function(wrapped: Callable, evidence: str, instance: Any, args: Any, kwargs: Any) -> Any:
     if asm_config.is_iast_request_enabled:
-        WeakHash.report(
-            evidence_value=evidence,
-        )
+        if WeakHash.has_quota():
+            WeakHash.report(
+                evidence_value=evidence,
+            )
         # Reports Span Metrics
         increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
         # Report Telemetry Metrics

tests/appsec/iast/test_processor.py

@@ -3,7 +3,6 @@
 import pytest
 
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec._iast._iast_request_context import get_iast_reporter
 from ddtrace.appsec._iast._overhead_control_engine import oce
 from ddtrace.constants import _SAMPLING_PRIORITY_KEY
 from ddtrace.constants import AUTO_KEEP
@@ -27,7 +26,6 @@ def traced_function(tracer):
     return span
 
 
-@pytest.mark.skip_iast_check_logs
 def test_appsec_iast_processor(iast_context_defaults):
     """
     test_appsec_iast_processor.
@@ -38,11 +36,9 @@ def test_appsec_iast_processor(iast_context_defaults):
     span = traced_function(tracer)
     tracer._on_span_finish(span)
 
-    span_report = get_iast_reporter()
     result = span.get_tag(IAST.JSON)
 
     assert len(json.loads(result)["vulnerabilities"]) == 1
-    assert len(span_report.vulnerabilities) == 1
 
 
 @pytest.mark.parametrize("sampling_rate", ["0.0", "0.5", "1.0"])

tests/appsec/iast/test_telemetry.py

@@ -87,7 +87,7 @@ def test_metric_executed_sink(no_request_sampling, telemetry_writer, caplog):
     # the agent)
     filtered_metrics = [metric for metric in generate_metrics if metric["tags"][0] == "vulnerability_type:weak_hash"]
     assert [metric["tags"] for metric in filtered_metrics] == [["vulnerability_type:weak_hash"]]
-    assert span.get_metric("_dd.iast.telemetry.executed.sink.weak_hash") == 2
+    assert span.get_metric("_dd.iast.telemetry.executed.sink.weak_hash") == 10
     # request.tainted metric is None because AST is not running in this test
     assert span.get_metric(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED) is None