Initial serial verification code, need to test

Author: jonathanmelitski

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.11-buster
+FROM python:3.11-bullseye
 
 LABEL maintainer="Penn Labs"
 

src/auth.py

@@ -9,19 +9,26 @@
 JWKS_URL = settings.JWKS_URL
 
 
-def get_jwk():
+def get_jwks():
     if settings.JWKS_CACHE:
-        key = settings.JWKS_CACHE
-        return key
+        # Check to make sure we have a cached JWK for each key, otherwise refetch all.
+        missing = False
+        for key in settings.JWKS_URL.keys():
+            if key not in settings.JWKS_CACHE:
+                missing = True
+        
+        if not missing:
+            return settings.JWKS_CACHE
 
     # Make a request to get the JWKS
-    try:
-        response = requests.get(JWKS_URL)
-        jwks = jwk.JWKSet.from_json(response.text)
-        settings.JWKS_CACHE = jwks
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
+    for key in settings.JWKS_URL:
+        try:
+            response = requests.get(JWKS_URL)
+            jwks = jwk.JWKSet.from_json(response.text)
+            settings.JWKS_CACHE[key] = jwks
+        except Exception as e:
+            del settings.JWKS_CACHE[key]
+            raise HTTPException(status_code=500, detail=str(e))
     return settings.JWKS_CACHE
 
 
@@ -41,11 +48,12 @@ def get_token_from_header(request: Request):
 
 
 def verify_jwt(token: str = Depends(get_token_from_header)):
-    try:
-        # Load the public key
-        public_key = get_jwk()
-        # Decode and verify the JWT
-        decoded_token = jwt.JWT(key=public_key, jwt=token)
-        return decoded_token.claims
-    except Exception as e:
-        raise HTTPException(status_code=401, detail=str(e))
+    public_keys = get_jwks()
+    for key in public_keys:
+        try:
+            decoded_token = jwt.JWT(key=key, jwt=token)
+            return decoded_token.claims
+        except:
+            pass
+    
+    raise HTTPException(status_code=401, detail="Failed to verify JWT token against public keys array.")

src/config.py

@@ -11,8 +11,11 @@ class Config(BaseSettings):
     DATABASE_URL: PostgresDsn
     REDIS_URL: RedisDsn
 
-    JWKS_CACHE: JWKSet | None = None
-    JWKS_URL: str = "https://platform.pennlabs.org/identity/jwks/"
+    JWKS_CACHE: list[JWKSet] | None = None
+    JWKS_URL: dict[str, str] = {
+        "b2b":"https://platform.pennlabs.org/identity/jwks/",
+        "user":"https://platform.pennlabs.org/accounts/.well-known/jwks.json"
+    }
 
     SITE_DOMAIN: str = "analytics.pennlabs.org"