From d7a79469e62e8cc98e2ccff0b6f53b2f43a0058e Mon Sep 17 00:00:00 2001
From: jorgectf <jorgectf@protonmail.com>
Date: Fri, 5 Nov 2021 20:08:52 +0100
Subject: [PATCH] Improve tests

---
 .../Security/CWE-614/django_bad.py            | 17 ++++++++++++--
 .../query-tests/Security/CWE-614/flask_bad.py | 23 +++++++++++++------
 .../Security/CWE-614/flask_good.py            |  7 ------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/python/ql/test/experimental/query-tests/Security/CWE-614/django_bad.py b/python/ql/test/experimental/query-tests/Security/CWE-614/django_bad.py
index 340291a6b9c..6f1916930e5 100644
--- a/python/ql/test/experimental/query-tests/Security/CWE-614/django_bad.py
+++ b/python/ql/test/experimental/query-tests/Security/CWE-614/django_bad.py
@@ -8,8 +8,21 @@ def django_response(request):
     return resp
 
 
+def django_response():
+    response = django.http.HttpResponse()
+    response['Set-Cookie'] = "name=value; SameSite=None;"
+    return response
+
+
 def django_response(request):
     resp = django.http.HttpResponse()
-    resp.set_cookie("name", "value", secure=False,
-                    httponly=False, samesite='None')
+    resp.set_cookie(django.http.request.GET.get("name"),
+                    django.http.request.GET.get("value"),
+                    secure=False, httponly=False, samesite='None')
     return resp
+
+
+def django_response():
+    response = django.http.HttpResponse()
+    response['Set-Cookie'] = f"{django.http.request.GET.get('name')}={django.http.request.GET.get('value')}; SameSite=None;"
+    return response
diff --git a/python/ql/test/experimental/query-tests/Security/CWE-614/flask_bad.py b/python/ql/test/experimental/query-tests/Security/CWE-614/flask_bad.py
index fc0177e3012..740070a7b53 100644
--- a/python/ql/test/experimental/query-tests/Security/CWE-614/flask_bad.py
+++ b/python/ql/test/experimental/query-tests/Security/CWE-614/flask_bad.py
@@ -3,6 +3,21 @@ from flask import Flask, request, make_response, Response
 app = Flask(__name__)
 
 
+@app.route("/false")
+def false():
+    resp = make_response()
+    resp.set_cookie("name", value="value", secure=False,
+                    httponly=False, samesite='None')
+    return resp
+
+
+@app.route("/flask_Response")
+def flask_Response():
+    resp = Response()
+    resp.headers['Set-Cookie'] = "name=value; SameSite=None;"
+    return resp
+
+
 @app.route("/false")
 def false():
     resp = make_response()
@@ -14,15 +29,9 @@ def false():
 @app.route("/flask_Response")
 def flask_Response():
     resp = Response()
-    resp.headers['Set-Cookie'] = "name=value;"
+    resp.headers['Set-Cookie'] = f"{request.args['name']}={request.args['value']}; SameSite=None;"
     return resp
 
 
-@app.route("/flask_make_response")
-def flask_make_response():
-    resp = make_response("hello")
-    resp.headers['Set-Cookie'] = "name=value; SameSite=None;"
-    return resp
-
 # if __name__ == "__main__":
 #     app.run(debug=True)
diff --git a/python/ql/test/experimental/query-tests/Security/CWE-614/flask_good.py b/python/ql/test/experimental/query-tests/Security/CWE-614/flask_good.py
index 5b9f83e1a63..724f8de8289 100644
--- a/python/ql/test/experimental/query-tests/Security/CWE-614/flask_good.py
+++ b/python/ql/test/experimental/query-tests/Security/CWE-614/flask_good.py
@@ -18,13 +18,6 @@ def flask_Response():
     return resp
 
 
-@app.route("/flask_make_response")
-def flask_make_response():
-    resp = make_response("hello")
-    resp.headers['Set-Cookie'] = "name=value; Secure; HttpOnly; SameSite=Lax;"
-    return resp
-
-
 def indeterminate(secure):
     resp = make_response()
     resp.set_cookie("name", value="value", secure=secure)