chore(iast): add more iast tests types (#13132)

## 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/_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.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
+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,
-        evidence,
-        file_name,
-        line_number,
-        function_name: Text = "",
-        class_name: Text = "",
+        vulnerability_type: str,
+        evidence: Evidence,
+        file_name: Optional[str],
+        line_number: int,
+        function_name: str = "",
+        class_name: str = "",
         *args,
         **kwargs,
     ) -> bool:
@@ -125,7 +128,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)
@@ -145,17 +148,19 @@ 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, IAST.TEXT_TYPES):
+            if isinstance(evidence_value, (bytes, bytearray)):
+                evidence_value = evidence_value.decode("utf-8")
             evidence = Evidence(value=evidence_value, dialect=dialect)
         else:
             log.debug("Unexpected evidence_value type: %s", type(evidence_value))
@@ -165,7 +170,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
@@ -189,7 +194,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

@@ -1,17 +1,16 @@
-# -*- encoding: utf-8 -*-
 import re
-from typing import Any  # noqa:F401
-from typing import List  # noqa:F401
-from typing import NamedTuple  # noqa:F401
-from typing import Optional  # noqa:F401
-from typing import Text  # noqa:F401
+from typing import Any
+from typing import List
+from typing import NamedTuple
+from typing import Optional
 
 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
 
 
@@ -26,18 +25,16 @@
 TAINT_FORMAT_PATTERN_BYTES = re.compile(TAINT_FORMAT_CAPTURE_BYTES, re.MULTILINE | re.DOTALL)
 
 
-def create_taint_range_with_format(text_input, fn_origin=""):  # type: (Any, str) -> Any
+def create_taint_range_with_format(text_input: Any, fn_origin: str = "") -> Any:
     is_bytes = isinstance(text_input, (bytes, bytearray))
     taint_format_capture = TAINT_FORMAT_CAPTURE_BYTES if is_bytes else TAINT_FORMAT_CAPTURE
     taint_format_pattern = TAINT_FORMAT_PATTERN_BYTES if is_bytes else TAINT_FORMAT_PATTERN
 
-    text_output = re.sub(  # type: ignore
-        taint_format_capture, r"\2", text_input, flags=re.MULTILINE | re.DOTALL
-    )  # type: Any
+    text_output: Any = re.sub(taint_format_capture, r"\2", text_input, flags=re.MULTILINE | re.DOTALL)
     if isinstance(text_input, bytearray):
-        text_output = bytearray(text_output)  # type: ignore
+        text_output = bytearray(text_output)
 
-    ranges_ = []  # type: List[TaintRange]
+    ranges_: List[TaintRange] = []
     acc_input_id = 0
     for i, match in enumerate(taint_format_pattern.finditer(text_input)):  # type: ignore[attr-defined]
         match_start = match.start() - (i * 6) - acc_input_id
@@ -63,24 +60,23 @@ def create_taint_range_with_format(text_input, fn_origin=""):  # type: (Any, str
 
     taint_pyobject_with_ranges(
         text_output,
-        ranges_,
+        tuple(ranges_),
     )
     return text_output
 
 
-class BaseReplacement(object):
-    def _to_tainted_string_with_origin(self, text):
-        # type: (Text) -> Text
+class BaseReplacement:
+    def _to_tainted_string_with_origin(self, text: TEXT_TYPE) -> TEXT_TYPE:
         if not isinstance(text, (str, bytes, bytearray)):
             return text
 
         # CAVEAT: the sequences ":+-" and "-+:" can be escaped with  "::++--" and "--+*::"
         elements = re.split(r"(\:\+-<[0-9a-zA-Z\-]+>|<[0-9a-zA-Z\-]+>-\+\:)", text)
 
-        ranges = []  # type: List[TaintRange]
+        ranges: List[TaintRange] = []
         ranges_append = ranges.append
         new_text = text.__class__()
