chore(appsec): move iast sqli logic (#13748)

Remove appsec logic from contrib files


## 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/appsec/_iast/_listener.py

@@ -14,6 +14,7 @@
 from ddtrace.appsec._iast._handlers import _on_wsgi_environ
 from ddtrace.appsec._iast._iast_request_context import _iast_end_request
 from ddtrace.appsec._iast._langchain import langchain_listen
+from ddtrace.appsec._iast.taint_sinks.sql_injection import _on_report_sqli
 from ddtrace.internal import core
 
 
@@ -39,6 +40,9 @@ def iast_listen():
     core.on("context.ended.wsgi.__call__", _iast_end_request)
     core.on("context.ended.asgi.__call__", _iast_end_request)
 
+    # Sink points
+    core.on("db_query_check", _on_report_sqli)
+
     langchain_listen(core)
 
 

ddtrace/appsec/_iast/taint_sinks/sql_injection.py

@@ -1,9 +1,3 @@
-from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import Text
-from typing import Tuple
-
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._logs import iast_error
@@ -22,9 +16,7 @@ class SqlInjection(VulnerabilityBase):
     secure_mark = VulnerabilityType.SQL_INJECTION
 
 
-def check_and_report_sqli(
-    args: Tuple[Text, ...], kwargs: Dict[str, Any], integration_name: Text, method: Callable[..., Any]
-) -> bool:
+def _on_report_sqli(*args, **kwargs) -> bool:
     """Check for SQL injection vulnerabilities in database operations and report them.
 
     This function analyzes database operation arguments for potential SQL injection
@@ -36,18 +28,27 @@ def check_and_report_sqli(
         This function is part of the IAST (Interactive Application Security Testing)
         system and is used to detect potential SQL injection vulnerabilities at runtime.
     """
+
     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:
-                if SqlInjection.has_quota() and SqlInjection.is_tainted_pyobject(args[0]):
-                    SqlInjection.report(evidence_value=args[0], dialect=integration_name)
-                    reported = True
-
-                # Reports Span Metrics
-                increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, SqlInjection.vulnerability_type)
-                # Report Telemetry Metrics
-                _set_metric_iast_executed_sink(SqlInjection.vulnerability_type)
+        if asm_config._iast_enabled:
+            query_args, kwargs, integration_name, method = args
+
+            if supported_dbapi_integration(integration_name) and method.__name__ == "execute":
+                if (
+                    len(query_args)
+                    and query_args[0]
+                    and isinstance(query_args[0], IAST.TEXT_TYPES)
+                    and asm_config.is_iast_request_enabled
+                ):
+                    if SqlInjection.has_quota() and SqlInjection.is_tainted_pyobject(query_args[0]):
+                        SqlInjection.report(evidence_value=query_args[0], dialect=integration_name)
+                        reported = True
+
+                    # Reports Span Metrics
+                    increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, SqlInjection.vulnerability_type)
+                    # Report Telemetry Metrics
+                    _set_metric_iast_executed_sink(SqlInjection.vulnerability_type)
     except Exception as e:
         iast_error(f"propagation::sink_point::Error in check_and_report_sqli. {e}")
     return reported

ddtrace/contrib/dbapi.py

@@ -9,7 +9,6 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
-from ddtrace.settings.asm import config as asm_config
 
 from ..constants import _SPAN_MEASURED_KEY
 from ..constants import SPAN_KIND
@@ -100,10 +99,9 @@ def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *arg
             # set span.kind to the type of request being performed
             s.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
 
-            if asm_config._iast_enabled:
-                from ddtrace.appsec._iast.taint_sinks.sql_injection import check_and_report_sqli
+            # Security and IAST validations
+            core.dispatch("db_query_check", (args, kwargs, self._self_config.integration_name, method))
 
-                check_and_report_sqli(args, kwargs, self._self_config.integration_name, method)
             # dispatch DBM
             if dbm_propagator:
                 # this check is necessary to prevent fetch methods from trying to add dbm propagation

ddtrace/contrib/dbapi_async.py

@@ -4,7 +4,6 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
-from ddtrace.settings.asm import config as asm_config
 
 from ..constants import _SPAN_MEASURED_KEY
 from ..constants import SPAN_KIND
