Flask set-cookie codemod (#101)

* flask secure cookie

* add testing and docs

* add flask req

* remove open

* allow for strict samesite

integration_tests/test_secure_flask_cookie.py

@@ -0,0 +1,22 @@
+from core_codemods.secure_flask_cookie import SecureFlaskCookie
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestSecureFlaskCookie(BaseIntegrationTest):
+    codemod = SecureFlaskCookie
+    code_path = "tests/samples/flask_cookie.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path,
+        [
+            (
+                7,
+                """    resp.set_cookie('custom_cookie', 'value', secure=True, httponly=True, samesite='Lax')\n""",
+            ),
+        ],
+    )
+    expected_diff = "--- \n+++ \n@@ -5,5 +5,5 @@\n @app.route('/')\n def index():\n     resp = make_response('Custom Cookie Set')\n-    resp.set_cookie('custom_cookie', 'value')\n+    resp.set_cookie('custom_cookie', 'value', secure=True, httponly=True, samesite='Lax')\n     return resp\n"
+    expected_line_change = "8"
+    change_description = SecureFlaskCookie.CHANGE_DESCRIPTION

requirements/test.txt

@@ -1,4 +1,5 @@
 coverage==7.3.*
+Flask<4
 GitPython<4
 Jinja2~=3.1.2
 lxml~=4.9.3

src/codemodder/scripts/generate_docs.py

@@ -79,6 +79,10 @@ class DocMetadata:
         importance="High",
         guidance_explained="There may be times when setting `verify=False` is useful for testing though we discourage it. \nYou may also decide to set `verify=/path/to/ca/bundle`. This codemod will not attempt to modify the `verify` value if you do set it to a path.",
     ),
