Preserve custom loader with harden-pyyaml

pixee · Nov 14, 2023 · f743d0e · f743d0e
1 parent 954a20e
commit f743d0e
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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)