Django session codemod to use codemod API

src/codemodder/codemods/base_codemod.py

@@ -75,14 +75,17 @@ def node_position(self, node):
         # See https://github.com/Instagram/LibCST/blob/main/libcst/_metadata_dependent.py#L112
         return self.get_metadata(self.METADATA_DEPENDENCIES[0], node)
 
-    def add_change(self, node, description):
+    def add_change(self, node, description: str, start: bool = True):
         position = self.node_position(node)
-        self.add_change_from_position(position, description)
+        self.add_change_from_position(position, description, start)
 
-    def add_change_from_position(self, position: CodeRange, description):
+    def add_change_from_position(
+        self, position: CodeRange, description: str, start: bool = True
+    ):
+        lineno = position.start.line if start else position.end.line
         self.file_context.codemod_changes.append(
             Change(
-                lineNumber=position.start.line,
+                lineNumber=lineno,
                 description=description,
             )
         )

src/codemodder/codemods/utils.py

@@ -132,3 +132,10 @@ def get_function_name_node(call: cst.Call) -> Optional[cst.Name]:
         case cst.Attribute():
             return call.func.attr
     return None
+
+
+def is_assigned_to_True(original_node: cst.Assign):
+    return (
+        isinstance(original_node.value, cst.Name)
+        and original_node.value.value == "True"
+    )

src/core_codemods/django_session_cookie_secure_off.py

@@ -1,86 +1,66 @@
-from typing import List
 import libcst as cst
-from libcst.codemod import Codemod, CodemodContext
-from libcst.metadata import PositionProvider
-from codemodder.change import Change
-from codemodder.codemods.base_visitor import BaseTransformer
-from codemodder.codemods.base_codemod import (
-    SemgrepCodemod,
-    CodemodMetadata,
-    ReviewGuidance,
-)
-from codemodder.codemods.utils import is_django_settings_file
-from codemodder.file_context import FileContext
+from codemodder.codemods.api import SemgrepCodemod
+from codemodder.codemods.base_codemod import ReviewGuidance
+from codemodder.codemods.utils import is_django_settings_file, is_assigned_to_True
 
 
-class DjangoSessionCookieSecureOff(SemgrepCodemod, Codemod):
-    METADATA = CodemodMetadata(
-        DESCRIPTION=("Sets Django's `SESSION_COOKIE_SECURE` flag if off or missing."),
-        NAME="django-session-cookie-secure-off",
-        REVIEW_GUIDANCE=ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW,
-        REFERENCES=[
-            {
-                "url": "https://owasp.org/www-community/controls/SecureCookieAttribute",
-                "description": "",
-            },
-            {
-                "url": "https://docs.djangoproject.com/en/4.2/ref/settings/#session-cookie-secure",
-                "description": "",
-            },
-        ],
-    )
+class DjangoSessionCookieSecureOff(SemgrepCodemod):
+    NAME = "django-session-cookie-secure-off"
+    DESCRIPTION = "Sets Django's `SESSION_COOKIE_SECURE` flag if off or missing."
     SUMMARY = "Secure Setting for Django `SESSION_COOKIE_SECURE` flag"
