Handle login by auth managers (#32697)

airflow/auth/managers/base_auth_manager.py

@@ -18,9 +18,14 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException
 from airflow.utils.log.logging_mixin import LoggingMixin
 
+if TYPE_CHECKING:
+    from airflow.www.security import AirflowSecurityManager
+
 
 class BaseAuthManager(LoggingMixin):
     """
@@ -29,6 +34,9 @@ class BaseAuthManager(LoggingMixin):
     Auth managers are responsible for any user management related operation such as login, logout, authz, ...
     """
 
+    def __init__(self):
+        self._security_manager: AirflowSecurityManager | None = None
+
     @abstractmethod
     def get_user_name(self) -> str:
         """Return the username associated to the user in session."""
@@ -39,6 +47,11 @@ def is_logged_in(self) -> bool:
         """Return whether the user is logged in."""
         ...
 
+    @abstractmethod
+    def get_url_login(self, **kwargs) -> str:
+        """Return the login page url."""
+        ...
+
     def get_security_manager_override_class(self) -> type:
         """
         Return the security manager override class.
@@ -50,3 +63,19 @@ class airflow.www.security.AirflowSecurityManager with a custom implementation.
         By default, return an empty class.
         """
         return object
+
+    @property
+    def security_manager(self) -> AirflowSecurityManager:
+        """Get the security manager."""
+        if not self._security_manager:
+            raise AirflowException("Security manager not defined.")
+        return self._security_manager
+
+    @security_manager.setter
+    def security_manager(self, security_manager: AirflowSecurityManager):
+        """
+        Set the security manager.
+
+        :param security_manager: the security manager
+        """
+        self._security_manager = security_manager

airflow/auth/managers/fab/auth/__init__.py

@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/auth/managers/fab/auth/anonymous_user.py

@@ -0,0 +1,48 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from flask import current_app
+from flask_login import AnonymousUserMixin
+
+
+class AnonymousUser(AnonymousUserMixin):
+    """User object used when no active user is logged in."""
+
+    _roles: set[tuple[str, str]] = set()
+    _perms: set[tuple[str, str]] = set()
+
+    @property
+    def roles(self):
+        if not self._roles:
+            public_role = current_app.appbuilder.get_app.config["AUTH_ROLE_PUBLIC"]
+            self._roles = {current_app.appbuilder.sm.find_role(public_role)} if public_role else set()
+        return self._roles
+
+    @roles.setter
+    def roles(self, roles):
+        self._roles = roles
+        self._perms = set()
+
+    @property
+    def perms(self):
+        if not self._perms:
+            self._perms = {
+                (perm.action.name, perm.resource.name) for role in self.roles for perm in role.permissions
+            }
+        return self._perms

airflow/auth/managers/fab/fab_auth_manager.py

@@ -17,8 +17,10 @@
 # under the License.
 from __future__ import annotations
 
+from flask import url_for
 from flask_login import current_user
 
+from airflow import AirflowException
 from airflow.auth.managers.base_auth_manager import BaseAuthManager
 from airflow.auth.managers.fab.security_manager_override import FabAirflowSecurityManagerOverride
 
@@ -48,3 +50,12 @@ def is_logged_in(self) -> bool:
     def get_security_manager_override_class(self) -> type:
         """Return the security manager override."""
         return FabAirflowSecurityManagerOverride
+
+    def get_url_login(self, **kwargs) -> str:
+        """Return the login page url."""
+        if not self.security_manager.auth_view:
+            raise AirflowException("`auth_view` not defined in the security manager.")
+        if "next_url" in kwargs and kwargs["next_url"]:
+            return url_for(f"{self.security_manager.auth_view.endpoint}.login", next=kwargs["next_url"])
+        else:
+            return url_for(f"{self.security_manager.auth_view.endpoint}.login")

airflow/auth/managers/fab/security_manager_override.py

@@ -19,8 +19,14 @@
 
 from functools import cached_property
 
+from flask import g
 from flask_appbuilder.const import AUTH_DB, AUTH_LDAP, AUTH_OAUTH, AUTH_OID, AUTH_REMOTE_USER
 from flask_babel import lazy_gettext
+from flask_jwt_extended import JWTManager
+from flask_login import LoginManager
+from werkzeug.security import generate_password_hash
+
+from airflow.auth.managers.fab.auth.anonymous_user import AnonymousUser
 
 
 class FabAirflowSecurityManagerOverride:
@@ -47,6 +53,7 @@ class FabAirflowSecurityManagerOverride:
     :param resetmypasswordview: The class for reset my password view.
     :param resetpasswordview: The class for reset password view.
     :param rolemodelview: The class for role model view.
+    :param user_model: The user model.
     :param userinfoeditview: The class for user info edit view.
     :param userdbmodelview: The class for user db model view.
     :param userldapmodelview: The class for user ldap model view.
@@ -80,6 +87,7 @@ def __init__(self, **kwargs):
         self.resetmypasswordview = kwargs["resetmypasswordview"]
         self.resetpasswordview = kwargs["resetpasswordview"]
         self.rolemodelview = kwargs["rolemodelview"]
+        self.user_model = kwargs["user_model"]
         self.userinfoeditview = kwargs["userinfoeditview"]
         self.userdbmodelview = kwargs["userdbmodelview"]
         self.userldapmodelview = kwargs["userldapmodelview"]
@@ -88,6 +96,12 @@ def __init__(self, **kwargs):
         self.userremoteusermodelview = kwargs["userremoteusermodelview"]
         self.userstatschartview = kwargs["userstatschartview"]
 
+        # Setup Flask login
+        self.lm = self.create_login_manager()
+
+        # Setup Flask-Jwt-Extended
+        self.create_jwt_manager()
+
     def register_views(self):
         """Register FAB auth manager related views."""
         if not self.appbuilder.app.config.get("FAB_ADD_SECURITY_VIEWS", True):
@@ -192,6 +206,46 @@ def register_views(self):
                 category="Security",
             )
 
