allow for strict samesite

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

@@ -1,5 +1,5 @@
 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 against your system.
+application cookies may be vulnerable to being intercepted and used to gain access to sensitive data.
 
 The changes from this codemod look like this:
 

src/core_codemods/secure_flask_cookie.py

@@ -1,3 +1,5 @@
+from libcst import matchers
+import libcst as cst
 from codemodder.codemods.base_codemod import ReviewGuidance
 from codemodder.codemods.api import SemgrepCodemod
 from codemodder.codemods.api.helpers import NewArg
@@ -41,15 +43,33 @@ def rule(cls):
               - 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,
-            [
-                NewArg(name="secure", value="True", add_if_missing=True),
-                NewArg(name="httponly", value="True", add_if_missing=True),
-                NewArg(name="samesite", value="'Lax'", add_if_missing=True),
-            ],
+            original_node, self._choose_new_args(original_node)
         )
         return self.update_arg_target(updated_node, new_args)

tests/codemods/test_secure_flask_cookie.py

@@ -75,7 +75,11 @@ def test_import_alias(self, tmpdir, func):
             ),
             (
                 "secure=True, httponly=True, samesite='Strict'",
-                "secure=True, httponly=True, samesite='Lax'",
+                "secure=True, httponly=True, samesite='Strict'",
+            ),
+            (
+                "secure=False, httponly=True, samesite='Strict'",
+                "secure=True, httponly=True, samesite='Strict'",
             ),
             (
                 "httponly=True, samesite='Lax', secure=True",