-    CHANGE_DESCRIPTION = METADATA.DESCRIPTION
-    YAML_FILES = [
-        "detect-django-settings.yaml",
+    REVIEW_GUIDANCE = ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW
+    REFERENCES = [
+        {
+            "url": "https://owasp.org/www-community/controls/SecureCookieAttribute",
+            "description": "",
+        },
+        {
+            "url": "https://docs.djangoproject.com/en/4.2/ref/settings/#session-cookie-secure",
+            "description": "",
+        },
     ]
 
-    METADATA_DEPENDENCIES = (PositionProvider,)
-
-    def __init__(self, codemod_context: CodemodContext, *args):
-        Codemod.__init__(self, codemod_context)
-        SemgrepCodemod.__init__(self, *args)
-
-    def transform_module_impl(self, tree: cst.Module) -> cst.Module:
-        if is_django_settings_file(self.file_context.file_path):
-            transformer = SessionCookieSecureTransformer(
-                self.context, self.file_context, self.file_context.findings
-            )
-            new_tree = transformer.transform_module(tree)
-            if transformer.changes_in_file:
-                self.file_context.codemod_changes.extend(transformer.changes_in_file)
-            return new_tree
-        return tree
-
+    @classmethod
+    def rule(cls):
+        return """
+        rules:
+          - id: django-session-cookie-secure-off
+            # This pattern creates one finding with no text for settings.py file.
+            pattern-regex: ^
+            paths:
+              include:
+               - settings.py
+        """
 
-class SessionCookieSecureTransformer(BaseTransformer):
-    def __init__(
-        self, codemod_context: CodemodContext, file_context: FileContext, results
-    ):
-        super().__init__(codemod_context, results)
-        self.line_exclude = file_context.line_exclude
-        self.line_include = file_context.line_include
-        self.changes_in_file: List[Change] = []
+    def __init__(self, *args):
+        super().__init__(*args)
+        self.is_django_settings_file = is_django_settings_file(
+            self.file_context.file_path
+        )
         self.flag_correctly_set = False
 
+    def visit_Module(self, _: cst.Module) -> bool:
+        """
+        Only visit module with this codemod if it's a settings.py file.
+        """
+        return self.is_django_settings_file
+
     def leave_Module(
         self, original_node: cst.Module, updated_node: cst.Module
     ) -> cst.Module:
         """
         Handle case for `SESSION_COOKIE_SECURE`  is missing from settings.py
         """
-        if self.flag_correctly_set or len(self.changes_in_file):
+        if not self.is_django_settings_file:
+            return updated_node
+
+        if self.flag_correctly_set or len(self.file_context.codemod_changes):
             # Nothing to do at the end of the module if
             # `SESSION_COOKIE_SECURE = True` or if assigned to
             # something else and we changed it in `leave_Assign`.
             return updated_node
 
-        # line_number is the end of the module where we will insert the new flag.
-        pos_to_match = self.node_position(original_node)
-        line_number = pos_to_match.end.line
-        self.changes_in_file.append(
-            Change(line_number, DjangoSessionCookieSecureOff.CHANGE_DESCRIPTION)
-        )
+        self.add_change(original_node, self.CHANGE_DESCRIPTION, start=False)
         final_line = cst.parse_statement("SESSION_COOKIE_SECURE = True")
         new_body = updated_node.body + (final_line,)
         return updated_node.with_changes(body=new_body)
@@ -100,10 +80,7 @@ def leave_Assign(
                 return updated_node
 
             # SESSION_COOKIE_SECURE = anything other than True
-            line_number = pos_to_match.start.line
-            self.changes_in_file.append(
-                Change(line_number, DjangoSessionCookieSecureOff.CHANGE_DESCRIPTION)
-            )
+            self.add_change(original_node, self.CHANGE_DESCRIPTION)
             return updated_node.with_changes(value=cst.Name("True"))
         return updated_node
 
@@ -116,10 +93,3 @@ def is_session_cookie_secure(original_node: cst.Assign):
     return (
         isinstance(target_var, cst.Name) and target_var.value == "SESSION_COOKIE_SECURE"
     )
-
-
-def is_assigned_to_True(original_node: cst.Assign):
-    return (
-        isinstance(original_node.value, cst.Name)
-        and original_node.value.value == "True"
-    )

tests/codemods/test_django_session_cookie_secure_off.py

@@ -18,13 +18,15 @@ def test_not_settings_dot_py(self, tmpdir):
         input_code = """SESSION_COOKIE_SECURE = True"""
         expected = input_code
         self.run_and_assert_filepath(django_root, file_path, input_code, expected)
+        assert len(self.file_context.codemod_changes) == 0
 
     def test_no_manage_dot_py(self, tmpdir):
         django_root, settings_folder = self.create_dir_structure(tmpdir)
         file_path = settings_folder / "settings.py"
         input_code = """SESSION_COOKIE_SECURE = True"""
         expected = input_code
         self.run_and_assert_filepath(django_root, file_path, input_code, expected)
+        assert len(self.file_context.codemod_changes) == 0
 
     def test_settings_dot_py_secure_true(self, tmpdir):
         django_root, settings_folder = self.create_dir_structure(tmpdir)
@@ -34,6 +36,7 @@ def test_settings_dot_py_secure_true(self, tmpdir):
 SESSION_COOKIE_SECURE = True
 """
         self.run_and_assert_filepath(django_root, file_path, input_code, input_code)
+        assert len(self.file_context.codemod_changes) == 0
 
     @pytest.mark.parametrize("value", ["False", "gibberish"])
     def test_settings_dot_py_secure_bad(self, tmpdir, value):
@@ -47,6 +50,7 @@ def test_settings_dot_py_secure_bad(self, tmpdir, value):
 SESSION_COOKIE_SECURE = True
 """
         self.run_and_assert_filepath(django_root, file_path, input_code, expected)
+        assert len(self.file_context.codemod_changes) == 1
 
     def test_settings_dot_py_secure_missing(self, tmpdir):
         django_root, settings_folder = self.create_dir_structure(tmpdir)
@@ -58,3 +62,4 @@ def test_settings_dot_py_secure_missing(self, tmpdir):
 SESSION_COOKIE_SECURE = True
 """
         self.run_and_assert_filepath(django_root, file_path, input_code, expected)
+        assert len(self.file_context.codemod_changes) == 1