fixup! Initial implementation of transformation

Author: andrecsilva

src/core_codemods/docs/sonar_python_use-secure-protocols.md

@@ -9,13 +9,16 @@ Our changes look like the following:
 
 - ftp_con = ftplib.FTP("ftp.example.com")
 + ftp_con = ftplib.FTP_TLS("ftp.example.com")
-smtp_con = smtplib.SMTP("smtp.example.com", port=587)
 + smtp_context = ssl.create_default_context()
 + smtp_context.verify_mode = ssl.CERT_REQUIRED
 + smtp_context.check_hostname = True
+smtp_con = smtplib.SMTP("smtp.example.com", port=587)
 + smtp.starttls(context=smtp_context)
 
 
++ smtp_context_1 = ssl.create_default_context()
++ smtp_context_1.verify_mode = ssl.CERT_REQUIRED
++ smtp_context_1.check_hostname = True
 - smtp_con_2 = smtplib.SMTP("smtp.gmail.com")
-+ smtp_con_2 = smtplib.SMTP_SSL("smtp.gmail.com")
++ smtp_con_2 = smtplib.SMTP_SSL("smtp.gmail.com", context=smtp_context_1)
 ```

src/core_codemods/sonar/sonar_use_secure_protocols.py

@@ -1,4 +1,5 @@
 import libcst as cst
+from libcst.codemod import CodemodContext
 
 from codemodder.codemods.base_codemod import Metadata, ReviewGuidance, ToolRule
 from codemodder.codemods.libcst_transformer import (
@@ -7,6 +8,8 @@
 )
 from codemodder.codemods.utils_mixin import NameAndAncestorResolutionMixin
 from codemodder.codetf import Reference
+from codemodder.file_context import FileContext
+from codemodder.result import Result
 from core_codemods.sonar.api import SonarCodemod
 
 rules = [
@@ -23,28 +26,55 @@ class SonarUseSecureProtocolsTransformer(
 ):
     change_description = "Modified URLs or calls to use secure protocols"
 
+    def __init__(
+        self,
+        context: CodemodContext,
+        results: list[Result] | None,
+        file_context: FileContext,
+        _transformer: bool = False,
+    ):
+        self.nodes_memory_with_context_name: dict[cst.CSTNode, str] = {}
+        super().__init__(context, results, file_context, _transformer)
+
     def _match_and_handle_statement(
         self, possible_smtp_call, original_node_statement, updated_node_statement
     ):
         maybe_name = self.find_base_name(possible_smtp_call)
         match possible_smtp_call:
             case cst.Call() if maybe_name == "smtplib.SMTP":
+                # get the stored context_name or create a new one:
+                if possible_smtp_call in self.nodes_memory_with_context_name:
+                    context_name = self.nodes_memory_with_context_name[
+                        possible_smtp_call
+                    ]
+                else:
+                    context_name = self.generate_available_name(
+                        original_node_statement, ["smtp_context"]
+                    )
+
                 new_statements = []
                 new_statements.append(
-                    cst.parse_statement("smtp_context = ssl.create_default_context()")
+                    cst.parse_statement(
+                        f"{context_name} = ssl.create_default_context()"
+                    )
                 )
                 new_statements.append(
-                    cst.parse_statement("smtp_context.verify_mode = ssl.CERT_REQUIRED")
+                    cst.parse_statement(
+                        f"{context_name}.verify_mode = ssl.CERT_REQUIRED"
+                    )
                 )
                 new_statements.append(
-                    cst.parse_statement("smtp_context.check_hostname = True")
+                    cst.parse_statement(f"{context_name}.check_hostname = True")
                 )
                 new_statements.append(updated_node_statement)
-                # TODO don't append this if we changed the call to SSL version
-                new_statements.append(
-                    cst.parse_statement("smtp.starttls(context=smtp_context)")
-                )
-                self.add_needed_import("smtp")
+                # don't append this if we changed the call to SSL version
+                if possible_smtp_call in self.nodes_memory_with_context_name:
+                    self.nodes_memory_with_context_name.pop(possible_smtp_call)
+                else:
+                    new_statements.append(
+                        cst.parse_statement(f"smtplib.starttls(context={context_name})")
+                    )
+                self.add_needed_import("smtplib")
                 self.add_needed_import("ssl")
                 self.report_change(possible_smtp_call)
                 return cst.FlattenSentinel(new_statements)
@@ -70,6 +100,8 @@ def leave_Call(self, original_node, updated_node):
                     self.report_change(original_node)
                     self.add_needed_import("ftplib")
                     return updated_node.with_changes(func=new_func)
+                # Just using ssl.create_default_context() may not be enough for older python versions
+                # See https://stackoverflow.com/questions/33857698/sending-email-from-python-using-starttls
                 case "smtplib.SMTP":
                     # port is the second positional, check that
                     maybe_port_value = (
@@ -94,37 +126,28 @@ def leave_Call(self, original_node, updated_node):
                     )
                     match maybe_port_value:
                         case None:
-                            new_func = cst.parse_expression("smtplib.SMTP_SSL")
-                            self.report_change(original_node)
-                            self.add_needed_import("smtplib")
-                            new_args = [
-                                *original_node.args,
-                                cst.Arg(
-                                    keyword=cst.Name("context"),
-                                    value=cst.Name("smtp_context"),
-                                ),
-                            ]
-                            return updated_node.with_changes(
-                                func=new_func, args=new_args
-                            )
-                        # TODO still needs the context object statements here
-                        # TODO only change this if it mathces the statement pattern in leave_statemenet
-                        # TODO use a flag for this in visit_SimpleStatement
+                            self._change_to_smtp_ssl(original_node, updated_node)
                         case cst.Integer() if maybe_port_value == "0":
-                            new_func = cst.parse_expression("smtplib.SMTP_SSL")
-                            self.report_change(original_node)
-                            self.add_needed_import("smtplib")
-                            new_args = [
-                                *original_node.args,
-                                cst.Arg(
-                                    keyword=cst.Name("context"),
-                                    value=cst.Name("smtp_context"),
-                                ),
-                            ]
-                            return updated_node.with_changes(func=new_func)
+                            self._change_to_smtp_ssl(original_node, updated_node)
 
         return updated_node
 
+    def _change_to_smtp_ssl(self, original_node, updated_node):
+        # remember this node so we don't add the starttls
+        new_func = cst.parse_expression("smtplib.SMTP_SSL")
+
+        context_name = self.generate_available_name(original_node, ["smtp_context"])
+        self.nodes_memory_with_context_name[original_node] = context_name
+
+        new_args = [
+            *original_node.args,
+            cst.Arg(
+                keyword=cst.Name("context"),
+                value=cst.Name(context_name),
+            ),
+        ]
+        return updated_node.with_changes(func=new_func, args=new_args)
+
     def leave_SimpleString(
         self, original_node: cst.SimpleString, updated_node: cst.SimpleString
     ) -> cst.BaseExpression: