From 6509081c72a2f92c1500b3f09aa063441ea60031 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Sat, 1 May 2021 15:17:33 +0100
Subject: [PATCH] Support token auth with custom header in MultiAuth class
 (Fixes #125)

---
 flask_httpauth.py   | 33 ++++++++++++++++++---------------
 tests/test_multi.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/flask_httpauth.py b/flask_httpauth.py
index a1ff75b..e03ac8b 100644
--- a/flask_httpauth.py
+++ b/flask_httpauth.py
@@ -37,6 +37,18 @@ def default_auth_error(status):
         self.get_password(default_get_password)
         self.error_handler(default_auth_error)
 
+    def is_compatible_auth(self, headers):
+        if self.header is None or self.header == 'Authorization':
+            try:
+                scheme, _ = request.headers.get('Authorization', '').split(
+                    None, 1)
+            except ValueError:
+                # malformed Authorization header
+                return False
+            return scheme == self.scheme
+        else:
+            return self.header in headers
+
     def get_password(self, f):
         self.get_password_callback = f
         return f
@@ -368,21 +380,12 @@ def login_required(self, f=None, role=None, optional=None):
         def login_required_internal(f):
             @wraps(f)
             def decorated(*args, **kwargs):
-                selected_auth = None
-                if 'Authorization' in request.headers:
-                    try:
-                        scheme, creds = request.headers[
-                            'Authorization'].split(None, 1)
-                    except ValueError:
-                        # malformed Authorization header
-                        pass
-                    else:
-                        for auth in self.additional_auth:
-                            if auth.scheme == scheme:
-                                selected_auth = auth
-                                break
-                if selected_auth is None:
-                    selected_auth = self.main_auth
+                selected_auth = self.main_auth
+                if not self.main_auth.is_compatible_auth(request.headers):
+                    for auth in self.additional_auth:
+                        if auth.is_compatible_auth(request.headers):
+                            selected_auth = auth
+                            break
                 return selected_auth.login_required(role=role,
                                                     optional=optional
                                                     )(f)(*args, **kwargs)
diff --git a/tests/test_multi.py b/tests/test_multi.py
index f9d3253..7356ff7 100644
--- a/tests/test_multi.py
+++ b/tests/test_multi.py
@@ -11,7 +11,8 @@ def setUp(self):
 
         basic_auth = HTTPBasicAuth()
         token_auth = HTTPTokenAuth('MyToken')
-        multi_auth = MultiAuth(basic_auth, token_auth)
+        custom_token_auth = HTTPTokenAuth(header='X-Token')
+        multi_auth = MultiAuth(basic_auth, token_auth, custom_token_auth)
 
         @basic_auth.verify_password
         def verify_password(username, password):
@@ -37,6 +38,16 @@ def get_token_role(auth):
         def error_handler():
             return 'error', 401, {'WWW-Authenticate': 'MyToken realm="Foo"'}
 
+        @custom_token_auth.verify_token
+        def verify_custom_token(token):
+            return token == 'this-is-the-custom-token!'
+
+        @custom_token_auth.get_user_roles
+        def get_custom_token_role(auth):
+            if auth['token'] == 'this-is-the-custom-token!':
+                return 'foo'
+            return
+
         @app.route('/')
         def index():
             return 'index'
@@ -91,6 +102,19 @@ def test_multi_auth_login_invalid_token(self):
         self.assertEqual(response.headers['WWW-Authenticate'],
                          'MyToken realm="Foo"')
 
+    def test_multi_auth_login_valid_custom_token(self):
+        response = self.client.get(
+            '/protected', headers={'X-Token': 'this-is-the-custom-token!'})
+        self.assertEqual(response.data.decode('utf-8'), 'access granted:None')
+
+    def test_multi_auth_login_invalid_custom_token(self):
+        response = self.client.get(
+            '/protected', headers={'X-Token': 'this-is-not-the-token!'})
+        self.assertEqual(response.status_code, 401)
+        self.assertTrue('WWW-Authenticate' in response.headers)
+        self.assertEqual(response.headers['WWW-Authenticate'],
+                         'Bearer realm="Authentication Required"')
+
     def test_multi_auth_login_invalid_scheme(self):
         response = self.client.get(
             '/protected', headers={'Authorization': 'Foo this-is-the-token!'})
@@ -116,3 +140,9 @@ def test_multi_auth_login_valid_token_role(self):
             '/protected-with-role', headers={'Authorization':
                                              'MyToken this-is-the-token!'})
         self.assertEqual(response.data.decode('utf-8'), 'role access granted')
+
+    def test_multi_auth_login_valid_custom_token_role(self):
+        response = self.client.get(
+            '/protected-with-role', headers={'X-Token':
+                                             'this-is-the-custom-token!'})
+        self.assertEqual(response.data.decode('utf-8'), 'role access granted')