From f743d0e51f76a64c3c14c32e4faf8f50ec705be3 Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Tue, 14 Nov 2023 14:44:39 -0500
Subject: [PATCH] Preserve custom loader with harden-pyyaml

---
 src/core_codemods/harden_pyyaml.py   | 14 ++++++++++----
 tests/codemods/base_codemod_test.py  |  6 +++---
 tests/codemods/test_harden_pyyaml.py | 20 ++++++++++++++++++++
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/core_codemods/harden_pyyaml.py b/src/core_codemods/harden_pyyaml.py
index 5638eccd..2bec5395 100644
--- a/src/core_codemods/harden_pyyaml.py
+++ b/src/core_codemods/harden_pyyaml.py
@@ -31,8 +31,11 @@ def rule(cls):
                       - metavariable-pattern:
                           metavariable: $ARG
                           patterns:
-                            - pattern-not:
-                                pattern: yaml.SafeLoader
+                            - pattern-either:
+                                - pattern: yaml.Loader
+                                - pattern: yaml.BaseLoader
+                                - pattern: yaml.FullLoader
+                                - pattern: yaml.UnsafeLoader
                   - patterns:
                       - pattern: yaml.load(...)
                       - pattern-inside: |
@@ -42,8 +45,11 @@ def rule(cls):
                       - metavariable-pattern:
                           metavariable: $ARG
                           patterns:
-                            - pattern-not:
-                                pattern: yaml.SafeLoader
+                            - pattern-either:
+                                - pattern: yaml.Loader
+                                - pattern: yaml.BaseLoader
+                                - pattern: yaml.FullLoader
+                                - pattern: yaml.UnsafeLoader
 
         """
 
diff --git a/tests/codemods/base_codemod_test.py b/tests/codemods/base_codemod_test.py
index 78d7fbf0..1a80b099 100644
--- a/tests/codemods/base_codemod_test.py
+++ b/tests/codemods/base_codemod_test.py
@@ -68,7 +68,7 @@ def setup_class(cls):
 
     def results_by_id_filepath(self, input_code, file_path):
         with open(file_path, "w", encoding="utf-8") as tmp_file:
-            tmp_file.write(input_code)
+            tmp_file.write(dedent(input_code))
 
         name = self.codemod.name()
         results = self.registry.match_codemods(codemod_include=[name])
@@ -82,7 +82,7 @@ def run_and_assert_filepath(self, root, file_path, input_code, expected):
             registry=mock.MagicMock(),
             repo_manager=mock.MagicMock(),
         )
-        input_tree = cst.parse_module(input_code)
+        input_tree = cst.parse_module(dedent(input_code))
         all_results = self.results_by_id_filepath(input_code, file_path)
         results = all_results.results_for_rule_and_file(self.codemod.name(), file_path)
         self.file_context = FileContext(
@@ -99,7 +99,7 @@ def run_and_assert_filepath(self, root, file_path, input_code, expected):
         )
         output_tree = command_instance.transform_module(input_tree)
 
-        assert output_tree.code == expected
+        assert output_tree.code == dedent(expected)
 
 
 class BaseDjangoCodemodTest(BaseSemgrepCodemodTest):
diff --git a/tests/codemods/test_harden_pyyaml.py b/tests/codemods/test_harden_pyyaml.py
index 0e3e47f3..d783e2f1 100644
--- a/tests/codemods/test_harden_pyyaml.py
+++ b/tests/codemods/test_harden_pyyaml.py
@@ -60,3 +60,23 @@ def test_import_alias(self, tmpdir):
 deserialized_data = yam.load(data, yam.SafeLoader)
 """
         self.run_and_assert(tmpdir, input_code, expected)
+
+    def test_preserve_custom_loader(self, tmpdir):
+        expected = input_code = """
+        import yaml
+        from custom import CustomLoader
+
+        yaml.load(data, CustomLoader)
+        """
+
+        self.run_and_assert(tmpdir, input_code, expected)
+
+    def test_preserve_custom_loader_kwarg(self, tmpdir):
+        expected = input_code = """
+        import yaml
+        from custom import CustomLoader
+
+        yaml.load(data, Loader=CustomLoader)
+        """
+
+        self.run_and_assert(tmpdir, input_code, expected)