ci: fix type

Author: avara1986

ddtrace/appsec/_iast/_ast/visitor.py

@@ -227,7 +227,7 @@ def _merge_dicts(*args_functions: Set[str]) -> Set[str]:
 
     @staticmethod
     def _is_string_node(node: Any) -> bool:
-        if PY3 and (isinstance(node, ast.Constant) and isinstance(node.value, (str, bytes, bytearray))):
+        if PY3 and (isinstance(node, ast.Constant) and isinstance(node.value, IAST.TEXT_TYPES)):
             return True
 
         return False

ddtrace/appsec/_iast/_handlers.py

@@ -4,6 +4,7 @@
 from wrapt import when_imported
 from wrapt import wrap_function_wrapper as _w
 
+from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._iast_request_context import get_iast_stacktrace_reported
 from ddtrace.appsec._iast._iast_request_context import set_iast_stacktrace_reported
 from ddtrace.appsec._iast._logs import iast_instrumentation_wrapt_debug_log
@@ -249,7 +250,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
 
 def _custom_protobuf_getattribute(self, name):
     ret = type(self).__saved_getattr(self, name)
-    if isinstance(ret, (str, bytes, bytearray)):
+    if isinstance(ret, IAST.TEXT_TYPES):
         ret = taint_pyobject(
             pyobject=ret,
             source_name=OriginType.GRPC_BODY,

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -4,6 +4,7 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
+from ..._constants import IAST
 from .._patch import set_and_check_module_is_patched
 from .._patch import set_module_unpatched
 from .._patch import try_wrap_function_wrapper
@@ -55,7 +56,7 @@ def wrapped_loads(wrapped, instance, args, kwargs):
                     obj = taint_structure(obj, source.origin, source.origin)
                 elif isinstance(obj, list):
                     obj = taint_structure(obj, source.origin, source.origin)
-                elif isinstance(obj, (str, bytes, bytearray)):
+                elif isinstance(obj, IAST.TEXT_TYPES):
                     obj = taint_pyobject(obj, source.name, source.value, source.origin)
             except Exception:
                 log.debug("Unexpected exception while reporting vulnerability", exc_info=True)

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -435,7 +435,7 @@ def format_value_aspect(
     element: Any,
     options: int = 0,
     format_spec: Optional[str] = None,
-) -> Union[str, bytes, bytearray]:
+) -> TEXT_TYPES:
     if options == 115:
         new_text = str_aspect(str, 0, element)
     elif options == 114:

ddtrace/appsec/_iast/_taint_utils.py

@@ -5,6 +5,7 @@
 from typing import Optional
 from typing import Union
 
