From 3030e881d2e44f4021764e18e489fe940a9b3636 Mon Sep 17 00:00:00 2001
From: Daniel Vaz Gaspar <danielvazgaspar@gmail.com>
Date: Wed, 21 Aug 2024 19:00:21 +0100
Subject: [PATCH] feat: add no cache directive to login forms (#2266)

* feat: add no cache directive to login forms

* add test
---
 flask_appbuilder/security/decorators.py | 14 ++++++++++++++
 flask_appbuilder/security/views.py      |  5 ++++-
 tests/security/test_mvc_security.py     | 13 +++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/flask_appbuilder/security/decorators.py b/flask_appbuilder/security/decorators.py
index 740580b08c..46287e7235 100644
--- a/flask_appbuilder/security/decorators.py
+++ b/flask_appbuilder/security/decorators.py
@@ -30,6 +30,20 @@
 P = ParamSpec("P")
 
 
+def no_cache(view: Callable[..., Response]) -> Callable[..., Response]:
+    @functools.wraps(view)
+    def wrapped_view(*args, **kwargs) -> Response:
+        response = make_response(view(*args, **kwargs))
+        response.headers[
+            "Cache-Control"
+        ] = "no-store, no-cache, must-revalidate, max-age=0"
+        response.headers["Pragma"] = "no-cache"
+        response.headers["Expires"] = "0"
+        return response
+
+    return wrapped_view
+
+
 def response_unauthorized_mvc(status_code: int) -> Response:
     response = make_response(
         jsonify({"message": str(FLAMSG_ERR_SEC_ACCESS_DENIED), "severity": "danger"}),
diff --git a/flask_appbuilder/security/views.py b/flask_appbuilder/security/views.py
index 86fa1021fe..ddc1d0b5e8 100644
--- a/flask_appbuilder/security/views.py
+++ b/flask_appbuilder/security/views.py
@@ -9,7 +9,7 @@
 from flask_appbuilder.baseviews import BaseView
 from flask_appbuilder.charts.views import DirectByChartView
 from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget
-from flask_appbuilder.security.decorators import has_access
+from flask_appbuilder.security.decorators import has_access, no_cache
 from flask_appbuilder.security.forms import (
     DynamicForm,
     LoginForm_db,
@@ -520,6 +520,7 @@ class AuthDBView(AuthView):
     login_template = "appbuilder/general/security/login_db.html"
 
     @expose("/login/", methods=["GET", "POST"])
+    @no_cache
     def login(self):
         if g.user is not None and g.user.is_authenticated:
             return redirect(self.appbuilder.get_url_for_index)
@@ -543,6 +544,7 @@ class AuthLDAPView(AuthView):
     login_template = "appbuilder/general/security/login_ldap.html"
 
     @expose("/login/", methods=["GET", "POST"])
+    @no_cache
     def login(self):
         if g.user is not None and g.user.is_authenticated:
             return redirect(self.appbuilder.get_url_for_index)
@@ -568,6 +570,7 @@ class AuthOIDView(AuthView):
     oid_ask_for_optional: List[str] = []
 
     @expose("/login/", methods=["GET", "POST"])
+    @no_cache
     def login(self, flag=True) -> WerkzeugResponse:
         @self.appbuilder.sm.oid.loginhandler
         def login_handler(self):
diff --git a/tests/security/test_mvc_security.py b/tests/security/test_mvc_security.py
index d7adea63fc..bb022fc911 100644
--- a/tests/security/test_mvc_security.py
+++ b/tests/security/test_mvc_security.py
@@ -63,6 +63,19 @@ class Model1View(ModelView):
 
         self.appbuilder.add_view(Model1View, "Model1", category="Model1")
 
+    def test_sec_login_no_cache(self):
+        """
+        Test Security Login, no cache directives
+        """
+        rv = self.client.get("/login/")
+        assert rv.status_code == 200
+        assert (
+            rv.headers.get("Cache-Control")
+            == "no-store, no-cache, must-revalidate, max-age=0"
+        )
+        assert rv.headers["Pragma"] == "no-cache"
+        assert rv.headers["Expires"] == "0"
+
     def test_sec_login(self):
         """
         Test Security Login, Logout, invalid login, invalid access