-        context = None  # type: Optional[EscapeContext]
+        context: Optional[EscapeContext] = None
         for index, element in enumerate(elements):
             if index % 2 == 0:
                 element = element.replace("::++--", ":+-")
@@ -112,11 +108,11 @@ def _to_tainted_string_with_origin(self, text):
 
     def _assert_format_result(
         self,
-        taint_escaped_template,  # type: Text
-        taint_escaped_parameter,  # type: Any
-        expected_result,  # type: Text
-        escaped_expected_result,  # type: Text
-    ):  # type: (...) -> None
+        taint_escaped_template: TEXT_TYPE,
+        taint_escaped_parameter: Any,
+        expected_result: TEXT_TYPE,
+        escaped_expected_result: TEXT_TYPE,
+    ) -> None:
         template = self._to_tainted_string_with_origin(taint_escaped_template)
         parameter = self._to_tainted_string_with_origin(taint_escaped_parameter)
         result = mod.do_format_with_positional_parameter(template, parameter)

tests/appsec/iast/aspects/test_add_aspect.py

@@ -1,5 +1,10 @@
 import logging
+from pathlib import Path
+from pathlib import PosixPath
 
+from hypothesis import given
+from hypothesis.strategies import from_type
+from hypothesis.strategies import one_of
 import pytest
 
 from ddtrace.appsec._iast._taint_tracking import OriginType
@@ -13,6 +18,7 @@
 from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
 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
 
@@ -34,6 +40,30 @@ def test_add_aspect_successful(obj1, obj2):
     assert ddtrace_aspects.add_aspect(obj1, obj2) == obj1 + obj2
 
 
+@given(one_of(string_strategies))
+def test_add_aspect_text_successful(text):
+    assert ddtrace_aspects.add_aspect(text, text) == text + text
+
+
+@given(from_type(int))
+def test_add_aspect_int_successful(text):
+    assert ddtrace_aspects.add_aspect(text, text) == text + text
+
+
+@given(from_type(Path))
+def test_add_aspect_path_error(path):
+    with pytest.raises(TypeError) as exc_info:
+        ddtrace_aspects.add_aspect(path, path)
+    assert str(exc_info.value) == "unsupported operand type(s) for +: 'PosixPath' and 'PosixPath'"
+
+
+@given(from_type(PosixPath))
+def test_add_aspect_posixpath_error(posixpath):
+    with pytest.raises(TypeError) as exc_info:
+        ddtrace_aspects.add_aspect(posixpath, posixpath)
+    assert str(exc_info.value) == "unsupported operand type(s) for +: 'PosixPath' and 'PosixPath'"
+
+
 @pytest.mark.parametrize(
     "obj1, obj2",
     [(b"Hi", ""), ("Hi", b""), ({"a", "b"}, {"c", "d"}), (dict(), dict())],

tests/appsec/iast/aspects/test_common_aspects.py

@@ -4,10 +4,13 @@
 """
 import os
 
+from hypothesis import given
+from hypothesis.strategies import one_of
 import pytest
 
 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=""):
@@ -48,6 +51,16 @@ def callee_{function}_direct(*args, **kwargs):
 from tests.appsec.iast.fixtures.aspects import unpatched_callers  # type: ignore[attr-defined] # noqa: E402
 
 
+@given(one_of(string_strategies))
+def test_aspect_patched_result_hypothesis(text_input):
+    """
+    Test that the result of the patched aspect call is the same as the unpatched one
+    using Hypothesis-generated inputs.
+    """
+    for aspect in [x for x in dir(unpatched_callers) if not x.startswith(("_", "@"))]:
+        assert getattr(patched_callers, aspect)(text_input) == getattr(unpatched_callers, aspect)(text_input)
+
+
 @pytest.mark.parametrize("aspect", [x for x in dir(unpatched_callers) if not x.startswith(("_", "@"))])
 @pytest.mark.parametrize("args", [(), ("a"), ("a", "b")])
 @pytest.mark.parametrize("kwargs", [{}, {"dry_run": False}, {"dry_run": True}])