refactor: replace python-jose with pyjwt

api/pyproject.toml

@@ -9,7 +9,7 @@ license = ""
 cachetools = "^5.3.0"
 python = "^3.10"
 fastapi = "^0.101.0"
-python-jose = {extras = ["cryptography"], version = "^3.3.0"}
+pyjwt = "^2.8.0"
 uvicorn = {extras = ["standard"], version = "^0.21.1"}
 pymongo = "4.1.1"
 certifi = "^2023.7.22"

api/src/authentication/authentication.py

@@ -1,8 +1,8 @@
 import httpx
+import jwt
 from cachetools import TTLCache, cached
 from fastapi import Security
 from fastapi.security import OAuth2AuthorizationCodeBearer
-from jose import JWTError, jwt
 
 from authentication.mock_token_generator import mock_rsa_public_key
 from authentication.models import User
@@ -18,18 +18,12 @@
 
 
 @cached(cache=TTLCache(maxsize=32, ttl=86400))
-def fetch_openid_configuration() -> dict[str, str]:
+def fetch_openid_configuration() -> jwt.PyJWKClient:
     try:
         oid_conf_response = httpx.get(config.OAUTH_WELL_KNOWN)
         oid_conf_response.raise_for_status()
         oid_conf = oid_conf_response.json()
-        json_web_key_set_response = httpx.get(oid_conf["jwks_uri"])
-        json_web_key_set_response.raise_for_status()
-        return {
-            "authorization_endpoint": oid_conf["authorization_endpoint"],
-            "token_endpoint": oid_conf["token_endpoint"],
-            "jwks": json_web_key_set_response.json()["keys"],
-        }
+        return jwt.PyJWKClient(oid_conf["jwks_uri"])
     except Exception as error:
         logger.error(f"Failed to fetch OpenId Connect configuration for '{config.OAUTH_WELL_KNOWN}': {error}")
         raise credentials_exception
@@ -41,7 +35,11 @@ def auth_with_jwt(jwt_token: str = Security(oauth2_scheme)) -> User:
     if not jwt_token:
         raise credentials_exception
     # If TEST_TOKEN is true, we are running tests. Use the self-signed keys. If not, get keys from auth server.
-    key = mock_rsa_public_key if config.TEST_TOKEN else {"keys": fetch_openid_configuration()["jwks"]}
+    key = (
+        mock_rsa_public_key
+        if config.TEST_TOKEN
+        else fetch_openid_configuration().get_signing_key_from_jwt(jwt_token).key
+    )
 
     try:
         payload = jwt.decode(jwt_token, key, algorithms=["RS256"], audience=config.OAUTH_AUDIENCE)
@@ -50,7 +48,7 @@ def auth_with_jwt(jwt_token: str = Security(oauth2_scheme)) -> User:
             user = User(user_id=payload["oid"], **payload)
         else:
             user = User(user_id=payload["sub"], **payload)
-    except JWTError as error:
+    except jwt.exceptions.InvalidTokenError as error:
         logger.warning(f"Failed to decode JWT: {error}")
         raise credentials_exception
 

api/src/authentication/mock_token_generator.py

@@ -1,4 +1,4 @@
-from jose import jwt
+import jwt
 
 from authentication.models import User
 from config import default_user
@@ -63,5 +63,6 @@ def generate_mock_token(user: User = default_user) -> str:
         "sub": user.user_id,
         "roles": user.roles,
         "iss": "mock-auth-server",
+        "aud": "TEST",
     }
     return jwt.encode(payload, mock_rsa_private_key, algorithm="RS256")