Tests and docs for flask-enable-csrf-protection

integration_tests/base_test.py

@@ -204,7 +204,7 @@ def _lines_from_codepath(code_path):
 def _replace_lines_with(lines, replacements):
     total_lines = len(lines)
     for lineno, replacement in replacements:
-        if lineno > total_lines:
+        if lineno >= total_lines:
             lines.extend(replacement)
             continue
         lines[lineno] = replacement

integration_tests/test_flask_enable_csrf_protection.py

@@ -0,0 +1,36 @@
+from core_codemods.flask_enable_csrf_protection import FlaskEnableCSRFProtection
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestFlaskEnableCSRFProtection(BaseIntegrationTest):
+    codemod = FlaskEnableCSRFProtection
+    code_path = "tests/samples/flask_enable_csrf_protection.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path,
+        [
+            (1, """from flask_wtf.csrf import CSRFProtect\n"""),
+            (2, """\n"""),
+            (3, """app = Flask(__name__)\n"""),
+            (4, """csrf_app = CSRFProtect(app)\n"""),
+        ],
+    )
+
+    # fmt: off
+    expected_diff =(
+    """--- \n"""
+    """+++ \n"""
+    """@@ -1,3 +1,5 @@\n"""
+    """ from flask import Flask\n"""
+    """+from flask_wtf.csrf import CSRFProtect\n"""
+    """ \n"""
+    """ app = Flask(__name__)\n"""
+    """+csrf_app = CSRFProtect(app)\n"""
+    )
+    # fmt: on
+
+    expected_line_change = "3"
+    change_description = FlaskEnableCSRFProtection.CHANGE_DESCRIPTION
+    num_changed_files = 2

src/codemodder/scripts/generate_docs.py

@@ -202,6 +202,10 @@ class DocMetadata:
         importance="Low",
         guidance_explained="This change fixes deprecated uses and is safe.",
     ),
+    "flask-enable-csrf-protection": DocMetadata(
+        importance="High",
+        guidance_explained="Flask views may require proper handling of CSRF to function as expected and thus this change may break some views.",
+    ),
 }
 
 

src/core_codemods/docs/pixee_python_flask-enable-csrf-protection.md

@@ -0,0 +1,10 @@
+Flask views using `FlaskForm` have CSRF protection enabled by default. However other views may use AJAX to perform unsafe HTTP methods. FlaskWTF provides a way to enable CSRF protection globally for all views of a Flask app.
+Our changes look something like this:
+
+```diff
+from flask import Flask
++from flask_wtf.csrf import CSRFProtect
+
+app = Flask(__name__)
++csrf_app = CSRFProtect(app)
+```

tests/codemods/test_flask_enable_csrf_protection.py

@@ -0,0 +1,97 @@
+from core_codemods.flask_enable_csrf_protection import FlaskEnableCSRFProtection
+from tests.codemods.base_codemod_test import BaseCodemodTest
+from textwrap import dedent
+
+
+class TestFlaskEnableCSRFProtection(BaseCodemodTest):
+    codemod = FlaskEnableCSRFProtection
+
+    def test_name(self):
+        assert self.codemod.name() == "flask-enable-csrf-protection"
+
+    def test_simple(self, tmpdir):
+        input_code = """\
+        from flask import Flask
+        app = Flask(__name__)
+        """
+        expected = """\
+        from flask import Flask
+        from flask_wtf.csrf import CSRFProtect
+
+        app = Flask(__name__)
+        csrf_app = CSRFProtect(app)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_simple_alias(self, tmpdir):
+        input_code = """\
+        from flask import Flask as Flosk
+        app = Flosk(__name__)
+        """
+        expected = """\
+        from flask import Flask as Flosk
+        from flask_wtf.csrf import CSRFProtect
+
+        app = Flosk(__name__)
+        csrf_app = CSRFProtect(app)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_multiple(self, tmpdir):
+        input_code = """\
+        from flask import Flask
+        app = Flask(__name__)
+        app2 = Flask(__name__)
+        """
+        expected = """\
+        from flask import Flask
+        from flask_wtf.csrf import CSRFProtect
+
+        app = Flask(__name__)
+        csrf_app = CSRFProtect(app)
+        app2 = Flask(__name__)
+        csrf_app2 = CSRFProtect(app2)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 2
+
+    def test_multiple_inline(self, tmpdir):
+        input_code = """\
+        from flask import Flask
+        app = Flask(__name__); app2 = Flask(__name__)
+        """
+        expected = """\
+        from flask import Flask
+        from flask_wtf.csrf import CSRFProtect
+
+        app = Flask(__name__); app2 = Flask(__name__); csrf_app = CSRFProtect(app); csrf_app2 = CSRFProtect(app2)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_multiple_inline_suite(self, tmpdir):
+        input_code = """\
+        from flask import Flask
+        if True: app = Flask(__name__); app2 = Flask(__name__)
+        """
+        expected = """\
+        from flask import Flask
+        from flask_wtf.csrf import CSRFProtect
+
+        if True: app = Flask(__name__); app2 = Flask(__name__); csrf_app = CSRFProtect(app); csrf_app2 = CSRFProtect(app2)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_simple_protected(self, tmpdir):
+        input_code = """\
+        from flask import Flask
+        from flask_wtf.csrf import CSRFProtect
+
+        app = Flask(__name__)
+        csrf_app = CSRFProtect(app)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(input_code))
+        assert len(self.file_context.codemod_changes) == 0

tests/samples/flask_enable_csrf_protection.py

@@ -1,4 +1,3 @@
 from flask import Flask
-from flask_wtf.csrf import CSRFProtect
 
 app = Flask(__name__)