Jinja autoescape (#50)

---------

Co-authored-by: Dan D'Avella <[email protected]>

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,
+        [
+            (2, "env = Environment(autoescape=True)\n"),
+            (3, "env = Environment(autoescape=True)\n"),
+        ],
+    )
+    expected_diff = "--- \n+++ \n@@ -1,4 +1,4 @@\n from jinja2 import Environment\n \n-env = Environment()\n-env = Environment(autoescape=False)\n+env = Environment(autoescape=True)\n+env = Environment(autoescape=True)\n"
+    expected_line_change = "3"
+    num_changes = 2
+    change_description = EnableJinja2Autoescape.CHANGE_DESCRIPTION

src/codemodder/codemods/api/__init__.py

@@ -121,6 +121,15 @@ def __init__(
         _SemgrepCodemod.__init__(self, execution_context, file_context)
         BaseTransformer.__init__(self, codemod_context, self._results)
 
+    def _new_or_updated_node(self, original_node, updated_node):
+        if self.node_is_selected(original_node):
+            self.report_change(original_node)
+            if (attr := getattr(self, "on_result_found", None)) is not None:
+                # pylint: disable=not-callable
+                new_node = attr(original_node, updated_node)
+                return new_node
+        return updated_node
+
     # TODO: there needs to be a way to generalize this so that it applies
     # more broadly than to just a specific kind of node. There's probably a
     # decent way to do this with metaprogramming. We could either apply it
@@ -130,24 +139,7 @@ def __init__(
     # similar when they define their `on_result_found` method.
     # Right now this is just to demonstrate a particular use case.
     def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
-        if self.node_is_selected(original_node):
-            self.report_change(original_node)
-            if (attr := getattr(self, "on_result_found", None)) is not None:
-                # pylint: disable=not-callable
-                new_node = attr(original_node, updated_node)
-                return new_node
-
-        return updated_node
+        return self._new_or_updated_node(original_node, updated_node)
 
     def leave_Assign(self, original_node, updated_node):
-        pos_to_match = self.node_position(original_node)
-        if self.filter_by_result(
-            pos_to_match
-        ) and self.filter_by_path_includes_or_excludes(pos_to_match):
-            self.report_change(original_node)
-            if (attr := getattr(self, "on_result_found", None)) is not None:
-                # pylint: disable=not-callable
-                new_node = attr(original_node, updated_node)
-                return new_node
-
-        return updated_node
+        return self._new_or_updated_node(original_node, updated_node)

src/codemodder/codemods/api/helpers.py

@@ -47,18 +47,47 @@ 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,
+                new = self.make_new_arg(
+                    target_arg_name, target_arg_replacement_val, arg
                 )
+                arg_added = True
             else:
                 new = arg
             new_args.append(new)
+
+        if add_if_missing and not arg_added:
+            new = self.make_new_arg(target_arg_name, target_arg_replacement_val)
+            new_args.append(new)
         return new_args
+
+    def make_new_arg(self, name, value, existing_arg=None):
+        equal = (
+            existing_arg.equal
+            if existing_arg
+            else cst.AssignEqual(
+                whitespace_before=cst.SimpleWhitespace(""),
+                whitespace_after=cst.SimpleWhitespace(""),
+            )
+        )
+        return cst.Arg(
+            keyword=cst.parse_expression(name),
+            value=cst.parse_expression(value),
+            equal=equal,
+        )

src/codemodder/codemods/base_visitor.py

@@ -26,6 +26,8 @@ def filter_by_path_includes_or_excludes(self, pos_to_match):
         return True
 
     def node_is_selected(self, node) -> bool:
+        if not self.results:
+            return False
         pos_to_match = self.node_position(node)
         return self.filter_by_result(
             pos_to_match

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/docs/pixee_python_enable-jinja2-autoescape.md

@@ -0,0 +1,16 @@
+This codemod enables autoescaping of HTML content in `jinja2`. Unfortunately, the jinja2
+default behavior is to not autoescape when rendering templates, which makes your applications
+vulnerable to Cross-Site Scripting (XSS) attacks.
+
+Our codemod checks if you forgot to enable autoescape or if you explicitly disabled it. The change looks as follows:
+
+```diff
+  from jinja2 import Environment
+
+- env = Environment()
+- env = Environment(autoescape=False, loader=some_loader)
++ env = Environment(autoescape=True)
++ env = Environment(autoescape=True, loader=some_loader)
+  ...
+```
+

src/core_codemods/enable_jinja2_autoescape.py

@@ -0,0 +1,27 @@
+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: jinja2.Environment(...)
+                  - 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)

src/core_codemods/jwt_decode_verify.py

@@ -64,7 +64,13 @@ def replace_options_arg(self, node_args):
             new_args.append(new)
         return new_args
 
-    def replace_arg(self, original_node, target_arg_name, target_arg_replacement_val):
+    def replace_arg(
+        self,
+        original_node,
+        target_arg_name,
+        target_arg_replacement_val,
+        add_if_missing=False,
+    ):
         new_args = super().replace_arg(original_node, "verify", "True")
         return self.replace_options_arg(new_args)
 

tests/codemods/test_enable_jinja2_autoescape.py

@@ -0,0 +1,87 @@
+from core_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.NAME == "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)