Task/cdd 3145 validate cognito jwt (#3092)

* Add common/auth/cognito_jwt 
* Uses custom X-UHD-AUTH header for token

Author: mattjreynolds

common/__init__.py

@@ -0,0 +1 @@
+from .backend import JSONWebTokenAuthentication  # noqa

common/auth/__init__.py

@@ -0,0 +1,103 @@
+from django.apps import apps as django_apps
+from django.conf import settings
+from django.utils.encoding import force_str
+from django.utils.module_loading import import_string
+from django.utils.translation import gettext as _
+from rest_framework import HTTP_HEADER_ENCODING, exceptions
+from rest_framework.authentication import BaseAuthentication
+
+from .validator import TokenError, TokenValidator
+
+# 2 objects expected when parsing Auth Header: 'Bearer' + token
+VALID_AUTH_HEADER_LENGTH = 2
+
+
+def get_authorization_header(request):
+    """
+    Return request's 'X-UHD-AUTH:' header, as a bytestring.
+
+    Hide some test client ickyness where the header can be unicode.
+    """
+    auth = request.META.get("HTTP_X_UHD_AUTH", b"")
+    if isinstance(auth, str):
+        # Work around django test client oddness
+        auth = auth.encode(HTTP_HEADER_ENCODING)
+    return auth
+
+
+class JSONWebTokenAuthentication(BaseAuthentication):
+    """Token based authentication using the JSON Web Token standard.
+    Based on https://github.com/labd/django-cognito-jwt and modified
+    to suit our use case
+    """
+
+    def authenticate(self, request):
+        """Entrypoint for Django Rest Framework"""
+        jwt_token = self.get_jwt_token(request)
+        if jwt_token is None:
+            return None
+
+        # Authenticate token
+        try:
+            token_validator = self.get_token_validator(request)
+            jwt_payload = token_validator.validate(jwt_token)
+        except TokenError:
+            raise exceptions.AuthenticationFailed from None
+
+        custom_user_manager = self.get_custom_user_manager()
+        if custom_user_manager:
+            user = custom_user_manager.get_or_create_for_cognito(jwt_payload)
+        else:
+            user_model = self.get_user_model()
+            user = user_model.objects.get_or_create_for_cognito(jwt_payload)
+        return (user, jwt_token)
+
+    @staticmethod
+    def get_custom_user_manager():
+        """If COGNITO_USER_MANAGER is set, then the user object is obtained
+        via get_or_create_for_cognito on the user manager, this allows use
+        of the default unmodified Django User model"""
+        result = None
+        custom_user_manager_path = getattr(settings, "COGNITO_USER_MANAGER", False)
+        if custom_user_manager_path:
+            result = import_string(custom_user_manager_path)()
+        return result
+
+    @staticmethod
+    def get_user_model():
+        user_model = getattr(settings, "COGNITO_USER_MODEL", settings.AUTH_USER_MODEL)
+        return django_apps.get_model(user_model, require_ready=False)
+
+    @staticmethod
+    def get_jwt_token(request):
+        auth = get_authorization_header(request).split()
+        if not auth or force_str(auth[0].lower()) != "bearer":
+            return None
+
+        if len(auth) == 1:
+            msg = _("Invalid Authorization header. No credentials provided.")
+            raise exceptions.AuthenticationFailed(msg)
+        if len(auth) > VALID_AUTH_HEADER_LENGTH:
+            msg = _(
+                "Invalid Authorization header. Credentials string "
+                "should not contain spaces."
+            )
+            raise exceptions.AuthenticationFailed(msg)
+
+        return auth[1]
+
+    @staticmethod
+    def get_token_validator(request):
+        return TokenValidator(
+            settings.COGNITO_AWS_REGION,
+            settings.COGNITO_USER_POOL,
+            settings.COGNITO_AUDIENCE,
+        )
+
+    @staticmethod
+    def authenticate_header(request):
+        """
+        Method required by the DRF in order to return 401 responses for authentication failures, instead of 403.
+        More details in https://www.django-rest-framework.org/api-guide/authentication/#custom-authentication.
+        """
+        return "Bearer: api"

common/auth/cognito_jwt/__init__.py

@@ -0,0 +1,26 @@
+import logging
+
+from django.contrib.auth import get_user_model
+from django.contrib.auth.models import BaseUserManager, User
+
+logger = logging.getLogger(__name__)
+
+
+class CognitoManager(BaseUserManager):
+
+    @staticmethod
+    def get_or_create_for_cognito(jwt_payload):
+        username = jwt_payload["entraObjectId"]
+        try:
+            user = get_user_model().objects.get(username=username)
+            logger.debug("Found existing user %s", user.username)
+        except User.DoesNotExist:
+            password = None
+            user = get_user_model().objects.create_user(
+                username=username,
+                password=password,
+            )
+            logger.info("Created user %s", user.username)
+            user.is_active = True
+            user.save()
+        return user

common/auth/cognito_jwt/backend.py

@@ -0,0 +1,88 @@
+import json
+import logging
+
+import jwt
+import requests
+from django.conf import settings
+from django.core.cache import cache
+from django.utils.functional import cached_property
+from jwt.algorithms import RSAAlgorithm
+
+logger = logging.getLogger(__name__)
+
+
+class TokenError(Exception):
+    pass
+
+
+class TokenValidator:
+    def __init__(self, aws_region, aws_user_pool, audience):
+        self.aws_region = aws_region
+        self.aws_user_pool = aws_user_pool
+        self.audience = audience
+
+    @cached_property
+    def pool_url(self):
+        return (
+            f"https://cognito-idp.{self.aws_region}.amazonaws.com/{self.aws_user_pool}"
+        )
+
+    @cached_property
+    def _json_web_keys(self):
+        response = requests.get(self.pool_url + "/.well-known/jwks.json", timeout=10)
+        response.raise_for_status()
+        json_data = response.json()
+        return {item["kid"]: json.dumps(item) for item in json_data["keys"]}
+
+    def _get_public_key(self, token):
+        try:
+            headers = jwt.get_unverified_header(token)
+        except jwt.DecodeError as exc:
+            raise TokenError(str(exc)) from exc
+
+        if getattr(settings, "COGNITO_PUBLIC_KEYS_CACHING_ENABLED", False):
+            cache_key = "cognito_jwt:{}".format(headers["kid"])
+            jwk_data = cache.get(cache_key)
+
+            if not jwk_data:
+                jwk_data = self._json_web_keys.get(headers["kid"])
+                timeout = getattr(settings, "COGNITO_PUBLIC_KEYS_CACHING_TIMEOUT", 300)
+                cache.set(cache_key, jwk_data, timeout=timeout)
+        else:
+            jwk_data = self._json_web_keys.get(headers["kid"])
+
+        if jwk_data:
+            return RSAAlgorithm.from_jwk(jwk_data)
+        return None
+
+    def validate(self, token):
+        public_key = self._get_public_key(token)
+        if not public_key:
+            msg = "No key found for this token"
+            raise TokenError(msg)
+
+        params = {
+            "jwt": token,
+            "key": public_key,
+            "issuer": self.pool_url,
+            "algorithms": ["RS256"],
+        }
+
+        logger.debug("JWT - %s", params)
+        token_payload = jwt.decode(
+            token, options={"verify_signature": False}  # noqa: S5659
+        )
+        logger.debug("JWT decoded - %s", token_payload)
+
+        if "aud" in token_payload:
+            params.update({"audience": self.audience})
+
+        try:
+            jwt_data = jwt.decode(**params)
+        except (
+            jwt.InvalidTokenError,
+            jwt.ExpiredSignatureError,
+            jwt.DecodeError,
+        ) as exc:
+            raise TokenError(str(exc)) from exc
+        return jwt_data

common/auth/cognito_jwt/user_manager.py

@@ -61,6 +61,10 @@
 # The name of the AWS profile to use for the AWS client used for ingestion
 AWS_PROFILE_NAME = os.environ.get("AWS_PROFILE_NAME")
 
+# Cognito configuration
+COGNITO_AWS_REGION = os.environ.get("COGNITO_AWS_REGION")
+COGNITO_USER_POOL = os.environ.get("COGNITO_USER_POOL")
+
 # Database configuration
 POSTGRES_DB = os.environ.get("POSTGRES_DB")
 POSTGRES_USER = os.environ.get("POSTGRES_USER")

common/auth/cognito_jwt/validator.py

@@ -111,11 +111,18 @@
     },
 ]
 