+    def create_login_manager(self) -> LoginManager:
+        """Create the login manager."""
+        lm = LoginManager(self.appbuilder.app)
+        lm.anonymous_user = AnonymousUser
+        lm.login_view = "login"
+        lm.user_loader(self.load_user)
+        return lm
+
+    def create_jwt_manager(self):
+        """Create the JWT manager."""
+        jwt_manager = JWTManager()
+        jwt_manager.init_app(self.appbuilder.app)
+        jwt_manager.user_lookup_loader(self.load_user_jwt)
+
+    def reset_password(self, userid, password):
+        """
+        Change/Reset a user's password for authdb.
+
+        Password will be hashed and saved.
+        :param userid: the user id to reset the password
+        :param password: the clear text password to reset and save hashed on the db
+        """
+        user = self.get_user_by_id(userid)
+        user.password = generate_password_hash(password)
+        self.update_user(user)
+
+    def load_user(self, user_id):
+        """Load user by ID."""
+        return self.get_user_by_id(int(user_id))
+
+    def load_user_jwt(self, _jwt_header, jwt_data):
+        identity = jwt_data["sub"]
+        user = self.load_user(identity)
+        # Set flask g.user to JWT user, we can't do it on before request
+        g.user = user
+        return user
+
+    def get_user_by_id(self, pk):
+        return self.appbuilder.get_session.get(self.user_model, pk)
+
     @property
     def auth_user_registration(self):
         """Will user self registration be allowed."""

airflow/www/auth.py

@@ -19,7 +19,7 @@
 from functools import wraps
 from typing import Callable, Sequence, TypeVar, cast
 
-from flask import current_app, flash, g, redirect, render_template, request, url_for
+from flask import current_app, flash, g, redirect, render_template, request
 
 from airflow.configuration import conf
 from airflow.utils.net import get_hostname
@@ -61,12 +61,7 @@ def decorated(*args, **kwargs):
             else:
                 access_denied = "Access is Denied"
                 flash(access_denied, "danger")
-            return redirect(
-                url_for(
-                    appbuilder.sm.auth_view.__class__.__name__ + ".login",
-                    next=request.url,
-                )
-            )
+            return redirect(get_auth_manager().get_url_login(next=request.url))
 
         return cast(T, decorated)
 

airflow/www/extensions/init_appbuilder.py

@@ -41,6 +41,7 @@
 
 from airflow import settings
 from airflow.configuration import conf
+from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 # This product contains a modified portion of 'Flask App Builder' developed by Daniel Vaz Gaspar.
 # (https://github.com/dpgaspar/Flask-AppBuilder).
@@ -212,6 +213,8 @@ def init_app(self, app, session):
         self._addon_managers = app.config["ADDON_MANAGERS"]
         self.session = session
         self.sm = self.security_manager_class(self)
+        auth_manager = get_auth_manager()
+        auth_manager.security_manager = self.sm
         self.bm = BabelManager(self)
         self._add_global_static()
         self._add_global_filters()
@@ -583,13 +586,7 @@ def security_converge(self, dry=False) -> dict:
         return self.sm.security_converge(self.baseviews, self.menu, dry)
 
     def get_url_for_login_with(self, next_url: str | None = None) -> str:
-        if self.sm.auth_view is None:
-            return ""
-        return url_for(f"{self.sm.auth_view.endpoint}.{'login'}", next=next_url)
-
-    @property
-    def get_url_for_login(self):
-        return url_for(f"{self.sm.auth_view.endpoint}.login")
+        return get_auth_manager().get_url_login(next_url=next_url)
 
     @property
     def get_url_for_logout(self):

airflow/www/extensions/init_auth_manager.py

@@ -16,11 +16,15 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.auth.managers.base_auth_manager import BaseAuthManager
+from typing import TYPE_CHECKING
+
 from airflow.compat.functools import cache
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException
 
+if TYPE_CHECKING:
+    from airflow.auth.managers.base_auth_manager import BaseAuthManager
+
 
 @cache
 def get_auth_manager() -> BaseAuthManager:

airflow/www/extensions/init_security.py

@@ -19,7 +19,7 @@
 import logging
 from importlib import import_module
 
-from flask import g, redirect, url_for
+from flask import g, redirect
 from flask_login import logout_user
 
 from airflow.configuration import conf
@@ -71,4 +71,4 @@ def init_check_user_active(app):
     def check_user_active():
         if get_auth_manager().is_logged_in() and not g.user.is_active:
             logout_user()
-            return redirect(url_for(app.appbuilder.sm.auth_view.endpoint + ".login"))
+            return redirect(get_auth_manager().get_url_login())