add jinja2 autoescape codemod

integration_tests/test_jinja2_autoescape.py

@@ -0,0 +1,21 @@
+from core_codemods.enable_jinja2_autoescape import EnableJinja2Autoescape
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestEnableJinja2Autoescape(BaseIntegrationTest):
+    codemod = EnableJinja2Autoescape
+    code_path = "tests/samples/jinja2_autoescape.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path,
+        [
+            (1, "env = Environment(autoescape=True)\n"),
+            (2, "env = Environment(autoescape=True)\n"),
+        ],
+    )
+    expected_diff = "--- \n+++ \n@@ -1,3 +1,3 @@\n from jinja2 import Environment\n-env = Environment()\n-env = Environment(autoescape=False)\n+env = Environment(autoescape=True)\n+env = Environment(autoescape=True)\n"
+    expected_line_change = "2"
+    num_changes = 2
+    change_description = EnableJinja2Autoescape.CHANGE_DESCRIPTION

src/codemodder/codemods/api/helpers.py

@@ -47,18 +47,41 @@ def update_assign_rhs(self, updated_node: cst.Assign, rhs: str):
     def parse_expression(self, expression: str):
         return cst.parse_expression(expression)
 
-    def replace_arg(self, original_node, target_arg_name, target_arg_replacement_val):
-        """Given a node, return its args with one arg's value changed"""
+    def replace_arg(
+        self,
+        original_node,
+        target_arg_name,
+        target_arg_replacement_val,
+        add_if_missing=False,
+    ):
+        """Given a node, return its args with one arg's value changed.
+
+        If add_if_missing is True, then if target arg is not present, add it.
+        """
         assert hasattr(original_node, "args")
         new_args = []
+        arg_added = False
+
         for arg in original_node.args:
             if matchers.matches(arg.keyword, matchers.Name(target_arg_name)):
                 new = cst.Arg(
                     keyword=cst.parse_expression(target_arg_name),
                     value=cst.parse_expression(target_arg_replacement_val),
                     equal=arg.equal,
                 )
+                arg_added = True
             else:
                 new = arg
             new_args.append(new)
+
+        if add_if_missing and not arg_added:
+            new = cst.Arg(
+                keyword=cst.parse_expression(target_arg_name),
+                value=cst.parse_expression(target_arg_replacement_val),
+                equal=cst.AssignEqual(
+                    whitespace_before=cst.SimpleWhitespace(""),
+                    whitespace_after=cst.SimpleWhitespace(""),
+                ),
+            )
+            new_args.append(new)
         return new_args

src/core_codemods/__init__.py

@@ -2,6 +2,7 @@
 
 from .django_debug_flag_on import DjangoDebugFlagOn
 from .django_session_cookie_secure_off import DjangoSessionCookieSecureOff
+from .enable_jinja2_autoescape import EnableJinja2Autoescape
 from .fix_mutable_params import FixMutableParams
 from .harden_pyyaml import HardenPyyaml
 from .harden_ruamel import HardenRuamel