+    "secure-flask-cookie": DocMetadata(
+        importance="Medium",
+        guidance_explained="Our change provides the most secure way to create cookies in Flask. However, it's possible you have configured your Flask application configurations to use secure cookies. In these cases, using the default parameters for `set_cookie` is safe.",
+    ),
     "secure-random": DocMetadata(
         importance="High",
         guidance_explained="While most of the functions in the `random` module aren't cryptographically secure, there are still valid use cases for `random.random()` such as for simulations or games.",

src/core_codemods/__init__.py

@@ -17,6 +17,7 @@
 from .remove_unnecessary_f_str import RemoveUnnecessaryFStr
 from .remove_unused_imports import RemoveUnusedImports
 from .requests_verify import RequestsVerify
+from .secure_flask_cookie import SecureFlaskCookie
 from .secure_random import SecureRandom
 from .tempfile_mktemp import TempfileMktemp
 from .upgrade_sslcontext_minimum_version import UpgradeSSLContextMinimumVersion
@@ -47,6 +48,7 @@
         RemoveUnnecessaryFStr,
         RemoveUnusedImports,
         RequestsVerify,
+        SecureFlaskCookie,
         SecureRandom,
         TempfileMktemp,
         UpgradeSSLContextMinimumVersion,

src/core_codemods/docs/pixee_python_secure-flask-cookie.md

@@ -0,0 +1,15 @@
+This codemod sets the most secure parameters when Flask applications call `set_cookie` on a response object. Without these parameters, your Flask
+application cookies may be vulnerable to being intercepted and used to gain access to sensitive data.
+
+The changes from this codemod look like this:
+
+```diff
+  from flask import Flask, session, make_response
+  app = Flask(__name__)
+  @app.route('/')
+    def index():
+      resp = make_response('Custom Cookie Set')
+    - resp.set_cookie('custom_cookie', 'value')
+    + resp.set_cookie('custom_cookie', 'value', secure=True, httponly=True, samesite='Lax')
+      return resp
+```

src/core_codemods/secure_flask_cookie.py

@@ -0,0 +1,74 @@
+from libcst import matchers
+from codemodder.codemods.base_codemod import ReviewGuidance
+from codemodder.codemods.api import SemgrepCodemod
+from codemodder.codemods.api.helpers import NewArg
+
+
+class SecureFlaskCookie(SemgrepCodemod):
+    NAME = "secure-flask-cookie"
+    REVIEW_GUIDANCE = ReviewGuidance.MERGE_AFTER_CURSORY_REVIEW
+    SUMMARY = "Use Safe Parameters in `flask` Response `set_cookie` Call"
+    DESCRIPTION = "Flask response `set_cookie` call should be called with `secure=True`, `httponly=True`, and `samesite='Lax'`."
+    REFERENCES = [
+        {
+            "url": "https://flask.palletsprojects.com/en/3.0.x/api/#flask.Response.set_cookie",
+            "description": "",
+        },
+        {
+            "url": "https://owasp.org/www-community/controls/SecureCookieAttribute",
+            "description": "",
+        },
+    ]
+
+    @classmethod
+    def rule(cls):
+        return """
+        rules:
+          - id: secure-flask-cookie
+            mode: taint
+            pattern-sources:
+              - pattern-either:
+                  - patterns:
+                    - pattern: flask.make_response(...)
+                    - pattern-inside: |
+                        import flask
+                        ...
+                  - patterns:
+                    - pattern: flask.Response(...)
+                    - pattern-inside: |
+                        import flask
+                        ...
+            pattern-sinks:
+              - patterns:
+                - pattern: $SINK.set_cookie(...)
+                - pattern-not: $SINK.set_cookie(..., secure=True, ..., httponly=True, ..., samesite="Lax", ...)
+                - pattern-not: $SINK.set_cookie(..., secure=True, ..., httponly=True, ..., samesite="Strict", ...)
+        """
+
+    def _choose_new_args(self, original_node):
+        new_args = [
+            NewArg(name="secure", value="True", add_if_missing=True),
+            NewArg(name="httponly", value="True", add_if_missing=True),
+        ]
+
+        samesite = matchers.Arg(
+            keyword=matchers.Name(value="samesite"),
+            value=matchers.SimpleString(value="'Strict'"),
+        )
+
+        # samesite=Strict is OK because it's more restrictive than Lax.
+        strict_samesite_defined = any(
+            matchers.matches(arg, samesite) for arg in original_node.args
+        )
+        if not strict_samesite_defined:
+            new_args.append(
+                NewArg(name="samesite", value="'Lax'", add_if_missing=True),
+            )
+
+        return new_args
+
+    def on_result_found(self, original_node, updated_node):
+        new_args = self.replace_args(
+            original_node, self._choose_new_args(original_node)
+        )
+        return self.update_arg_target(updated_node, new_args)

tests/codemods/test_secure_flask_cookie.py

@@ -0,0 +1,114 @@
+import pytest
+from core_codemods.secure_flask_cookie import SecureFlaskCookie
+from tests.codemods.base_codemod_test import BaseSemgrepCodemodTest
+from textwrap import dedent
+
+each_func = pytest.mark.parametrize("func", ["make_response", "Response"])
+
+
+class TestSecureFlaskCookie(BaseSemgrepCodemodTest):
+    codemod = SecureFlaskCookie
+
+    def test_name(self):
+        assert self.codemod.name() == "secure-flask-cookie"
+
+    @each_func
+    def test_import(self, tmpdir, func):
+        input_code = f"""\
+        import flask
+
+        response = flask.{func}()
+        var = "hello"
+        response.set_cookie("name", "value")
+        """
+        expexted_output = f"""\
+        import flask
+
+        response = flask.{func}()
+        var = "hello"
+        response.set_cookie("name", "value", secure=True, httponly=True, samesite='Lax')
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expexted_output))
+
+    @each_func
+    def test_from_import(self, tmpdir, func):
+        input_code = f"""\
+        from flask import {func}
+
+        response = {func}()
+        var = "hello"
+        response.set_cookie("name", "value")
+        """
+        expexted_output = f"""\
+        from flask import {func}
+
+        response = {func}()
+        var = "hello"
+        response.set_cookie("name", "value", secure=True, httponly=True, samesite='Lax')
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expexted_output))
+
+    @each_func
+    def test_import_alias(self, tmpdir, func):
+        input_code = f"""\
+        from flask import {func} as flask_resp
+        var = "hello"
+        flask_resp().set_cookie("name", "value")
+        """
+        expexted_output = f"""\
+        from flask import {func} as flask_resp
+        var = "hello"
+        flask_resp().set_cookie("name", "value", secure=True, httponly=True, samesite='Lax')
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expexted_output))
+
+    @pytest.mark.parametrize(
+        "input_args,expected_args",
+        [
+            (
+                "secure=True",
+                "secure=True, httponly=True, samesite='Lax'",
+            ),
+            (
+                "secure=True, httponly=True, samesite='Lax'",
+                "secure=True, httponly=True, samesite='Lax'",
+            ),
+            (
+                "secure=True, httponly=True, samesite='Strict'",
+                "secure=True, httponly=True, samesite='Strict'",
+            ),
+            (
+                "secure=False, httponly=True, samesite='Strict'",
+                "secure=True, httponly=True, samesite='Strict'",
+            ),
+            (
+                "httponly=True, samesite='Lax', secure=True",
+                "httponly=True, samesite='Lax', secure=True",
+            ),
+            (
+                "secure=True, httponly=True",
+                "secure=True, httponly=True, samesite='Lax'",
+            ),
+            (
+                "secure=True, httponly=True, samesite=None",
+                "secure=True, httponly=True, samesite='Lax'",
+            ),
+        ],
+    )
+    @each_func
+    def test_verify_variations(self, tmpdir, func, input_args, expected_args):
+        input_code = f"""\
+        import flask
+
+        response = flask.{func}()
+        var = "hello"
+        response.set_cookie("name", "value", {input_args})
+        """
+        expexted_output = f"""\
+        import flask
+
+        response = flask.{func}()
+        var = "hello"
+        response.set_cookie("name", "value", {expected_args})
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expexted_output))

tests/samples/flask_cookie.py

@@ -0,0 +1,9 @@
+from flask import Flask, session, make_response
+
+app = Flask(__name__)
+
+@app.route('/')
+def index():
+    resp = make_response('Custom Cookie Set')
+    resp.set_cookie('custom_cookie', 'value')
+    return resp