Python: Model CookieWrite for flask

RasmusWL · RasmusWL · commit 4606444b859d · 2021-06-24T17:34:43.000+02:00
diff --git a/python/ql/src/semmle/python/frameworks/Flask.qll b/python/ql/src/semmle/python/frameworks/Flask.qll
@@ -465,4 +465,39 @@ module Flask {
       result = "text/html"
     }
   }
+
+  // ---------------------------------------------------------------------------
+  // flask.Response related
+  // ---------------------------------------------------------------------------
+  /**
+   * A call to `set_cookie` on a Flask HTTP Response.
+   *
+   * See https://flask.palletsprojects.com/en/2.0.x/api/#flask.Response.set_cookie
+   */
+  class FlaskResponseSetCookieCall extends HTTP::Server::CookieWrite::Range,
+    DataFlow::MethodCallNode {
+    FlaskResponseSetCookieCall() { this.calls(Flask::Response::instance(), "set_cookie") }
+
+    override DataFlow::Node getHeaderArg() { none() }
+
+    override DataFlow::Node getNameArg() { result in [this.getArg(0), this.getArgByName("key")] }
+
+    override DataFlow::Node getValueArg() { result in [this.getArg(1), this.getArgByName("value")] }
+  }
+
+  /**
+   * A call to `delete_cookie` on a Flask HTTP Response.
+   *
+   * See https://flask.palletsprojects.com/en/2.0.x/api/#flask.Response.delete_cookie
+   */
+  class FlaskResponseDeleteCookieCall extends HTTP::Server::CookieWrite::Range,
+    DataFlow::MethodCallNode {
+    FlaskResponseDeleteCookieCall() { this.calls(Flask::Response::instance(), "delete_cookie") }
+
+    override DataFlow::Node getHeaderArg() { none() }
+
+    override DataFlow::Node getNameArg() { result in [this.getArg(0), this.getArgByName("key")] }
+
+    override DataFlow::Node getValueArg() { none() }
+  }
 }
diff --git a/python/ql/test/library-tests/frameworks/flask/response_test.py b/python/ql/test/library-tests/frameworks/flask/response_test.py
@@ -191,10 +191,11 @@ def redirect_simple():  # $requestHandler
 @app.route("/setting_cookie")  # $routeSetup="/setting_cookie"
 def setting_cookie():  # $requestHandler
     resp = make_response() # $ HttpResponse mimetype=text/html
-    resp.set_cookie("key", "value") # $ MISSING: CookieWrite CookieName="key" CookieValue="value"
-    resp.set_cookie(key="key", value="value") # $ MISSING: CookieWrite CookieName="key" CookieValue="value"
+    resp.set_cookie("key", "value") # $ CookieWrite CookieName="key" CookieValue="value"
+    resp.set_cookie(key="key", value="value") # $ CookieWrite CookieName="key" CookieValue="value"
     resp.headers.add("Set-Cookie", "key2=value2") # $ MISSING: CookieWrite CookieRawHeader="key2=value2"
-    resp.delete_cookie("key3") # $ MISSING: CookieWrite CookieName="key3"
+    resp.delete_cookie("key3") # $ CookieWrite CookieName="key3"
+    resp.delete_cookie(key="key3") # $ CookieWrite CookieName="key3"
     return resp  # $ SPURIOUS: HttpResponse mimetype=text/html responseBody=resp
 
 ################################################################################
