chore: allow machine auth to generate CSRF token

Author: mapledan

superset/tasks/cache.py

@@ -291,11 +291,19 @@ def cache_warmup(
         return message
 
     user = security_manager.get_user_by_username(app.config["THUMBNAIL_SELENIUM_USER"])
-    cookies = MachineAuthProvider.get_auth_cookies(user)
-    headers = {
-        "Cookie": f"session={cookies.get('session', '')}",
-        "Content-Type": "application/json",
-    }
+    if app.config["WTF_CSRF_ENABLED"]:
+        cookies, csrf_token = MachineAuthProvider.get_auth_cookie_and_csrf_token(user)
+        headers = {
+            "Cookie": f"session={cookies.get('session', '')}",
+            "Content-Type": "application/json",
+            "X-CSRFToken": csrf_token,
+        }
+    else:
+        cookies = MachineAuthProvider.get_auth_cookies(user)
+        headers = {
+            "Cookie": f"session={cookies.get('session', '')}",
+            "Content-Type": "application/json",
+        }
 
     results: dict[str, list[str]] = {"scheduled": [], "errors": []}
     for payload in strategy.get_payloads():

superset/utils/machine_auth.py

@@ -23,6 +23,7 @@
 
 from flask import current_app, Flask, request, Response, session
 from flask_login import login_user
+from flask_wtf.csrf import generate_csrf
 from selenium.webdriver.remote.webdriver import WebDriver
 from werkzeug.http import parse_cookie
 
@@ -143,6 +144,31 @@ def get_auth_cookies(user: User) -> dict[str, str]:
 
         return cookies
 
+    @staticmethod
+    def get_auth_cookie_and_csrf_token(user: User) -> tuple[dict[str, str], str]:
+        # Login with the user specified to get the reports
+        with current_app.test_request_context("/login"):
+            login_user(user)
+            csrf_token = generate_csrf()
+            # A mock response object to get the cookie information from
+            response = Response()
+            current_app.session_interface.save_session(current_app, session, response)
+
+        cookies = {}
+
+        # Grab any "set-cookie" headers from the login response
+        for name, value in response.headers:
+            if name.lower() == "set-cookie":
+                # This yields a MultiDict, which is ordered -- something like
+                # MultiDict([('session', 'value-we-want), ('HttpOnly', ''), etc...
+                # Therefore, we just need to grab the first tuple and add it to our
+                # final dict
+                cookie = parse_cookie(value)
+                cookie_tuple = list(cookie.items())[0]
+                cookies[cookie_tuple[0]] = cookie_tuple[1]
+
+        return cookie, csrf_token
+
 
 class MachineAuthProviderFactory:
     def __init__(self) -> None:

tests/integration_tests/utils/machine_auth_tests.py

@@ -27,6 +27,15 @@ def test_get_auth_cookies(self):
         auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user)
         self.assertIsNotNone(auth_cookies["session"])
 
+    def test_get_auth_cookie_and_csrf_token(self):
+        user = self.get_user("admin")
+        (
+            auth_cookies,
+            csrf_token,
+        ) = machine_auth_provider_factory.instance.get_auth_cookie_and_csrf_token(user)
+        self.assertIsNotNone(auth_cookies["session"])
+        self.assertIsNotNone(csrf_token)
+
     @patch("superset.utils.machine_auth.MachineAuthProvider.get_auth_cookies")
     def test_auth_driver_user(self, get_auth_cookies):
         user = self.get_user("admin")