fix: Validate Azure JWTs using authlib

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
@@ -632,11 +633,9 @@ def get_oauth_user_info(self, provider, resp):
         # 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 +682,17 @@ def get_oauth_user_info(self, provider, resp):
         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(
+                "https://login.microsoftonline.com/common/discovery/keys"
+            ).json()
+        )
+        claims = jwt.decode(id_token, keyset)
+        claims.validate()
+        log.info("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

setup.py

@@ -46,6 +46,7 @@ def desc():
     platforms="any",
     install_requires=[
         "apispec[yaml]>=6.0.0, <7",
+        "authlib>=1.2.1",
         "colorama>=0.3.9, <1",
         "click>=8, <9",
         "email_validator>=1.0.5, <2",