From ed7029f4c1a3f1737bd5a3237be6743a9e1c5d02 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Tue, 26 Nov 2024 16:46:42 +0100
Subject: [PATCH] chore(iast): check the return value of astpatch_module
 [backport 2.17] (#11525)

Backport 6205242103a2f232001152a749b8c296540cdc05 from #11520 to 2.17.

## Description

Some testing code were not checking the return value of
`astpatch_module` as it should, making some package tests fail after a
previous PR that fixed the return value of `astpatch_module` when no
patching was done.

Also disable `google.*` from package tests since we have disabled
patching for those.

Signed-off-by: Juanjo Alvarez <juanjo.alvarezmartinez@datadoghq.com>

## 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)

Co-authored-by: Juanjo Alvarez Martinez <juanjo.alvarezmartinez@datadoghq.com>
---
 ddtrace/appsec/_iast/__init__.py              |  4 ++
 .../appsec/iast_packages/inside_env_runner.py |  3 ++
 tests/appsec/iast_packages/test_packages.py   | 42 +++++++++----------
 3 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/ddtrace/appsec/_iast/__init__.py b/ddtrace/appsec/_iast/__init__.py
index a77f166292b..5aab86cf783 100644
--- a/ddtrace/appsec/_iast/__init__.py
+++ b/ddtrace/appsec/_iast/__init__.py
@@ -64,6 +64,10 @@ def ddtrace_iast_flask_patch():
         log.debug("Unexpected exception while AST patching", exc_info=True)
         return
 
+    if not patched_ast:
+        log.debug("Main flask module not patched, probably it was not needed")
+        return
+
     compiled_code = compile(patched_ast, module_path, "exec")
     exec(compiled_code, module.__dict__)  # nosec B102
     sys.modules[module_name] = compiled_code
diff --git a/tests/appsec/iast_packages/inside_env_runner.py b/tests/appsec/iast_packages/inside_env_runner.py
index bfd8f6d2eb8..77f7c52c0e5 100644
--- a/tests/appsec/iast_packages/inside_env_runner.py
+++ b/tests/appsec/iast_packages/inside_env_runner.py
@@ -16,6 +16,9 @@ def _iast_patched_module_and_patched_source(module_name):
     module = importlib.import_module(module_name)
     module_path, patched_module = astpatch_module(module)
 
+    if not patched_module:
+        assert False, "Module %s was not patched" % module_name
+
     compiled_code = compile(patched_module, module_path, "exec")
     exec(compiled_code, module.__dict__)
     return module, patched_module
diff --git a/tests/appsec/iast_packages/test_packages.py b/tests/appsec/iast_packages/test_packages.py
index cde0bf5a89c..048fa2abe05 100644
--- a/tests/appsec/iast_packages/test_packages.py
+++ b/tests/appsec/iast_packages/test_packages.py
@@ -266,27 +266,27 @@ def uninstall(self, python_cmd):
     ),
     PackageForTesting("flask", "2.3.3", "", "", "", test_e2e=False, import_module_to_validate="flask.app"),
     PackageForTesting("fsspec", "2024.5.0", "", "/", ""),
-    PackageForTesting(
-        "google-auth",
-        "2.35.0",
-        "",
-        "",
-        "",
-        import_name="google.auth.crypt.rsa",
-        import_module_to_validate="google.auth.crypt.rsa",
-        expect_no_change=True,
-    ),
-    PackageForTesting(
-        "google-api-core",
-        "2.22.0",
-        "",
-        "",
-        "",
-        import_name="google",
-        import_module_to_validate="google.auth.iam",
-        extras=[("google-cloud-storage", "2.18.2")],
-        test_e2e=True,
-    ),
+    # PackageForTesting(
+    # "google-auth",
+    # "2.35.0",
+    # "",
+    # "",
+    # "",
+    # import_name="google.auth.crypt.rsa",
+    # import_module_to_validate="google.auth.crypt.rsa",
+    # expect_no_change=True,
+    # ),
+    # PackageForTesting(
+    # "google-api-core",
+    # "2.22.0",
+    # "",
+    # "",
+    # "",
+    # import_name="google",
+    # import_module_to_validate="google.auth.iam",
+    # extras=[("google-cloud-storage", "2.18.2")],
+    # test_e2e=True,
+    # ),
     PackageForTesting(
         "google-api-python-client",
         "2.111.0",