@@ -27,6 +28,7 @@
     codemods=[
         DjangoDebugFlagOn,
         DjangoSessionCookieSecureOff,
+        EnableJinja2Autoescape,
         FixMutableParams,
         HardenPyyaml,
         HardenRuamel,

src/core_codemods/enable_jinja2_autoescape.py

@@ -0,0 +1,34 @@
+from codemodder.codemods.base_codemod import ReviewGuidance
+from codemodder.codemods.api import SemgrepCodemod
+
+
+class EnableJinja2Autoescape(SemgrepCodemod):
+    NAME = "enable-jinja2-autoescape"
+    REVIEW_GUIDANCE = ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW
+    SUMMARY = "Enable jinja2 autoescape"
+    DESCRIPTION = "Makes the `autoescape` parameter to jinja2.Environment be `True`."
+
+    @classmethod
+    def rule(cls):
+        return """
+            rules:
+                - patterns:
+                  - pattern-either:
+                    - pattern: |
+                        jinja2.Environment()
+                    - pattern: |
+                        jinja2.Environment(...)
+                    - pattern: |
+                        jinja2.Environment(..., autoescape=False, ...)
+                  - pattern-not: |
+                      jinja2.Environment(..., autoescape=True, ...)
+                  - pattern-inside: |
+                      import jinja2
+                        ...
+        """
+
+    def on_result_found(self, original_node, updated_node):
+        new_args = self.replace_arg(
+            original_node, "autoescape", "True", add_if_missing=True
+        )
+        return self.update_arg_target(updated_node, new_args)

tests/codemods/test_enable_jinja2_autoescape.py

@@ -0,0 +1,88 @@
+import pytest
+from codemodder.codemods.enable_jinja2_autoescape import EnableJinja2Autoescape
+from tests.codemods.base_codemod_test import BaseSemgrepCodemodTest
+
+
+class TestEnableJinja2Autoescape(BaseSemgrepCodemodTest):
+    codemod = EnableJinja2Autoescape
+
+    def test_rule_ids(self):
+        assert self.codemod.RULE_IDS == ["enable-jinja2-autoescape"]
+
+    def test_import(self, tmpdir):
+        input_code = """import jinja2
+env = jinja2.Environment()
+var = "hello"
+"""
+        expexted_output = """import jinja2
+env = jinja2.Environment(autoescape=True)
+var = "hello"
+"""
+
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_import_with_arg(self, tmpdir):
+        input_code = """import jinja2
+env = jinja2.Environment(autoescape=False)
+var = "hello"
+"""
+        expexted_output = """import jinja2
+env = jinja2.Environment(autoescape=True)
+var = "hello"
+"""
+
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_import_with_more_args(self, tmpdir):
+        input_code = """import jinja2
+env = jinja2.Environment(loader=some_loader, autoescape=False)
+var = "hello"
+"""
+        expexted_output = """import jinja2
+env = jinja2.Environment(loader=some_loader, autoescape=True)
+var = "hello"
+"""
+
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_import_with_other_args(self, tmpdir):
+        input_code = """import jinja2
+env = jinja2.Environment(loader=some_loader)
+var = "hello"
+"""
+        expexted_output = """import jinja2
+env = jinja2.Environment(loader=some_loader, autoescape=True)
+var = "hello"
+"""
+
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_from_import(self, tmpdir):
+        input_code = """from jinja2 import Environment
+env = Environment()
+var = "hello"
+"""
+        expexted_output = """from jinja2 import Environment
+env = Environment(autoescape=True)
+var = "hello"
+"""
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_import_alias(self, tmpdir):
+        input_code = """import jinja2 as templating
+env = templating.Environment()
+var = "hello"
+"""
+        expexted_output = """import jinja2 as templating
+env = templating.Environment(autoescape=True)
+var = "hello"
+"""
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_autoescape_enabled(self, tmpdir):
+        input_code = """import jinja2
+env = jinja2.Environment(autoescape=True)
+var = "hello"
+"""
+        expexted_output = input_code
+        self.run_and_assert(tmpdir, input_code, expexted_output)

tests/samples/jinja2_autoescape.py

@@ -0,0 +1,4 @@
+from jinja2 import Environment
+
+env = Environment()
+env = Environment(autoescape=False)

tests/samples/jwt_decode_verify.py

@@ -8,7 +8,11 @@
 
 encoded_jwt = jwt.encode(payload, SECRET_KEY, algorithm="HS256")
 
-decoded_payload = jwt.decode(encoded_jwt, SECRET_KEY, algorithms=["HS256"], verify=False)
-decoded_payload = jwt.decode(encoded_jwt, SECRET_KEY, algorithms=["HS256"], options={"verify_signature": False})
+decoded_payload = jwt.decode(
+    encoded_jwt, SECRET_KEY, algorithms=["HS256"], verify=False
+)
+decoded_payload = jwt.decode(
+    encoded_jwt, SECRET_KEY, algorithms=["HS256"], options={"verify_signature": False}
+)
 
 var = "something"