fix: Validate Azure JWTs using authlib

flask_appbuilder/const.py

@@ -191,3 +191,9 @@
 API_ADD_TITLE_RIS_KEY = "add_title"
 API_EDIT_TITLE_RIS_KEY = "edit_title"
 API_SHOW_TITLE_RIS_KEY = "show_title"
+
+# -----------------------------------
+#  OAuth Provider Constants
+# -----------------------------------
+
+MICROSOFT_KEY_SET_URL = "https://login.microsoftonline.com/common/discovery/keys"

flask_appbuilder/security/manager.py

@@ -1,17 +1,18 @@
-import base64
 import datetime
 import json
 import logging
 import re
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
+from authlib.jose import JsonWebKey, jwt
 from flask import Flask, g, session, url_for
 from flask_babel import lazy_gettext as _
 from flask_jwt_extended import current_user as current_user_jwt
 from flask_jwt_extended import JWTManager
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
 from flask_login import current_user, LoginManager
+import requests
 from werkzeug.security import check_password_hash, generate_password_hash
 
 from .api import SecurityApi
@@ -54,6 +55,7 @@
     LOGMSG_WAR_SEC_LOGIN_FAILED,
     LOGMSG_WAR_SEC_NO_USER,
     LOGMSG_WAR_SEC_NOLDAP_OBJ,
+    MICROSOFT_KEY_SET_URL,
     PERMISSION_PREFIX,
 )
 
@@ -632,11 +634,9 @@
         # https://docs.microsoft.com/en-us/azure/active-directory/develop/
         # active-directory-protocols-oauth-code
         if provider == "azure":
-            log.debug("Azure response received : %s", resp)
-            id_token = resp["id_token"]
-            log.debug(str(id_token))
-            me = self._azure_jwt_token_parse(id_token)
-            log.debug("Parse JWT token : %s", me)
+            log.debug("Azure response received:\n%s", json.dumps(resp, indent=4))
+            me = self._decode_and_validate_azure_jwt(resp["id_token"])
+            log.debug("Decoded JWT:\n%s", json.dumps(me, indent=4))
             return {
                 "name": me.get("name", ""),
                 "email": me["upn"],
@@ -683,36 +683,13 @@
         else:
             return {}
 
-    def _azure_parse_jwt(self, id_token):
-        jwt_token_parts = r"^([^\.\s]*)\.([^\.\s]+)\.([^\.\s]*)$"
-        matches = re.search(jwt_token_parts, id_token)
-        if not matches or len(matches.groups()) < 3:
-            log.error("Unable to parse token.")
-            return {}
-        return {
-            "header": matches.group(1),
-            "Payload": matches.group(2),
-            "Sig": matches.group(3),
-        }
-
-    def _azure_jwt_token_parse(self, id_token):
-        jwt_split_token = self._azure_parse_jwt(id_token)
-        if not jwt_split_token:
-            return
-
-        jwt_payload = jwt_split_token["Payload"]
-        # Prepare for base64 decoding
-        payload_b64_string = jwt_payload
-        payload_b64_string += "=" * (4 - ((len(jwt_payload) % 4)))
-        decoded_payload = base64.urlsafe_b64decode(payload_b64_string.encode("ascii"))
-
-        if not decoded_payload:
-            log.error("Payload of id_token could not be base64 url decoded.")
-            return
-
-        jwt_decoded_payload = json.loads(decoded_payload.decode("utf-8"))
+    def _decode_and_validate_azure_jwt(self, id_token):
+        keyset = JsonWebKey.import_key_set(requests.get(MICROSOFT_KEY_SET_URL).json())
+        claims = jwt.decode(id_token, keyset)
+        claims.validate()
+        log.debug("Decoded JWT:\n%s", json.dumps(claims, indent=4))
 
-        return jwt_decoded_payload
+        return claims
 
     def register_views(self):
         if not self.appbuilder.app.config.get("FAB_ADD_SECURITY_VIEWS", True):

requirements.txt

@@ -6,6 +6,8 @@
 #
 apispec[yaml]==6.3.0
     # via Flask-AppBuilder (setup.py)
+authlib==1.2.1
+    # via Flask-AppBuilder (setup.py)
 attrs==21.4.0
     # via jsonschema
 babel==2.9.1

tests/security/test_base_security_manager.py

@@ -1,8 +1,12 @@
 import datetime
+import json
 import unittest
 from unittest.mock import MagicMock, patch
 
 from flask_appbuilder.security.manager import BaseSecurityManager
+from flask_appbuilder.security.manager import JsonWebKey, jwt
+
+JWTClaimsMock = MagicMock()
 
 
 @patch.object(BaseSecurityManager, "update_user")
@@ -67,3 +71,12 @@ def test_subsequent_unsuccessful_auth(self, mock1, mock2):
         self.assertEqual(user_mock.fail_login_count, 10)
         self.assertEqual(user_mock.last_login, None)
         self.assertTrue(bsm.update_user.called_once)
+
+    @patch.object(JsonWebKey, "import_key_set", MagicMock())
+    @patch.object(jwt, "decode", MagicMock(return_value=JWTClaimsMock))
+    @patch.object(json, "dumps", MagicMock(return_value="DecodedExampleAzureJWT"))
+    def test_azure_jwt_validated(self, mock1, mock2):
+        bsm = BaseSecurityManager()
+
+        bsm._decode_and_validate_azure_jwt("ExampleAzureJWT")
+        JWTClaimsMock.validate.assert_called()