From af8c3a484c1bc4056b9f52086d3905eff5e05541 Mon Sep 17 00:00:00 2001
From: imacat <imacat@mail.imacat.idv.tw>
Date: Fri, 25 Nov 2022 09:14:57 +1100
Subject: [PATCH] Revised so that you always call digest_auth.init_app(), to
 avoid confusion.  It remembers the current application.  The logout() method
 no longer need current_app for the current application.

---
 README.rst                    |  5 +++++
 src/flask_digest_auth/auth.py | 20 +++++++-------------
 tests/test_auth.py            |  1 +
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/README.rst b/README.rst
index ca7864f..6c63239 100644
--- a/README.rst
+++ b/README.rst
@@ -75,6 +75,7 @@ In your ``my_app.py``:
     ... (Configure the Flask application) ...
 
     auth: DigestAuth = DigestAuth(realm="Admin")
+    auth.init_app(app)
 
     @auth.register_get_password
     def get_password_hash(username: str) -> t.Optional[str]:
@@ -113,6 +114,7 @@ In your ``my_app/__init__.py``:
         ... (Configure the Flask application) ...
 
         auth.realm = app.config["REALM"]
+        auth.init_app(app)
 
         @auth.register_get_password
         def get_password_hash(username: str) -> t.Optional[str]:
@@ -156,6 +158,9 @@ module that requires log in, without specifying the authentication
 mechanism.  The Flask application can specify the actual
 authentication mechanism as it sees fit.
 
+``login_manager.init_app(app)`` must be called before
+``auth.init_app(app)``.
+
 
 Example for Simple Applications with Flask-Login Integration
 ------------------------------------------------------------
diff --git a/src/flask_digest_auth/auth.py b/src/flask_digest_auth/auth.py
index df94654..bb53c8a 100644
--- a/src/flask_digest_auth/auth.py
+++ b/src/flask_digest_auth/auth.py
@@ -27,8 +27,7 @@ from functools import wraps
 from random import random
 from secrets import token_urlsafe
 
-from flask import g, request, Response, session, abort, Flask, Request, \
-    current_app
+from flask import g, request, Response, session, abort, Flask, Request
 from itsdangerous import URLSafeTimedSerializer, BadData
 from werkzeug.datastructures import Authorization
 
@@ -55,6 +54,7 @@ class DigestAuth:
         self.__get_password_hash: t.Callable[[str], t.Optional[str]] \
             = lambda x: None
         self.__get_user: t.Callable[[str], t.Optional] = lambda x: None
+        self.app: t.Optional[Flask] = None
 
     def login_required(self, view) -> t.Callable:
         """The view decorator for HTTP digest authentication.
@@ -205,13 +205,12 @@ class DigestAuth:
         :param app: The Flask application.
         :return: None.
         """
+        app.digest_auth = self
+        self.app = app
 
-        try:
+        if hasattr(app, "login_manager"):
             from flask_login import LoginManager, login_user
 
-            if not hasattr(app, "login_manager"):
-                raise AttributeError(
-                    "Please run the Flask-Login init-app() first")
             login_manager: LoginManager = getattr(app, "login_manager")
 
             @login_manager.unauthorized_handler
@@ -252,12 +251,7 @@ class DigestAuth:
                         app.logger.warning(str(e))
                     return None
 
-        except ModuleNotFoundError:
-            raise ModuleNotFoundError(
-                "init_app() is only for Flask-Login integration")
-
-    @staticmethod
-    def logout() -> None:
+    def logout(self) -> None:
         """Logs out the user.
         This actually causes the next authentication to fail, which forces
         the browser to ask the user for the username and password again.
@@ -267,7 +261,7 @@ class DigestAuth:
         if "user" in session:
             del session["user"]
         try:
-            if hasattr(current_app, "login_manager"):
+            if hasattr(self.app, "login_manager"):
                 from flask_login import logout_user
                 logout_user()
         except ModuleNotFoundError:
diff --git a/tests/test_auth.py b/tests/test_auth.py
index 5467048..fcb6751 100644
--- a/tests/test_auth.py
+++ b/tests/test_auth.py
@@ -49,6 +49,7 @@ class AuthenticationTestCase(TestCase):
         app.test_client_class = Client
 
         auth: DigestAuth = DigestAuth(realm=_REALM)
+        auth.init_app(app)
         user_db: t.Dict[str, str] \
             = {_USERNAME: make_password_hash(_REALM, _USERNAME, _PASSWORD)}