Redo tempfile.mktemp codemod (#577)

* Change tempfile.mktemp codemod to correctly create a NamedTemporary

* refactor

* handle import from cases

* test all other import types

* handle with block

* handle args

* change all tests

* add docs

* remove check for alias

* add demonstrative unit test

* check all simplestatements, not just directly under a module

Author: clavedeluna

integration_tests/sonar/test_sonar_tempfile_mktemp.py

@@ -6,7 +6,10 @@
 class TestTempfileMktemp(SonarIntegrationTest):
     codemod = SonarTempfileMktemp
     code_path = "tests/samples/tempfile_mktemp.py"
-    replacement_lines = [(3, "tempfile.mkstemp()\n")]
-    expected_diff = "--- \n+++ \n@@ -1,3 +1,3 @@\n import tempfile\n \n-tempfile.mktemp()\n+tempfile.mkstemp()\n"
+    replacement_lines = [
+        (3, "with tempfile.NamedTemporaryFile(delete=False) as tf:\n"),
+        (4, "    filename = tf.name\n"),
+    ]
+    expected_diff = "--- \n+++ \n@@ -1,3 +1,4 @@\n import tempfile\n \n-filename = tempfile.mktemp()\n+with tempfile.NamedTemporaryFile(delete=False) as tf:\n+    filename = tf.name\n"
     expected_line_change = "3"
     change_description = TempfileMktempTransformer.change_description

integration_tests/test_tempfile_mktemp.py

@@ -7,10 +7,12 @@ class TestTempfileMktemp(BaseIntegrationTest):
     original_code = """
     import tempfile
 
-    tempfile.mktemp()
-    var = "hello"
+    filename = tempfile.mktemp()
     """
-    replacement_lines = [(3, "tempfile.mkstemp()\n")]
-    expected_diff = '--- \n+++ \n@@ -1,4 +1,4 @@\n import tempfile\n \n-tempfile.mktemp()\n+tempfile.mkstemp()\n var = "hello"\n'
+    replacement_lines = [
+        (3, "with tempfile.NamedTemporaryFile(delete=False) as tf:\n"),
+        (4, "    filename = tf.name\n"),
+    ]
+    expected_diff = "--- \n+++ \n@@ -1,3 +1,4 @@\n import tempfile\n \n-filename = tempfile.mktemp()\n+with tempfile.NamedTemporaryFile(delete=False) as tf:\n+    filename = tf.name\n"
     expected_line_change = "3"
     change_description = TempfileMktempTransformer.change_description

src/codemodder/scripts/generate_docs.py

@@ -94,7 +94,7 @@ class DocMetadata:
     ),
     "secure-tempfile": DocMetadata(
         importance="High",
-        guidance_explained="We believe this codemod is safe and will cause no unexpected errors.",
+        guidance_explained="We believe this codemod is safe. You should review this code before merging to make sure temporary files are created, used, and closed according to your expectations.",
     ),
     "upgrade-sslcontext-minimum-version": DocMetadata(
         importance="High",

src/core_codemods/docs/pixee_python_secure-tempfile.md

@@ -1,12 +1,14 @@
-This codemod replaces all `tempfile.mktemp` calls to the more secure `tempfile.mkstemp`.
+This codemod replaces all `tempfile.mktemp` calls with the more secure `tempfile.NamedTemporaryFile`
 
-The Python [tempfile documentation](https://docs.python.org/3/library/tempfile.html#tempfile.mktemp) is explicit
-that `tempfile.mktemp` should be deprecated to avoid an unsafe and unexpected race condition.
-The changes from this codemod look like this:
+The Python [tempfile documentation](https://docs.python.org/3/library/tempfile.html#tempfile.mktemp) is explicit that `tempfile.mktemp` should be deprecated to avoid an unsafe and unexpected race condition. `tempfile.mktemp` does not handle the possibility that the returned file name could already be used by another process by the time your code opens the file. A more secure approach to create temporary files is to use `tempfile.NamedTemporaryFile` which will create the file for you and handle all security conditions. 
 
+The changes from this codemod look like this:
 
 ```diff
   import tempfile
-- tempfile.mktemp(...)
-+ tempfile.mkstemp(...)
+- filename = tempfile.mktemp()
++ with tempfile.NamedTemporaryFile(delete=False) as tf:
++     filename = tf.name
 ```
+
+The change sets `delete=False` to closely follow your code's intention when calling `tempfile.mktemp`. However, you should use this as a starting point to determine when your temporary file should be deleted.

src/core_codemods/tempfile_mktemp.py

@@ -1,44 +1,130 @@
+from textwrap import dedent
+from typing import Optional
+
+import libcst as cst
+from libcst import matchers
+
 from codemodder.codemods.libcst_transformer import (
     LibcstResultTransformer,
     LibcstTransformerPipeline,
 )
-from codemodder.codemods.semgrep import SemgrepRuleDetector
-from codemodder.codemods.utils_mixin import NameResolutionMixin
+from codemodder.codemods.utils_mixin import NameAndAncestorResolutionMixin
+from codemodder.utils.utils import clean_simplestring
 from core_codemods.api import CoreCodemod, Metadata, Reference, ReviewGuidance
 
 
-class TempfileMktempTransformer(LibcstResultTransformer, NameResolutionMixin):
+class TempfileMktempTransformer(
+    LibcstResultTransformer, NameAndAncestorResolutionMixin
+):
     change_description = "Replaces `tempfile.mktemp` with `tempfile.mkstemp`."
     _module_name = "tempfile"
 
-    def on_result_found(self, original_node, updated_node):
-        maybe_name = self.get_aliased_prefix_name(original_node, self._module_name)
-        if (maybe_name := maybe_name or self._module_name) == self._module_name:
-            self.add_needed_import(self._module_name)
-        self.remove_unused_import(original_node)
-        return self.update_call_target(updated_node, maybe_name, "mkstemp")
+    def leave_SimpleStatementLine(self, original_node, updated_node):
+        match original_node:
+            case cst.SimpleStatementLine(body=[bsstmt]):
+                return self.check_mktemp(original_node, bsstmt)
+        return updated_node
+
+    def check_mktemp(
+        self, original_node: cst.SimpleStatementLine, bsstmt: cst.BaseSmallStatement
+    ) -> cst.SimpleStatementLine | cst.FlattenSentinel:
+        if maybe_tuple := self._is_assigned_to_mktemp(bsstmt):  # type: ignore
+            assign_name, call = maybe_tuple
+            return self.report_and_change(call, assign_name)
+        if maybe_tuple := self._mktemp_is_sink(bsstmt):
+            wrapper_func_name, call = maybe_tuple
+            return self.report_and_change(call, wrapper_func_name, assignment=False)
+        return original_node
+
+    def report_and_change(
+        self, node: cst.Call, name: cst.Name, assignment=True
+    ) -> cst.FlattenSentinel:
+        self.report_change(node)
+        self.add_needed_import(self._module_name)
+        self.remove_unused_import(node)
+        with_block = (
+            f"{name.value} = tf.name" if assignment else f"{name.value}(tf.name)"
+        )
+        new_stmt = dedent(
+            f"""
+        with tempfile.NamedTemporaryFile({self._make_args(node)}) as tf:
+            {with_block}
+        """
+        ).rstrip()
+        return cst.FlattenSentinel(
+            [
+                cst.parse_statement(new_stmt),
+            ]
+        )
+
+    def _make_args(self, node: cst.Call) -> str:
+        """Convert args passed to tempfile.mktemp() to string for args to tempfile.NamedTemporaryFile"""
+
+        default = "delete=False"
+        if not node.args:
+            return default
+        new_args = ""
+        arg_keys = ("suffix", "prefix", "dir")
+        for idx, arg in enumerate(node.args):
+            cst.ensure_type(val := arg.value, cst.SimpleString)
+            new_args += f'{arg_keys[idx]}="{clean_simplestring(val)}", '
+        return f"{new_args}{default}"
+
+    def _is_assigned_to_mktemp(
+        self, bsstmt: cst.BaseSmallStatement
+    ) -> Optional[tuple[cst.Name, cst.Call]]:
+        match bsstmt:
+            case cst.Assign(value=value, targets=targets):
+                maybe_value = self._is_mktemp_call(value)  # type: ignore
+                if maybe_value and all(
+                    map(
+                        lambda t: matchers.matches(
+                            t, matchers.AssignTarget(target=matchers.Name())
+                        ),
+                        targets,  # type: ignore
+                    )
+                ):
+                    # # Todo: handle multiple potential targets
+                    return (targets[0].target, maybe_value)
+            case cst.AnnAssign(target=target, value=value):
+                maybe_value = self._is_mktemp_call(value)  # type: ignore
+                if maybe_value and isinstance(target, cst.Name):  # type: ignore
+                    return (target, maybe_value)
+        return None
+
+    def _is_mktemp_call(self, value) -> Optional[cst.Call]:
+        match value:
+            case cst.Call() if self.find_base_name(value.func) == "tempfile.mktemp":
+                return value
+        return None
+
+    def _mktemp_is_sink(
+        self, bsstmt: cst.BaseSmallStatement
+    ) -> Optional[tuple[cst.Name, cst.Call]]:
+        match bsstmt:
+            case cst.Expr(value=cst.Call() as call):
+                if not (args := call.args):
+                    return None
+
+                # todo: handle more complex cases of mktemp in different arg pos
+                match first_arg_call := args[0].value:
+                    case cst.Call():
+                        if maybe_value := self._is_mktemp_call(first_arg_call):  # type: ignore
+                            wrapper_func = call.func
+                            return (wrapper_func, maybe_value)
+        return None
 
 
 TempfileMktemp = CoreCodemod(
     metadata=Metadata(
         name="secure-tempfile",
         summary="Upgrade and Secure Temp File Creation",
-        review_guidance=ReviewGuidance.MERGE_WITHOUT_REVIEW,
+        review_guidance=ReviewGuidance.MERGE_AFTER_REVIEW,
         references=[
             Reference(
                 url="https://docs.python.org/3/library/tempfile.html#tempfile.mktemp"
             ),
         ],
     ),
-    detector=SemgrepRuleDetector(
-        """
-        rules:
-          - patterns:
-            - pattern: tempfile.mktemp(...)
-            - pattern-inside: |
-                import tempfile
-                ...
-        """
-    ),
     transformer=LibcstTransformerPipeline(TempfileMktempTransformer),
 )

tests/codemods/sonar/test_sonar_tempfile_mktemp.py

@@ -15,12 +15,13 @@ def test_simple(self, tmpdir):
         input_code = """
         import tempfile
         
-        tempfile.mktemp()
+        filename = tempfile.mktemp()
         """
         expected = """
         import tempfile
         
-        tempfile.mkstemp()
+        with tempfile.NamedTemporaryFile(delete=False) as tf:
+            filename = tf.name
         """
         issues = {
             "issues": [