Use Keycloak for JWT validation

src/mainframe/constants.py

@@ -34,9 +34,14 @@ class Mainframe(EnvConfig):
 
     job_timeout: int = 60 * 2
 
-
 mainframe_settings = Mainframe()  # pyright: ignore
 
+class Keycloak(EnvConfig, env_prefix="kc_"): # pyright: ignore
+    issuer_url: str = "https://keycloak.vipyrsec.com/realms/dragonfly"
+    audience: str = "Dragonfly Mainframe"
+    jwks_uri: str = "https://keycloak.vipyrsec.com/realms/dragonfly/protocol/openid-connect/certs"
+
+keycloak_settings = Keycloak()
 
 class _Sentry(EnvConfig, env_prefix="sentry_"):  # pyright: ignore
     dsn: str = ""

src/mainframe/dependencies.py

@@ -40,7 +40,6 @@ def validate_token_override():
         audience="DEVELOPMENT AUDIENCE",
         issued_at=datetime.now() - timedelta(seconds=10),
         expires_at=datetime.now() + timedelta(seconds=10),
-        grant_type="DEVELOPMENT GRANT TYPE",
     )
 
 

src/mainframe/json_web_token.py

@@ -4,21 +4,19 @@
 
 import jwt
 
-from mainframe.constants import mainframe_settings
+from mainframe.constants import keycloak_settings
 from mainframe.custom_exceptions import (
     BadCredentialsException,
     UnableCredentialsException,
 )
 
-
 @dataclass
 class AuthenticationData:
     issuer: str
     subject: str
     audience: str
     issued_at: datetime
     expires_at: datetime
-    grant_type: str
 
     @classmethod
     def from_dict(cls, data: dict[Any, Any]):
@@ -28,7 +26,6 @@ def from_dict(cls, data: dict[Any, Any]):
             audience=data["aud"],
             issued_at=datetime.fromtimestamp(data["iat"]),
             expires_at=datetime.fromtimestamp(data["exp"]),
-            grant_type=data["gty"],
         )
 
 
@@ -37,21 +34,18 @@ class JsonWebToken:
     """Perform JSON Web Token (JWT) validation using PyJWT"""
 
     jwt_access_token: str
-    auth0_issuer_url: str = f"https://{mainframe_settings.auth0_domain}/"
-    auth0_audience: str = mainframe_settings.auth0_audience
     algorithm: str = "RS256"
-    jwks_uri: str = f"{auth0_issuer_url}.well-known/jwks.json"
 
     def validate(self) -> AuthenticationData:
         try:
-            jwks_client = jwt.PyJWKClient(self.jwks_uri)
+            jwks_client = jwt.PyJWKClient(keycloak_settings.jwks_uri, headers={"User-Agent": "Dragonfly Mainframe"})
             jwt_signing_key = jwks_client.get_signing_key_from_jwt(self.jwt_access_token).key
             payload = jwt.decode(
                 self.jwt_access_token,
                 jwt_signing_key,
                 algorithms=self.algorithm,  # type: ignore
-                audience=self.auth0_audience,
-                issuer=self.auth0_issuer_url,
+                audience=keycloak_settings.audience,
+                issuer=keycloak_settings.issuer_url,
             )
         except jwt.exceptions.PyJWKClientError:
             raise UnableCredentialsException

tests/conftest.py

@@ -65,7 +65,6 @@ def auth() -> AuthenticationData:
         audience="DEVELOPMENT AUDIENCE",
         issued_at=datetime.now() - timedelta(seconds=10),
         expires_at=datetime.now() + timedelta(seconds=10),
-        grant_type="DEVELOPMENT GRANT TYPE",
     )