refactor(iast): refactor imports to make them local (#6899)

gnufede · web-flow · commit 6b16ac19956c · 2023-09-13T13:33:49.000+02:00
## Checklist
- [x] This PR doesn't touch any of that.
diff --git a/ddtrace/appsec/ddwaf/__init__.py b/ddtrace/appsec/ddwaf/__init__.py
@@ -12,7 +12,7 @@
     from typing import Any
     from typing import Union
 
-    from ddtrace.appsec.ddwaf.ddwaf_types import DDWafRulesType
+    from .ddwaf_types import DDWafRulesType
 
 
 LOGGER = get_logger(__name__)
diff --git a/ddtrace/appsec/iast/__init__.py b/ddtrace/appsec/iast/__init__.py
@@ -31,11 +31,12 @@ def wrapped_function(wrapped, instance, args, kwargs):
 import inspect
 import sys
 
-from ddtrace.appsec.iast._ast.ast_patching import astpatch_module
-from ddtrace.appsec.iast._overhead_control_engine import OverheadControl
-from ddtrace.appsec.iast._utils import _is_iast_enabled
 from ddtrace.internal.logger import get_logger
 
+from ._ast.ast_patching import astpatch_module
+from ._overhead_control_engine import OverheadControl
+from ._utils import _is_iast_enabled
+
 
 log = get_logger(__name__)
 
diff --git a/ddtrace/appsec/iast/_ast/ast_patching.py b/ddtrace/appsec/iast/_ast/ast_patching.py
@@ -15,10 +15,11 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._python_info.stdlib import _stdlib_for_python_version
-from ddtrace.appsec.iast._ast.visitor import AstVisitor
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.module import origin
 
+from .visitor import AstVisitor
+
 
 # Prefixes for modules where IAST patching is allowed
 IAST_ALLOWLIST = ("tests.appsec.iast",)  # type: tuple[str, ...]
diff --git a/ddtrace/appsec/iast/_ast/visitor.py b/ddtrace/appsec/iast/_ast/visitor.py
@@ -8,7 +8,7 @@
 
 from six import iteritems
 
-from ddtrace.appsec.iast._metrics import _set_metric_iast_instrumented_propagation
+from .._metrics import _set_metric_iast_instrumented_propagation
 
 
 if TYPE_CHECKING:  # pragma: no cover
diff --git a/ddtrace/appsec/iast/_metrics.py b/ddtrace/appsec/iast/_metrics.py
@@ -63,7 +63,7 @@ def _set_iast_error_metric(msg, stack_trace):
 
 @metric_verbosity(TELEMETRY_MANDATORY_VERBOSITY)
 def _set_metric_iast_instrumented_source(source_type):
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import origin_to_str  # noqa: F401
+    from ._taint_tracking._native.taint_tracking import origin_to_str  # noqa: F401
 
     telemetry_writer.add_count_metric(
         TELEMETRY_NAMESPACE_TAG_IAST, "instrumented.source", 1, (("source_type", origin_to_str(source_type)),)
@@ -84,7 +84,7 @@ def _set_metric_iast_instrumented_sink(vulnerability_type, counter=1):
 
 @metric_verbosity(TELEMETRY_INFORMATION_VERBOSITY)
 def _set_metric_iast_executed_source(source_type):
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import origin_to_str  # noqa: F401
+    from ._taint_tracking._native.taint_tracking import origin_to_str  # noqa: F401
 
     telemetry_writer.add_count_metric(
         TELEMETRY_NAMESPACE_TAG_IAST, "executed.source", 1, (("source_type", origin_to_str(source_type)),)
@@ -100,7 +100,7 @@ def _set_metric_iast_executed_sink(vulnerability_type):
 
 @metric_verbosity(TELEMETRY_INFORMATION_VERBOSITY)
 def _set_metric_iast_request_tainted():
-    from ddtrace.appsec.iast._taint_tracking import num_objects_tainted
+    from ._taint_tracking import num_objects_tainted
 
     total_objects_tainted = num_objects_tainted()
     if total_objects_tainted > 0:
diff --git a/ddtrace/appsec/iast/_patch.py b/ddtrace/appsec/iast/_patch.py
@@ -3,11 +3,12 @@
 import sys
 from typing import TYPE_CHECKING
 
-from ddtrace.appsec.iast._utils import _is_iast_enabled
 from ddtrace.internal.logger import get_logger
 from ddtrace.vendor.wrapt import FunctionWrapper
 from ddtrace.vendor.wrapt import resolve_path
 
+from ._utils import _is_iast_enabled
+
 
 if TYPE_CHECKING:  # pragma: no cover
     from typing import Any
@@ -142,8 +143,8 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
 
     if _is_iast_enabled():
         try:
-            from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
-            from ddtrace.appsec.iast._taint_tracking import taint_pyobject
+            from ._taint_tracking import is_pyobject_tainted
+            from ._taint_tracking import taint_pyobject
 
             if not is_pyobject_tainted(value):
                 name = str(args[0]) if len(args) else "http.request.body"
@@ -155,7 +156,7 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
 
 def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
     if _is_iast_enabled():
-        from ddtrace.appsec.iast._taint_tracking import taint_pyobject
+        from ._taint_tracking import taint_pyobject
 
         for key, value in wrapped(*args, **kwargs):
             new_key = taint_pyobject(pyobject=key, source_name=key, source_value=key, source_origin=origins[0])
diff --git a/ddtrace/appsec/iast/_patches/json_tainting.py b/ddtrace/appsec/iast/_patches/json_tainting.py
@@ -1,12 +1,13 @@
 from ddtrace import config
-from ddtrace.appsec.iast._patch import set_and_check_module_is_patched
-from ddtrace.appsec.iast._patch import set_module_unpatched
-from ddtrace.appsec.iast._patch import try_unwrap
-from ddtrace.appsec.iast._patch import try_wrap_function_wrapper
-from ddtrace.appsec.iast._taint_utils import LazyTaintDict
-from ddtrace.appsec.iast._taint_utils import LazyTaintList
 from ddtrace.internal.logger import get_logger
 
+from .._patch import set_and_check_module_is_patched
+from .._patch import set_module_unpatched
+from .._patch import try_unwrap
+from .._patch import try_wrap_function_wrapper
+from .._taint_utils import LazyTaintDict
+from .._taint_utils import LazyTaintList
+
 
 log = get_logger(__name__)
 
@@ -37,9 +38,9 @@ def wrapped_loads(wrapped, instance, args, kwargs):
     obj = wrapped(*args, **kwargs)
     if config._iast_enabled:
         try:
-            from ddtrace.appsec.iast._taint_tracking import get_tainted_ranges
-            from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
-            from ddtrace.appsec.iast._taint_tracking import taint_pyobject
+            from .._taint_tracking import get_tainted_ranges
+            from .._taint_tracking import is_pyobject_tainted
+            from .._taint_tracking import taint_pyobject
 
             if is_pyobject_tainted(args[0]) and obj:
                 # tainting object
diff --git a/ddtrace/appsec/iast/_taint_dict.py b/ddtrace/appsec/iast/_taint_dict.py
@@ -7,7 +7,7 @@
     from typing import Dict
     from typing import Tuple
 
-    from ddtrace.appsec.iast._taint_tracking import Source
+    from ._taint_tracking import Source
 
 _IAST_TAINT_DICT = {}  # type: Dict[int, Tuple[Tuple[Source, int, int],...]]
 
diff --git a/ddtrace/appsec/iast/_taint_tracking/__init__.py b/ddtrace/appsec/iast/_taint_tracking/__init__.py
@@ -2,38 +2,38 @@
 # flake8: noqa
 from typing import TYPE_CHECKING
 
-from ddtrace.appsec.iast._metrics import _set_metric_iast_executed_source
-from ddtrace.appsec.iast._utils import _is_python_version_supported
+from .._metrics import _set_metric_iast_executed_source
+from .._utils import _is_python_version_supported
 
 
 if _is_python_version_supported():
-    from ddtrace.appsec.iast import oce
-    from ddtrace.appsec.iast._taint_tracking._native import ops
-    from ddtrace.appsec.iast._taint_tracking._native.aspect_helpers import _convert_escaped_text_to_tainted_text
-    from ddtrace.appsec.iast._taint_tracking._native.aspect_helpers import as_formatted_evidence
-    from ddtrace.appsec.iast._taint_tracking._native.aspect_helpers import common_replace
-    from ddtrace.appsec.iast._taint_tracking._native.aspect_helpers import parse_params
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import active_map_addreses_size
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import create_context
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import destroy_context
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import initializer_size
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import num_objects_tainted
-    from ddtrace.appsec.iast._taint_tracking._native.initializer import reset_context
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import OriginType
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import Source
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import TagMappingMode
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import are_all_text_all_ranges
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import get_range_by_hash
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import get_ranges
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import is_notinterned_notfasttainted_unicode
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import is_tainted
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import origin_to_str
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import set_fast_tainted_if_notinterned_unicode
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import set_ranges
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import shift_taint_range
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import shift_taint_ranges
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import str_to_origin
-    from ddtrace.appsec.iast._taint_tracking._native.taint_tracking import taint_range as TaintRange
+    from .. import oce
+    from ._native import ops
+    from ._native.aspect_helpers import _convert_escaped_text_to_tainted_text
+    from ._native.aspect_helpers import as_formatted_evidence
+    from ._native.aspect_helpers import common_replace
+    from ._native.aspect_helpers import parse_params
+    from ._native.initializer import active_map_addreses_size
+    from ._native.initializer import create_context
+    from ._native.initializer import destroy_context
+    from ._native.initializer import initializer_size
+    from ._native.initializer import num_objects_tainted
+    from ._native.initializer import reset_context
+    from ._native.taint_tracking import OriginType
+    from ._native.taint_tracking import Source
+    from ._native.taint_tracking import TagMappingMode
+    from ._native.taint_tracking import are_all_text_all_ranges
+    from ._native.taint_tracking import get_range_by_hash
+    from ._native.taint_tracking import get_ranges
+    from ._native.taint_tracking import is_notinterned_notfasttainted_unicode
+    from ._native.taint_tracking import is_tainted
+    from ._native.taint_tracking import origin_to_str
+    from ._native.taint_tracking import set_fast_tainted_if_notinterned_unicode
+    from ._native.taint_tracking import set_ranges
+    from ._native.taint_tracking import shift_taint_range
+    from ._native.taint_tracking import shift_taint_ranges
+    from ._native.taint_tracking import str_to_origin
+    from ._native.taint_tracking import taint_range as TaintRange
 
     new_pyobject_id = ops.new_pyobject_id
     is_pyobject_tainted = is_tainted
diff --git a/ddtrace/appsec/iast/_taint_tracking/aspects.py b/ddtrace/appsec/iast/_taint_tracking/aspects.py
@@ -5,22 +5,23 @@
 import traceback
 from typing import TYPE_CHECKING
 
-from ddtrace.appsec.iast._metrics import _set_iast_error_metric
-from ddtrace.appsec.iast._taint_tracking import TagMappingMode
-from ddtrace.appsec.iast._taint_tracking import TaintRange
-from ddtrace.appsec.iast._taint_tracking import _convert_escaped_text_to_tainted_text
-from ddtrace.appsec.iast._taint_tracking import are_all_text_all_ranges
-from ddtrace.appsec.iast._taint_tracking import as_formatted_evidence
-from ddtrace.appsec.iast._taint_tracking import common_replace
-from ddtrace.appsec.iast._taint_tracking import get_ranges
-from ddtrace.appsec.iast._taint_tracking import get_tainted_ranges
-from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
-from ddtrace.appsec.iast._taint_tracking import parse_params
-from ddtrace.appsec.iast._taint_tracking import shift_taint_range
-from ddtrace.appsec.iast._taint_tracking import taint_pyobject_with_ranges
-from ddtrace.appsec.iast._taint_tracking._native import aspects  # noqa: F401
 from ddtrace.internal.compat import iteritems
 
+from .._metrics import _set_iast_error_metric
+from .._taint_tracking import TagMappingMode
+from .._taint_tracking import TaintRange
+from .._taint_tracking import _convert_escaped_text_to_tainted_text
+from .._taint_tracking import are_all_text_all_ranges
+from .._taint_tracking import as_formatted_evidence
+from .._taint_tracking import common_replace
+from .._taint_tracking import get_ranges
+from .._taint_tracking import get_tainted_ranges
+from .._taint_tracking import is_pyobject_tainted
+from .._taint_tracking import parse_params
+from .._taint_tracking import shift_taint_range
+from .._taint_tracking import taint_pyobject_with_ranges
+from .._taint_tracking._native import aspects  # noqa: F401
+
 
 if TYPE_CHECKING:
     from typing import Any
diff --git a/ddtrace/appsec/iast/_taint_utils.py b/ddtrace/appsec/iast/_taint_utils.py
@@ -30,8 +30,8 @@ def __init__(self, original_list, origins=(0, 0), override_pyobject_tainted=Fals
     def _taint(self, value):
         if value:
             if isinstance(value, (str, bytes, bytearray)):
-                from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
-                from ddtrace.appsec.iast._taint_tracking import taint_pyobject
+                from ._taint_tracking import is_pyobject_tainted
+                from ._taint_tracking import taint_pyobject
 
                 if not is_pyobject_tainted(value) or self._override_pyobject_tainted:
                     try:
@@ -201,8 +201,8 @@ def _taint(self, value, key, origin=None):
             origin = self._origin_value
         if value:
             if isinstance(value, (str, bytes, bytearray)):
-                from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
-                from ddtrace.appsec.iast._taint_tracking import taint_pyobject
+                from ._taint_tracking import is_pyobject_tainted
+                from ._taint_tracking import taint_pyobject
 
                 if not is_pyobject_tainted(value) or self._override_pyobject_tainted:
                     try:
@@ -386,7 +386,7 @@ def supported_dbapi_integration(integration_name):
 
 def check_tainted_args(args, kwargs, tracer, integration_name, method):
     if supported_dbapi_integration(integration_name) and method.__name__ == "execute":
-        from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted
+        from ._taint_tracking import is_pyobject_tainted
 
         return len(args) and args[0] and is_pyobject_tainted(args[0])
 
diff --git a/ddtrace/appsec/iast/processor.py b/ddtrace/appsec/iast/processor.py
@@ -4,17 +4,18 @@
 
 from ddtrace.appsec._constants import APPSEC
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec.iast import oce
-from ddtrace.appsec.iast._metrics import _set_metric_iast_request_tainted
-from ddtrace.appsec.iast._utils import _iast_report_to_str
-from ddtrace.appsec.iast._utils import _is_iast_enabled
-from ddtrace.appsec.trace_utils import _asm_manual_keep
 from ddtrace.constants import ORIGIN_KEY
 from ddtrace.ext import SpanTypes
 from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.processor import SpanProcessor
 
+from . import oce
+from ..trace_utils import _asm_manual_keep
+from ._metrics import _set_metric_iast_request_tainted
+from ._utils import _iast_report_to_str
+from ._utils import _is_iast_enabled
+
 
 if TYPE_CHECKING:  # pragma: no cover
     from ddtrace.span import Span
@@ -29,7 +30,7 @@ def on_span_start(self, span):
         if span.span_type != SpanTypes.WEB:
             return
         oce.acquire_request(span)
-        from ddtrace.appsec.iast._taint_tracking import create_context
+        from ._taint_tracking import create_context
 
         create_context()
 
@@ -49,7 +50,7 @@ def on_span_finish(self, span):
             span.set_metric(IAST.ENABLED, 0.0)
             return
 
-        from ddtrace.appsec.iast._taint_tracking import reset_context  # noqa: F401
+        from ._taint_tracking import reset_context  # noqa: F401
 
         span.set_metric(IAST.ENABLED, 1.0)
 
diff --git a/ddtrace/appsec/iast/taint_sinks/__init__.py b/ddtrace/appsec/iast/taint_sinks/__init__.py
@@ -1,4 +1,4 @@
-from ddtrace.appsec.iast.taint_sinks.path_traversal import open_path_traversal
+from .path_traversal import open_path_traversal
 
 
 __all__ = [
diff --git a/ddtrace/appsec/iast/taint_sinks/_base.py b/ddtrace/appsec/iast/taint_sinks/_base.py
@@ -4,30 +4,31 @@
 
 from ddtrace import tracer
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec.iast import oce
-from ddtrace.appsec.iast._metrics import _set_metric_iast_executed_sink
-from ddtrace.appsec.iast._overhead_control_engine import Operation
-from ddtrace.appsec.iast._utils import _has_to_scrub
-from ddtrace.appsec.iast._utils import _is_evidence_value_parts
-from ddtrace.appsec.iast._utils import _scrub
-from ddtrace.appsec.iast.reporter import Evidence
-from ddtrace.appsec.iast.reporter import IastSpanReporter
-from ddtrace.appsec.iast.reporter import Location
-from ddtrace.appsec.iast.reporter import Source
-from ddtrace.appsec.iast.reporter import Vulnerability
 from ddtrace.internal import core
 from ddtrace.internal.compat import six
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils.cache import LFUCache
 from ddtrace.settings import _config
 
+from .. import oce
+from .._metrics import _set_metric_iast_executed_sink
+from .._overhead_control_engine import Operation
+from .._utils import _has_to_scrub
+from .._utils import _is_evidence_value_parts
+from .._utils import _scrub
+from ..reporter import Evidence
+from ..reporter import IastSpanReporter
+from ..reporter import Location
+from ..reporter import Source
+from ..reporter import Vulnerability
+
 
 try:
     # Python >= 3.4
-    from ddtrace.appsec.iast._stacktrace import get_info_frame
+    from .._stacktrace import get_info_frame
 except ImportError:
     # Python 2
-    from ddtrace.appsec.iast._stacktrace_py2 import get_info_frame
+    from .._stacktrace_py2 import get_info_frame
 
 if TYPE_CHECKING:  # pragma: no cover
     from typing import Any
diff --git a/ddtrace/appsec/iast/taint_sinks/command_injection.py b/ddtrace/appsec/iast/taint_sinks/command_injection.py
diff --git a/ddtrace/appsec/iast/taint_sinks/insecure_cookie.py b/ddtrace/appsec/iast/taint_sinks/insecure_cookie.py
diff --git a/ddtrace/appsec/iast/taint_sinks/path_traversal.py b/ddtrace/appsec/iast/taint_sinks/path_traversal.py
diff --git a/ddtrace/appsec/iast/taint_sinks/sql_injection.py b/ddtrace/appsec/iast/taint_sinks/sql_injection.py
diff --git a/ddtrace/appsec/iast/taint_sinks/weak_cipher.py b/ddtrace/appsec/iast/taint_sinks/weak_cipher.py
diff --git a/ddtrace/appsec/iast/taint_sinks/weak_hash.py b/ddtrace/appsec/iast/taint_sinks/weak_hash.py