+COGNITO_USER_MANAGER = "common.auth.cognito_jwt.user_manager.CognitoManager"
+COGNITO_AWS_REGION = config.COGNITO_AWS_REGION
+COGNITO_USER_POOL = config.COGNITO_USER_POOL
+COGNITO_AUDIENCE = None
+COGNITO_PUBLIC_KEYS_CACHING_ENABLED = True
+COGNITO_PUBLIC_KEYS_CACHING_TIMEOUT = 60 * 60 * 24  # 24h caching, default is 300s
 
 REST_FRAMEWORK = {
     "DEFAULT_SCHEMA_CLASS": "drf_spectacular.openapi.AutoSchema",
     "DEFAULT_AUTHENTICATION_CLASSES": [
         "rest_framework.authentication.SessionAuthentication",
+        "common.auth.cognito_jwt.JSONWebTokenAuthentication",
     ],
 }
 

config.py

@@ -20,6 +20,7 @@ pytest==9.0.2
 pytest-cov==7.1.0
 pytest-django==4.12.0
 pytest-random-order==1.2.0
+pytest-responses==0.5.1
 ruff==0.15.7
 stevedore==5.7.0
-django-debug-toolbar==6.2.0
+django-debug-toolbar==6.2.0

metrics/api/settings/default.py

@@ -10,6 +10,7 @@ click==8.3.1
 colorama==0.4.6
 coreapi==2.3.3
 coreschema==0.0.4
+cryptography==46.0.5
 defusedxml==0.7.1
 distlib==0.4.0
 django-cors-headers==4.8.0
@@ -58,6 +59,7 @@ Pillow==12.1.1
 platformdirs==4.9.4
 plotly==6.6.0
 pluggy==1.6.0
+pyjwt==2.12.1
 pyparsing==3.3.2
 pyrsistent==0.20.0
 python-dateutil==2.9.0.post0