+from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.internal.logger import get_logger
@@ -101,7 +102,7 @@ def taint_structure(main_obj, source_key, source_value, override_pyobject_tainte
             if command.pre:  # first processing of the object
                 if not command.obj:
                     command.store(command.obj)
-                elif isinstance(command.obj, (str, bytes, bytearray)):
+                elif isinstance(command.obj, IAST.TEXT_TYPES):
                     if override_pyobject_tainted or not is_pyobject_tainted(command.obj):
                         new_obj = taint_pyobject(
                             pyobject=command.obj,
@@ -161,7 +162,7 @@ def __init__(self, original_list, origins=(0, 0), override_pyobject_tainted=Fals
 
     def _taint(self, value):
         if value:
-            if isinstance(value, (str, bytes, bytearray)):
+            if isinstance(value, IAST.TEXT_TYPES):
                 if not is_pyobject_tainted(value) or self._override_pyobject_tainted:
                     try:
                         # TODO: migrate this part to shift ranges instead of creating a new one
@@ -342,7 +343,7 @@ def _taint(self, value, key, origin=None):
         if origin is None:
             origin = self._origin_value
         if value:
-            if isinstance(value, (str, bytes, bytearray)):
+            if isinstance(value, IAST.TEXT_TYPES):
                 if not is_pyobject_tainted(value) or self._override_pyobject_tainted:
                     try:
                         # TODO: migrate this part to shift ranges instead of creating a new one

ddtrace/appsec/_iast/taint_sinks/_base.py

@@ -2,8 +2,8 @@
 from typing import Any
 from typing import Callable
 from typing import Optional
-from typing import Text
 from typing import Tuple
+from typing import Union
 
 from ddtrace.appsec._deduplications import deduplication
 from ddtrace.appsec._iast._taint_tracking import get_ranges
@@ -12,6 +12,7 @@
 from ddtrace.settings.asm import config as asm_config
 from ddtrace.trace import tracer
 
+from ..._constants import IAST
 from .._iast_request_context import get_iast_reporter
 from .._iast_request_context import set_iast_reporter
 from .._overhead_control_engine import Operation
@@ -27,6 +28,8 @@
 
 CWD = os.path.abspath(os.getcwd())
 
+TEXT_TYPES = Union[str, bytes, bytearray]
+
 
 class taint_sink_deduplication(deduplication):
     def _check_deduplication(self):
@@ -77,12 +80,12 @@ def wrapper(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
     @taint_sink_deduplication
     def _prepare_report(
         cls,
-        vulnerability_type: Text,
+        vulnerability_type: str,
         evidence: Evidence,
-        file_name: Optional[Text],
+        file_name: Optional[str],
         line_number: int,
-        function_name: Text = "",
-        class_name: Text = "",
+        function_name: str = "",
+        class_name: str = "",
         *args,
         **kwargs,
     ) -> bool:
@@ -126,7 +129,7 @@ def _prepare_report(
         return True
 
     @classmethod
-    def _compute_file_line(cls) -> Tuple[Optional[Text], Optional[int], Optional[Text], Optional[Text]]:
+    def _compute_file_line(cls) -> Tuple[Optional[str], Optional[int], Optional[str], Optional[str]]:
         file_name = line_number = function_name = class_name = None
 
         frame_info = get_info_frame(CWD)
@@ -146,18 +149,18 @@ def _compute_file_line(cls) -> Tuple[Optional[Text], Optional[int], Optional[Tex
     @classmethod
     def _create_evidence_and_report(
         cls,
-        vulnerability_type: Text,
-        evidence_value: Text = "",
-        dialect: Optional[Text] = None,
-        file_name: Optional[Text] = None,
+        vulnerability_type: str,
+        evidence_value: TEXT_TYPES = "",
+        dialect: Optional[str] = None,
+        file_name: Optional[str] = None,
         line_number: Optional[int] = None,
-        function_name: Optional[Text] = None,
-        class_name: Optional[Text] = None,
+        function_name: Optional[str] = None,
+        class_name: Optional[str] = None,
         *args,
         **kwargs,
     ):
-        if isinstance(evidence_value, (str, bytes, bytearray)):
-            if isinstance(evidence_value, bytearray):
+        if isinstance(evidence_value, IAST.TEXT_TYPES):
+            if isinstance(evidence_value, (bytes, bytearray)):
                 evidence_value = evidence_value.decode("utf-8")
             evidence = Evidence(value=evidence_value, dialect=dialect)
         else:
@@ -168,7 +171,7 @@ def _create_evidence_and_report(
         )
 
     @classmethod
-    def report(cls, evidence_value: Text = "", dialect: Optional[Text] = None) -> None:
+    def report(cls, evidence_value: TEXT_TYPES = "", dialect: Optional[str] = None) -> None:
         """Build a IastSpanReporter instance to report it in the `AppSecIastSpanProcessor` as a string JSON"""
         if cls.acquire_quota():
             file_name = line_number = function_name = class_name = None
@@ -192,7 +195,7 @@ def report(cls, evidence_value: Text = "", dialect: Optional[Text] = None) -> No
                 cls.increment_quota()
 
     @classmethod
-    def is_tainted_pyobject(cls, string_to_check: Text) -> bool:
+    def is_tainted_pyobject(cls, string_to_check: TEXT_TYPES) -> bool:
         """Check if a string contains tainted ranges that are not marked as secure.
 
         A string is considered tainted when:

tests/appsec/iast/aspects/aspect_utils.py

@@ -3,46 +3,17 @@
 from typing import List
 from typing import NamedTuple
 from typing import Optional
-from typing import Union
-
-from hypothesis.strategies import binary
-from hypothesis.strategies import builds
-from hypothesis.strategies import text
 
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import Source
 from ddtrace.appsec._iast._taint_tracking import TaintRange
 from ddtrace.appsec._iast._taint_tracking import as_formatted_evidence
 from ddtrace.appsec._iast._taint_tracking import set_ranges
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject_with_ranges
+from tests.appsec.iast.iast_utils import TEXT_TYPE
 from tests.appsec.iast.iast_utils import _iast_patched_module
 
 
-TEXT_TYPE = Union[str, bytes, bytearray]
-
-
-class CustomStr(str):
-    pass
-
-
-class CustomBytes(bytes):
-    pass
-
-
-class CustomBytearray(bytearray):
-    pass
-
-
-string_strategies: List[Any] = [
-    text(),  # regular str
-    binary(),  # regular bytes
-    builds(bytearray, binary()),  # regular bytearray
-    builds(CustomStr, text()),  # custom str subclass
-    builds(CustomBytes, binary()),  # custom bytes subclass
-    builds(CustomBytearray, binary()),  # custom bytearray subclass
-]
-
-
 mod = _iast_patched_module("benchmarks.bm.iast_fixtures.str_methods")
 
 EscapeContext = NamedTuple("EscapeContext", [("id", Any), ("position", int)])

tests/appsec/iast/aspects/test_add_aspect.py

@@ -16,9 +16,9 @@
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 import ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects
 from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
-from tests.appsec.iast.aspects.aspect_utils import string_strategies
 from tests.appsec.iast.conftest import _end_iast_context_and_oce
 from tests.appsec.iast.conftest import _start_iast_context_and_oce
+from tests.appsec.iast.iast_utils import string_strategies
 from tests.utils import override_env
 from tests.utils import override_global_config
 

tests/appsec/iast/aspects/test_common_aspects.py

@@ -8,9 +8,9 @@
 from hypothesis.strategies import one_of
 import pytest
 
-from tests.appsec.iast.aspects.aspect_utils import string_strategies
 import tests.appsec.iast.fixtures.aspects.callees
 from tests.appsec.iast.iast_utils import _iast_patched_module
+from tests.appsec.iast.iast_utils import string_strategies
 
 
 def generate_callers_from_callees(callees_module, callers_file="", callees_module_str=""):

tests/appsec/iast/iast_utils.py

@@ -1,10 +1,17 @@
 import importlib
 import re
 import types
+from typing import Any
+from typing import List
 from typing import Optional
 from typing import Text
+from typing import Union
 import zlib
 
+from hypothesis.strategies import binary
+from hypothesis.strategies import builds
+from hypothesis.strategies import text
+
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._ast.ast_patching import astpatch_module
 from ddtrace.appsec._iast._ast.ast_patching import iastpatch
@@ -60,3 +67,30 @@ def _iast_patched_module(module_name, new_module_object=False):
     else:
         raise IastTestException(f"IAST Test Error: module {module_name} was excluded: {res}")
     return module
+
+
+TEXT_TYPE = Union[str, bytes, bytearray]
+
+
+class CustomStr(str):
+    pass
+
+
+class CustomBytes(bytes):
+    pass
+
+
+class CustomBytearray(bytearray):
+    pass
+
+
+non_empty_text = text().filter(lambda x: x not in ("", "0", "1", "2"))
+
+string_strategies: List[Any] = [
+    text(),  # regular str
+    binary(),  # regular bytes
+    builds(bytearray, binary()),  # regular bytearray
+    builds(CustomStr, text()),  # custom str subclass
+    builds(CustomBytes, binary()),  # custom bytes subclass
+    builds(CustomBytearray, binary()),  # custom bytearray subclass
+]