@@ -75,16 +74,14 @@ async def _trace_method(self, method, name, resource, extra_tags, dbm_propagator
             # set span.kind to the type of request being performed
             s.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
 
-            if asm_config._iast_enabled:
-                from ddtrace.appsec._iast.taint_sinks.sql_injection import check_and_report_sqli
-
-                check_and_report_sqli(args, kwargs, self._self_config.integration_name, method)
+            # Security and IAST validations
+            core.dispatch("db_query_check", (args, kwargs, self._self_config.integration_name, method))
 
             # dispatch DBM
             if dbm_propagator:
                 # this check is necessary to prevent fetch methods from trying to add dbm propagation
                 result = core.dispatch_with_results(
-                    f"{self._self_config.integration_name}.execute", [self._self_config, s, args, kwargs]
+                    f"{self._self_config.integration_name}.execute", (self._self_config, s, args, kwargs)
                 ).result
                 if result:
                     s, args, kwargs = result.value

tests/appsec/iast/taint_sinks/test_sql_injection.py

@@ -6,7 +6,7 @@
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
-from ddtrace.appsec._iast.taint_sinks.sql_injection import check_and_report_sqli
+from ddtrace.appsec._iast.taint_sinks.sql_injection import _on_report_sqli
 
 
 def test_checked_tainted_args(iast_context_deduplication_enabled):
@@ -21,45 +21,24 @@ def test_checked_tainted_args(iast_context_deduplication_enabled):
     untainted_arg = "gallahad the pure"
 
     # Returns False: Untainted first argument
-    assert not check_and_report_sqli(
-        args=(untainted_arg,), kwargs=None, integration_name="sqlite", method=cursor.execute
-    )
+    assert not _on_report_sqli((untainted_arg,), None, "sqlite", cursor.execute)
 
     # Returns False: Untainted first argument
-    assert not check_and_report_sqli(
-        args=(untainted_arg, tainted_arg), kwargs=None, integration_name="sqlite", method=cursor.execute
-    )
-
+    assert not _on_report_sqli((untainted_arg, tainted_arg), None, "sqlite", cursor.execute)
     # Returns False: Integration name not in list
-    assert not check_and_report_sqli(
-        args=(tainted_arg,),
-        kwargs=None,
-        integration_name="nosqlite",
-        method=cursor.execute,
-    )
+    assert not _on_report_sqli((tainted_arg,), None, "nosqlite", cursor.execute)
 
     # Returns False: Wrong function name
-    assert not check_and_report_sqli(
-        args=(tainted_arg,),
-        kwargs=None,
-        integration_name="sqlite",
-        method=cursor.executemany,
-    )
+    assert not _on_report_sqli((tainted_arg,), None, "sqlite", cursor.executemany)
 
     # Returns True:
-    assert check_and_report_sqli(
-        args=(tainted_arg, untainted_arg), kwargs=None, integration_name="sqlite", method=cursor.execute
-    )
+    assert _on_report_sqli((tainted_arg, untainted_arg), None, "sqlite", cursor.execute)
 
     # Returns True:
-    assert check_and_report_sqli(
-        args=(tainted_arg, untainted_arg), kwargs=None, integration_name="mysql", method=cursor.execute
-    )
+    assert _on_report_sqli((tainted_arg, untainted_arg), None, "mysql", cursor.execute)
 
     # Returns False: No more QUOTA
-    assert not check_and_report_sqli(
-        args=(tainted_arg, untainted_arg), kwargs=None, integration_name="psycopg", method=cursor.execute
-    )
+    assert not _on_report_sqli((tainted_arg, untainted_arg), None, "psycopg", cursor.execute)
 
 
 @pytest.mark.parametrize(
@@ -116,7 +95,7 @@ def test_check_and_report_sqli_metrics(args, integration_name, expected_result,
         "ddtrace.appsec._iast.taint_sinks.sql_injection._set_metric_iast_executed_sink"
     ) as mock_set_metric:
         # Call with tainted argument that should trigger metrics
-        result = check_and_report_sqli(args=args, kwargs={}, integration_name=integration_name, method=cursor.execute)
+        result = _on_report_sqli(args, {}, integration_name, cursor.execute)
 
         assert result is expected_result
         mock_increment.assert_called_once()
@@ -154,7 +133,7 @@ def test_check_and_report_sqli_no_metrics(args, integration_name, iast_context_d
         "ddtrace.appsec._iast.taint_sinks.sql_injection._set_metric_iast_executed_sink"
     ) as mock_set_metric:
         # Call with untainted argument that should not trigger metrics
-        result = check_and_report_sqli(args=args, kwargs={}, integration_name=integration_name, method=cursor.execute)
+        result = _on_report_sqli(args, {}, integration_name, cursor.execute)
 
         assert result is False
         mock_increment.assert_not_called()

tests/contrib/dbapi/test_dbapi_appsec.py renamed to tests/appsec/iast/taint_sinks/test_sql_injection_dbapi.py

@@ -1,6 +1,8 @@
-import mock
+from unittest import mock
+
 import pytest
 
+from ddtrace.appsec._iast import load_iast
 from ddtrace.appsec._iast._overhead_control_engine import oce
 from ddtrace.contrib.dbapi import TracedCursor
 from ddtrace.settings._config import Config
@@ -16,6 +18,7 @@
 class TestTracedCursor(TracerTestCase):
     def setUp(self):
         super(TestTracedCursor, self).setUp()
+        load_iast()
         with override_global_config(
             dict(
                 _iast_enabled=True,

tests/appsec/integrations/django_tests/conftest.py

@@ -5,6 +5,7 @@
 import pytest
 
 from ddtrace.appsec._iast import enable_iast_propagation
+from ddtrace.appsec._iast import load_iast
 from ddtrace.appsec._iast.main import patch_iast
 from ddtrace.contrib.internal.django.patch import patch as django_patch
 from ddtrace.contrib.internal.requests.patch import patch as requests_patch
@@ -32,6 +33,7 @@ def pytest_configure():
     ), override_env(dict(_DD_IAST_PATCH_MODULES="tests.appsec.integrations")):
         settings.DEBUG = False
         patch_iast()
+        load_iast()
         requests_patch()
         django_patch()
         enable_iast_propagation()

tests/appsec/integrations/packages_tests/test_iast_sql_injection.py

@@ -1,6 +1,7 @@
 import pytest
 
 from ddtrace import patch
+from ddtrace.appsec._iast import load_iast
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import is_pyobject_tainted
@@ -37,6 +38,7 @@
 
 def setup_module():
     patch(pymysql=True, mysqldb=True)
+    load_iast()
 
 
 @pytest.mark.parametrize("fixture_path,fixture_module", DDBBS)