diff --git a/ddtrace/appsec/_constants.py b/ddtrace/appsec/_constants.py index 4b7aafa3e10..8a7f6f079a6 100644 --- a/ddtrace/appsec/_constants.py +++ b/ddtrace/appsec/_constants.py @@ -134,8 +134,10 @@ class IAST(metaclass=Constant_Class): JSON: Literal["_dd.iast.json"] = "_dd.iast.json" ENABLED: Literal["_dd.iast.enabled"] = "_dd.iast.enabled" PATCH_MODULES: Literal["_DD_IAST_PATCH_MODULES"] = "_DD_IAST_PATCH_MODULES" + ENV_NO_DIR_PATCH: Literal["_DD_IAST_NO_DIR_PATCH"] = "_DD_IAST_NO_DIR_PATCH" DENY_MODULES: Literal["_DD_IAST_DENY_MODULES"] = "_DD_IAST_DENY_MODULES" SEP_MODULES: Literal[","] = "," + PATCH_ADDED_SYMBOL_PREFIX: Literal["_ddtrace_"] = "_ddtrace_" METRICS_REPORT_LVLS = ( (TELEMETRY_DEBUG_VERBOSITY, TELEMETRY_DEBUG_NAME), diff --git a/ddtrace/appsec/_iast/_ast/ast_patching.py b/ddtrace/appsec/_iast/_ast/ast_patching.py index 4ae0c31233e..eca157eb3a5 100644 --- a/ddtrace/appsec/_iast/_ast/ast_patching.py +++ b/ddtrace/appsec/_iast/_ast/ast_patching.py @@ -5,6 +5,8 @@ import os import re from sys import builtin_module_names +from sys import version_info +import textwrap from types import ModuleType from typing import Optional from typing import Text @@ -14,12 +16,14 @@ from ddtrace.appsec._python_info.stdlib import _stdlib_for_python_version from ddtrace.internal.logger import get_logger from ddtrace.internal.module import origin +from ddtrace.internal.utils.formats import asbool from .visitor import AstVisitor _VISITOR = AstVisitor() +_PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX # Prefixes for modules where IAST patching is allowed IAST_ALLOWLIST: Tuple[Text, ...] = ("tests.appsec.iast.",) @@ -278,6 +282,7 @@ "protobuf.", "pycparser.", # this package is called when a module is imported, propagation is not needed "pytest.", # Testing framework + "_pytest.", "setuptools.", "sklearn.", # Machine learning library "sqlalchemy.orm.interfaces.", # Performance optimization @@ -368,7 +373,7 @@ def visit_ast( source_text: Text, module_path: Text, module_name: Text = "", -) -> Optional[str]: +) -> Optional[ast.Module]: parsed_ast = ast.parse(source_text, module_path) _VISITOR.update_location(filename=module_path, module_name=module_name) modified_ast = _VISITOR.visit(parsed_ast) @@ -401,23 +406,56 @@ def _remove_flask_run(text: Text) -> Text: return new_text -def astpatch_module(module: ModuleType, remove_flask_run: bool = False) -> Tuple[str, str]: +_DIR_WRAPPER = textwrap.dedent( + f""" + + +def {_PREFIX}dir(): + orig_dir = globals().get("{_PREFIX}orig_dir__") + + if orig_dir: + # Use the original __dir__ method and filter the results + results = [name for name in orig_dir() if not name.startswith("{_PREFIX}")] + else: + # List names from the module's __dict__ and filter out the unwanted names + results = [ + name for name in globals() + if not (name.startswith("{_PREFIX}") or name == "__dir__") + ] + + return results + +def {_PREFIX}set_dir_filter(): + if "__dir__" in globals(): + # Store the original __dir__ method + globals()["{_PREFIX}orig_dir__"] = __dir__ + + # Replace the module's __dir__ with the custom one + globals()["__dir__"] = {_PREFIX}dir + +{_PREFIX}set_dir_filter() + + """ +) + + +def astpatch_module(module: ModuleType, remove_flask_run: bool = False) -> Tuple[str, Optional[ast.Module]]: module_name = module.__name__ module_origin = origin(module) if module_origin is None: log.debug("astpatch_source couldn't find the module: %s", module_name) - return "", "" + return "", None module_path = str(module_origin) try: if module_origin.stat().st_size == 0: # Don't patch empty files like __init__.py log.debug("empty file: %s", module_path) - return "", "" + return "", None except OSError: log.debug("astpatch_source couldn't find the file: %s", module_path, exc_info=True) - return "", "" + return "", None # Get the file extension, if it's dll, os, pyd, dyn, dynlib: return # If its pyc or pyo, change to .py and check that the file exists. If not, @@ -427,30 +465,35 @@ def astpatch_module(module: ModuleType, remove_flask_run: bool = False) -> Tuple if module_ext.lower() not in {".pyo", ".pyc", ".pyw", ".py"}: # Probably native or built-in module log.debug("extension not supported: %s for: %s", module_ext, module_path) - return "", "" + return "", None with open(module_path, "r", encoding=get_encoding(module_path)) as source_file: try: source_text = source_file.read() except UnicodeDecodeError: log.debug("unicode decode error for file: %s", module_path, exc_info=True) - return "", "" + return "", None if len(source_text.strip()) == 0: # Don't patch empty files like __init__.py log.debug("empty file: %s", module_path) - return "", "" + return "", None if remove_flask_run: source_text = _remove_flask_run(source_text) - new_source = visit_ast( + if not asbool(os.environ.get(IAST.ENV_NO_DIR_PATCH, "false")) and version_info > (3, 7): + # Add the dir filter so __ddtrace stuff is not returned by dir(module) + # does not work in 3.7 because it enters into infinite recursion + source_text += _DIR_WRAPPER + + new_ast = visit_ast( source_text, module_path, module_name=module_name, ) - if new_source is None: + if new_ast is None: log.debug("file not ast patched: %s", module_path) - return "", "" + return "", None - return module_path, new_source + return module_path, new_ast diff --git a/ddtrace/appsec/_iast/_ast/visitor.py b/ddtrace/appsec/_iast/_ast/visitor.py index decf62d1384..6283d6cf87b 100644 --- a/ddtrace/appsec/_iast/_ast/visitor.py +++ b/ddtrace/appsec/_iast/_ast/visitor.py @@ -12,6 +12,7 @@ from typing import Text from typing import Tuple # noqa:F401 +from ..._constants import IAST from .._metrics import _set_metric_iast_instrumented_propagation from ..constants import DEFAULT_PATH_TRAVERSAL_FUNCTIONS from ..constants import DEFAULT_WEAK_RANDOMNESS_FUNCTIONS @@ -22,11 +23,12 @@ PY38_PLUS = sys.version_info >= (3, 8, 0) PY39_PLUS = sys.version_info >= (3, 9, 0) +_PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX CODE_TYPE_FIRST_PARTY = "first_party" CODE_TYPE_DD = "datadog" CODE_TYPE_SITE_PACKAGES = "site_packages" CODE_TYPE_STDLIB = "stdlib" -TAINT_SINK_FUNCTION_REPLACEMENT = "ddtrace_taint_sinks.ast_function" +TAINT_SINK_FUNCTION_REPLACEMENT = _PREFIX + "taint_sinks.ast_function" def _mark_avoid_convert_recursively(node): @@ -38,71 +40,71 @@ def _mark_avoid_convert_recursively(node): _ASPECTS_SPEC: Dict[Text, Any] = { "definitions_module": "ddtrace.appsec._iast._taint_tracking.aspects", - "alias_module": "ddtrace_aspects", + "alias_module": _PREFIX + "aspects", "functions": { - "StringIO": "ddtrace_aspects.stringio_aspect", - "BytesIO": "ddtrace_aspects.bytesio_aspect", - "str": "ddtrace_aspects.str_aspect", - "bytes": "ddtrace_aspects.bytes_aspect", - "bytearray": "ddtrace_aspects.bytearray_aspect", - "ddtrace_iast_flask_patch": "ddtrace_aspects.empty_func", # To avoid recursion + "StringIO": _PREFIX + "aspects.stringio_aspect", + "BytesIO": _PREFIX + "aspects.bytesio_aspect", + "str": _PREFIX + "aspects.str_aspect", + "bytes": _PREFIX + "aspects.bytes_aspect", + "bytearray": _PREFIX + "aspects.bytearray_aspect", + "ddtrace_iast_flask_patch": _PREFIX + "aspects.empty_func", # To avoid recursion }, "stringalike_methods": { - "StringIO": "ddtrace_aspects.stringio_aspect", - "BytesIO": "ddtrace_aspects.bytesio_aspect", - "decode": "ddtrace_aspects.decode_aspect", - "join": "ddtrace_aspects.join_aspect", - "encode": "ddtrace_aspects.encode_aspect", - "extend": "ddtrace_aspects.bytearray_extend_aspect", - "upper": "ddtrace_aspects.upper_aspect", - "lower": "ddtrace_aspects.lower_aspect", - "replace": "ddtrace_aspects.replace_aspect", - "swapcase": "ddtrace_aspects.swapcase_aspect", - "title": "ddtrace_aspects.title_aspect", - "capitalize": "ddtrace_aspects.capitalize_aspect", - "casefold": "ddtrace_aspects.casefold_aspect", - "translate": "ddtrace_aspects.translate_aspect", - "format": "ddtrace_aspects.format_aspect", - "format_map": "ddtrace_aspects.format_map_aspect", - "zfill": "ddtrace_aspects.zfill_aspect", - "ljust": "ddtrace_aspects.ljust_aspect", - "split": "ddtrace_aspects.split_aspect", # Both regular split and re.split - "rsplit": "ddtrace_aspects.rsplit_aspect", - "splitlines": "ddtrace_aspects.splitlines_aspect", + "StringIO": _PREFIX + "aspects.stringio_aspect", + "BytesIO": _PREFIX + "aspects.bytesio_aspect", + "decode": _PREFIX + "aspects.decode_aspect", + "join": _PREFIX + "aspects.join_aspect", + "encode": _PREFIX + "aspects.encode_aspect", + "extend": _PREFIX + "aspects.bytearray_extend_aspect", + "upper": _PREFIX + "aspects.upper_aspect", + "lower": _PREFIX + "aspects.lower_aspect", + "replace": _PREFIX + "aspects.replace_aspect", + "swapcase": _PREFIX + "aspects.swapcase_aspect", + "title": _PREFIX + "aspects.title_aspect", + "capitalize": _PREFIX + "aspects.capitalize_aspect", + "casefold": _PREFIX + "aspects.casefold_aspect", + "translate": _PREFIX + "aspects.translate_aspect", + "format": _PREFIX + "aspects.format_aspect", + "format_map": _PREFIX + "aspects.format_map_aspect", + "zfill": _PREFIX + "aspects.zfill_aspect", + "ljust": _PREFIX + "aspects.ljust_aspect", + "split": _PREFIX + "aspects.split_aspect", # Both regular split and re.split + "rsplit": _PREFIX + "aspects.rsplit_aspect", + "splitlines": _PREFIX + "aspects.splitlines_aspect", # re module and re.Match methods - "findall": "ddtrace_aspects.re_findall_aspect", - "finditer": "ddtrace_aspects.re_finditer_aspect", - "fullmatch": "ddtrace_aspects.re_fullmatch_aspect", - "expand": "ddtrace_aspects.re_expand_aspect", - "group": "ddtrace_aspects.re_group_aspect", - "groups": "ddtrace_aspects.re_groups_aspect", - "match": "ddtrace_aspects.re_match_aspect", - "search": "ddtrace_aspects.re_search_aspect", - "sub": "ddtrace_aspects.re_sub_aspect", - "subn": "ddtrace_aspects.re_subn_aspect", + "findall": _PREFIX + "aspects.re_findall_aspect", + "finditer": _PREFIX + "aspects.re_finditer_aspect", + "fullmatch": _PREFIX + "aspects.re_fullmatch_aspect", + "expand": _PREFIX + "aspects.re_expand_aspect", + "group": _PREFIX + "aspects.re_group_aspect", + "groups": _PREFIX + "aspects.re_groups_aspect", + "match": _PREFIX + "aspects.re_match_aspect", + "search": _PREFIX + "aspects.re_search_aspect", + "sub": _PREFIX + "aspects.re_sub_aspect", + "subn": _PREFIX + "aspects.re_subn_aspect", }, # Replacement function for indexes and ranges "slices": { - "index": "ddtrace_aspects.index_aspect", - "slice": "ddtrace_aspects.slice_aspect", + "index": _PREFIX + "aspects.index_aspect", + "slice": _PREFIX + "aspects.slice_aspect", }, # Replacement functions for modules "module_functions": { "os.path": { - "basename": "ddtrace_aspects.ospathbasename_aspect", - "dirname": "ddtrace_aspects.ospathdirname_aspect", - "join": "ddtrace_aspects.ospathjoin_aspect", - "normcase": "ddtrace_aspects.ospathnormcase_aspect", - "split": "ddtrace_aspects.ospathsplit_aspect", - "splitext": "ddtrace_aspects.ospathsplitext_aspect", + "basename": _PREFIX + "aspects.ospathbasename_aspect", + "dirname": _PREFIX + "aspects.ospathdirname_aspect", + "join": _PREFIX + "aspects.ospathjoin_aspect", + "normcase": _PREFIX + "aspects.ospathnormcase_aspect", + "split": _PREFIX + "aspects.ospathsplit_aspect", + "splitext": _PREFIX + "aspects.ospathsplitext_aspect", } }, "operators": { - ast.Add: "ddtrace_aspects.add_aspect", - "INPLACE_ADD": "ddtrace_aspects.add_inplace_aspect", - "FORMAT_VALUE": "ddtrace_aspects.format_value_aspect", - ast.Mod: "ddtrace_aspects.modulo_aspect", - "BUILD_STRING": "ddtrace_aspects.build_string_aspect", + ast.Add: _PREFIX + "aspects.add_aspect", + "INPLACE_ADD": _PREFIX + "aspects.add_inplace_aspect", + "FORMAT_VALUE": _PREFIX + "aspects.format_value_aspect", + ast.Mod: _PREFIX + "aspects.modulo_aspect", + "BUILD_STRING": _PREFIX + "aspects.build_string_aspect", }, "excluded_from_patching": { # Key: module being patched @@ -123,6 +125,8 @@ def _mark_avoid_convert_recursively(node): }, "django.utils.html": {"": ("format_html", "format_html_join")}, "sqlalchemy.sql.compiler": {"": ("_requires_quotes",)}, + # Our added functions + "": {"": (f"{_PREFIX}dir", f"{_PREFIX}set_dir_filter")}, }, # This is a set since all functions will be replaced by taint_sink_functions "taint_sinks": { @@ -149,10 +153,10 @@ def _mark_avoid_convert_recursively(node): if sys.version_info >= (3, 12): - _ASPECTS_SPEC["module_functions"]["os.path"]["splitroot"] = "ddtrace_aspects.ospathsplitroot_aspect" + _ASPECTS_SPEC["module_functions"]["os.path"]["splitroot"] = _PREFIX + "aspects.ospathsplitroot_aspect" if sys.version_info >= (3, 12) or os.name == "nt": - _ASPECTS_SPEC["module_functions"]["os.path"]["splitdrive"] = "ddtrace_aspects.ospathsplitdrive_aspect" + _ASPECTS_SPEC["module_functions"]["os.path"]["splitdrive"] = _PREFIX + "aspects.ospathsplitdrive_aspect" class AstVisitor(ast.NodeTransformer): @@ -163,7 +167,7 @@ def __init__( ): self._sinkpoints_spec = { "definitions_module": "ddtrace.appsec._iast.taint_sinks", - "alias_module": "ddtrace_taint_sinks", + "alias_module": _PREFIX + "taint_sinks", "functions": {}, } self._sinkpoints_functions = self._sinkpoints_spec["functions"] @@ -458,6 +462,9 @@ def visit_FunctionDef(self, def_node: ast.FunctionDef) -> Any: Special case for some tests which would enter in a patching loop otherwise when visiting the check functions """ + if f"{_PREFIX}dir" in def_node.name or f"{_PREFIX}set_dir_filter" in def_node.name: + return def_node + self.replacements_disabled_for_functiondef = def_node.name in self.dont_patch_these_functionsdefs if hasattr(def_node.args, "vararg") and def_node.args.vararg: diff --git a/ddtrace/appsec/_iast/_loader.py b/ddtrace/appsec/_iast/_loader.py index e81024932f6..382060e3b7f 100644 --- a/ddtrace/appsec/_iast/_loader.py +++ b/ddtrace/appsec/_iast/_loader.py @@ -12,19 +12,19 @@ def _exec_iast_patched_module(module_watchdog, module): - patched_source = None + patched_ast = None compiled_code = None if IS_IAST_ENABLED: try: - module_path, patched_source = astpatch_module(module) + module_path, patched_ast = astpatch_module(module) except Exception: log.debug("Unexpected exception while AST patching", exc_info=True) - patched_source = None + patched_ast = None - if patched_source: + if patched_ast: try: # Patched source is compiled in order to execute it - compiled_code = compile(patched_source, module_path, "exec") + compiled_code = compile(patched_ast, module_path, "exec") except Exception: log.debug("Unexpected exception while compiling patched code", exc_info=True) compiled_code = None diff --git a/releasenotes/notes/code-security-patch-dir-54cc85f18e31f45c.yaml b/releasenotes/notes/code-security-patch-dir-54cc85f18e31f45c.yaml new file mode 100644 index 00000000000..a46aadc94b4 --- /dev/null +++ b/releasenotes/notes/code-security-patch-dir-54cc85f18e31f45c.yaml @@ -0,0 +1,4 @@ +--- +fixes: + - | + Code Security: patch the module dir function so original pre-patch results are not changed. diff --git a/tests/appsec/appsec_utils.py b/tests/appsec/appsec_utils.py index b7fd5d4743a..cbcc913a7af 100644 --- a/tests/appsec/appsec_utils.py +++ b/tests/appsec/appsec_utils.py @@ -100,6 +100,7 @@ def appsec_application_server( env[IAST.ENV] = iast_enabled env[IAST.ENV_REQUEST_SAMPLING] = "100" env["_DD_APPSEC_DEDUPLICATION_ENABLED"] = "false" + env[IAST.ENV_NO_DIR_PATCH] = "false" if assert_debug: env["_" + IAST.ENV_DEBUG] = iast_enabled env["_" + IAST.ENV_PROPAGATION_DEBUG] = iast_enabled diff --git a/tests/appsec/iast/_ast/test_ast_patching.py b/tests/appsec/iast/_ast/test_ast_patching.py index fdf9699231e..cf0fabd14e4 100644 --- a/tests/appsec/iast/_ast/test_ast_patching.py +++ b/tests/appsec/iast/_ast/test_ast_patching.py @@ -1,14 +1,21 @@ #!/usr/bin/env python3 import logging +import sys import astunparse import mock import pytest +from ddtrace.appsec._constants import IAST from ddtrace.appsec._iast._ast.ast_patching import _in_python_stdlib from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch from ddtrace.appsec._iast._ast.ast_patching import astpatch_module from ddtrace.appsec._iast._ast.ast_patching import visit_ast +from ddtrace.internal.utils.formats import asbool +from tests.utils import override_env + + +_PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX @pytest.mark.parametrize( @@ -55,12 +62,12 @@ def test_visit_ast_changed(source_text, module_path, module_name): ], ) def test_astpatch_module_changed(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects" + f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks" + f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects" ) assert "ddtrace_aspects.str_aspect(" in new_code @@ -72,12 +79,12 @@ def test_astpatch_module_changed(module_name): ], ) def test_astpatch_module_changed_add_operator(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects" + f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks" + f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects" ) assert "ddtrace_aspects.add_aspect(" in new_code @@ -89,12 +96,12 @@ def test_astpatch_module_changed_add_operator(module_name): ], ) def test_astpatch_module_changed_add_inplace_operator(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects" + f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks" + f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects" ) assert "ddtrace_aspects.add_inplace_aspect(" in new_code @@ -107,18 +114,18 @@ def test_astpatch_module_changed_add_inplace_operator(module_name): ], ) def test_astpatch_source_changed_with_future_imports(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - """ + f""" '\\nSome\\nmulti-line\\ndocstring\\nhere\\n' from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals -import ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks -import ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects +import ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks +import ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects import html""" ) assert "ddtrace_aspects.str_aspect(" in new_code @@ -136,7 +143,7 @@ def test_astpatch_source_changed_with_future_imports(module_name): ], ) def test_astpatch_source_unchanged(module_name): - assert ("", "") == astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) == astpatch_module(__import__(module_name, fromlist=[None])) def test_module_should_iast_patch(): @@ -170,7 +177,9 @@ def test_module_in_python_stdlib(module_name, result): def test_module_path_none(caplog): with caplog.at_level(logging.DEBUG), mock.patch("ddtrace.internal.module.Path.resolve", side_effect=AttributeError): - assert ("", "") == astpatch_module(__import__("tests.appsec.iast.fixtures.ast.str.class_str", fromlist=[None])) + assert ("", None) == astpatch_module( + __import__("tests.appsec.iast.fixtures.ast.str.class_str", fromlist=[None]) + ) assert "astpatch_source couldn't find the module: tests.appsec.iast.fixtures.ast.str.class_str" in caplog.text @@ -182,12 +191,12 @@ def test_module_path_none(caplog): ], ) def test_astpatch_stringio_module_changed(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects" + f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks" + f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects" ) assert "ddtrace_aspects.stringio_aspect(" in new_code @@ -200,12 +209,12 @@ def test_astpatch_stringio_module_changed(module_name): ], ) def test_astpatch_bytesio_module_changed(module_name): - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") != (module_path, new_source) - new_code = astunparse.unparse(new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) assert new_code.startswith( - "\nimport ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects" + f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks" + f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects" ) assert "ddtrace_aspects.bytesio_aspect(" in new_code @@ -221,5 +230,81 @@ def test_astpatch_globals_module_unchanged(module_name): This is a regression test for partially matching function names: ``globals()`` was being incorrectly patched with the aspect for ``glob()`` """ - module_path, new_source = astpatch_module(__import__(module_name, fromlist=[None])) - assert ("", "") == (module_path, new_source) + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) == (module_path, new_ast) + + +@pytest.mark.parametrize( + "module_name, env_var", + [ + ("tests.appsec.iast.fixtures.ast.other.with_implemented_dir", "false"), + ("tests.appsec.iast.fixtures.ast.other.with_implemented_dir", "true"), + ("tests.appsec.iast.fixtures.ast.other.without_implemented_dir", "false"), + ("tests.appsec.iast.fixtures.ast.other.without_implemented_dir", "true"), + ], +) +def test_astpatch_dir_patched_with_env_var(module_name, env_var): + """ + Check that the ast_patching._DIR_WRAPPER code is added to the end of the module if + the env var is False and not added otherwise + """ + with override_env({IAST.ENV_NO_DIR_PATCH: env_var}): + module_path, new_ast = astpatch_module(__import__(module_name, fromlist=[None])) + assert ("", None) != (module_path, new_ast) + new_code = astunparse.unparse(new_ast) + + if asbool(env_var): + # If ENV_NO_DIR_PATCH is set to True our added symbols are not filtered out + assert f"{_PREFIX}aspects" in new_code + assert f"{_PREFIX}taint_sinks" in new_code + else: + # Check that the added dir code is there + assert f"def {_PREFIX}dir" in new_code + assert f"def {_PREFIX}set_dir_filter()" in new_code + + +@pytest.mark.parametrize( + "module_name, expected_names", + [ + ( + "tests.appsec.iast.fixtures.ast.other.with_implemented_dir", + {"custom_added", "symbol1", "symbol2", "symbol3", "symbol4"}, + ), + ( + "tests.appsec.iast.fixtures.ast.other.without_implemented_dir", + {"symbol1", "symbol2", "symbol3", "symbol4"}, + ), + ], +) +@pytest.mark.skipif(sys.version_info < (3, 8), reason="the dir wrappers enters and infinite loop in 3.7") +def test_astpatch_dir_patched_with_or_without_custom_dir(module_name, expected_names): + """ + Check that the patched dir doesn't have any __ddtrace symbols and match the original + unpatched dir() output, both with or without a previous custom __dir__ implementation + """ + with override_env({IAST.ENV_NO_DIR_PATCH: "false"}): + imported_mod = __import__(module_name, fromlist=[None]) + orig_dir = set(dir(imported_mod)) + + for name in expected_names: + assert name in orig_dir + + module_path, new_ast = astpatch_module(imported_mod) + assert ("", None) != (module_path, new_ast) + + new_code = astunparse.unparse(new_ast) + assert f"def {_PREFIX}dir" in new_code + assert f"def {_PREFIX}set_dir_filter()" in new_code + + compiled_code = compile(new_ast, module_path, "exec") + exec(compiled_code, imported_mod.__dict__) + patched_dir = set(dir(imported_mod)) + for symbol in patched_dir: + assert not symbol.startswith(_PREFIX) + + # Check that there are no new symbols that were not in the original dir() call + assert len(patched_dir.difference(orig_dir)) == 0 + + # Check that all the symbols in the expected set are in the patched dir() result + for name in expected_names: + assert name in patched_dir diff --git a/tests/appsec/iast/aspects/test_add_aspect_fixtures.py b/tests/appsec/iast/aspects/test_add_aspect_fixtures.py index 1fba28d38c9..19a6a97dae7 100644 --- a/tests/appsec/iast/aspects/test_add_aspect_fixtures.py +++ b/tests/appsec/iast/aspects/test_add_aspect_fixtures.py @@ -38,7 +38,7 @@ def test_operator_add_dis( bytecode = dis.Bytecode(mod.do_operator_add_params) dis.dis(mod.do_operator_add_params) - assert bytecode.codeobj.co_names == ("ddtrace_aspects", "add_aspect") + assert bytecode.codeobj.co_names == ("_ddtrace_aspects", "add_aspect") def test_string_operator_add_one_tainted(self): # type: () -> None string_input = taint_pyobject( diff --git a/tests/appsec/iast/aspects/test_add_inplace_aspect_fixtures.py b/tests/appsec/iast/aspects/test_add_inplace_aspect_fixtures.py index 8c6bf538de4..1d59ba41dbc 100644 --- a/tests/appsec/iast/aspects/test_add_inplace_aspect_fixtures.py +++ b/tests/appsec/iast/aspects/test_add_inplace_aspect_fixtures.py @@ -39,7 +39,7 @@ def test_operator_add_inplace_dis( bytecode = dis.Bytecode(mod.do_operator_add_inplace_params) dis.dis(mod.do_operator_add_inplace_params) - assert bytecode.codeobj.co_names == ("ddtrace_aspects", "add_inplace_aspect") + assert bytecode.codeobj.co_names == ("_ddtrace_aspects", "add_inplace_aspect") def test_string_operator_add_inplace_one_tainted(self) -> None: string_input = taint_pyobject( diff --git a/tests/appsec/iast/fixtures/ast/other/with_implemented_dir.py b/tests/appsec/iast/fixtures/ast/other/with_implemented_dir.py new file mode 100644 index 00000000000..e43e0422346 --- /dev/null +++ b/tests/appsec/iast/fixtures/ast/other/with_implemented_dir.py @@ -0,0 +1,10 @@ +symbol1 = "foo" +symbol2 = "bar" +symbol3 = symbol1 + symbol2 +symbol4 = symbol3[1] + +__ddtrace_mark = "baz" + + +def __dir__(): + return ["custom_added"] + [i for i in globals()] diff --git a/tests/appsec/iast/fixtures/ast/other/without_implemented_dir.py b/tests/appsec/iast/fixtures/ast/other/without_implemented_dir.py new file mode 100644 index 00000000000..61cb47bb417 --- /dev/null +++ b/tests/appsec/iast/fixtures/ast/other/without_implemented_dir.py @@ -0,0 +1,6 @@ +symbol1 = "foo" +symbol2 = "bar" +symbol3 = symbol1 + symbol2 +symbol4 = symbol3[1] + +__ddtrace_mark = "baz" diff --git a/tests/appsec/iast_packages/inside_env_runner.py b/tests/appsec/iast_packages/inside_env_runner.py index 530f63152b1..bfd8f6d2eb8 100644 --- a/tests/appsec/iast_packages/inside_env_runner.py +++ b/tests/appsec/iast_packages/inside_env_runner.py @@ -47,8 +47,8 @@ def try_patched(module_name, expect_no_change=False): ), "Patched source is None after patching: Maybe not an error, but something fishy is going on" new_code = unparse(patched_module) assert ( - "import ddtrace.appsec._iast.taint_sinks as ddtrace_taint_sinks" - "\nimport ddtrace.appsec._iast._taint_tracking.aspects as ddtrace_aspects\n" + "import ddtrace.appsec._iast.taint_sinks as _ddtrace_taint_sinks" + "\nimport ddtrace.appsec._iast._taint_tracking.aspects as _ddtrace_aspects\n" ) in new_code, "Patched imports not found" assert "ddtrace_aspects." in new_code, "Patched aspects not found"