Add cognito user manager, use env vars

Add optional use of a user manager to hold the get_or_create function
Use env vars for aws region, aws user pool

Author: mattjreynolds

config.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")

metrics/api/django_cognito_jwt/__init__.py

@@ -1,3 +1 @@
-__version__ = "0.0.4"
-
 from .backend import JSONWebTokenAuthentication  # noqa

metrics/api/django_cognito_jwt/backend.py

@@ -1,33 +1,34 @@
-import logging
-
 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
 
-logger = logging.getLogger(__name__)
 VALID_AUTH_HEADER_LENGTH = 2
 
 
 def get_authorization_header(request):
     """
-    Return request's 'Authorization:' header, as a bytestring.
+    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'')
+    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."""
+    """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"""
@@ -42,10 +43,25 @@ def authenticate(self, request):
         except TokenError:
             raise exceptions.AuthenticationFailed from None
 
-        user_model = self.get_user_model()
-        user = user_model.objects.get_or_create_for_cognito(jwt_payload)
+        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)

metrics/api/django_cognito_jwt/user_manager.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

metrics/api/django_cognito_jwt/validator.py

@@ -1,4 +1,5 @@
 import json
+import logging
 
 import jwt
 import requests
@@ -7,6 +8,8 @@
 from django.utils.functional import cached_property
 from jwt.algorithms import RSAAlgorithm
 
+logger = logging.getLogger(__name__)
+
 
 class TokenError(Exception):
     pass
@@ -62,10 +65,13 @@ def validate(self, token):
             "jwt": token,
             "key": public_key,
             "issuer": self.pool_url,
-            "algorithms": ["RS256"]
+            "algorithms": ["RS256"],
         }
 
+        logger.debug("JWT - %s", params)
         token_payload = jwt.decode(token, options={"verify_signature": False})
+        logger.debug("JWT decoded - %s", token_payload)
+
         if "aud" in token_payload:
             params.update({"audience": self.audience})
 

metrics/api/settings/default.py

@@ -111,8 +111,9 @@
     },
 ]
 
-COGNITO_AWS_REGION = "eu-west-2"
-COGNITO_USER_POOL = "eu-west-2_m44aoyGeK"
+COGNITO_USER_MANAGER = "metrics.api.django_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

tests/unit/metrics/api/django_cognito_jwt/test_backend.py

@@ -17,30 +17,42 @@ def test_authenticate_no_token(rf):
     assert auth.authenticate(request) is None
 
 
+@pytest.mark.parametrize(
+    "cognito_user_manager",
+    ["metrics.api.django_cognito_jwt.user_manager.CognitoManager", None],
+)
 def test_authenticate_valid(
-    rf, monkeypatch, cognito_well_known_keys, jwk_private_key_one
+    rf, monkeypatch, cognito_well_known_keys, jwk_private_key_one, cognito_user_manager
 ):
+    settings.COGNITO_USER_MANAGER = cognito_user_manager
     token = create_jwt_token(
         jwk_private_key_one,
         {
             "iss": "https://cognito-idp.eu-central-1.amazonaws.com/bla",
             "aud": settings.COGNITO_AUDIENCE,
             "sub": "username",
+            "entraObjectId": "entraOID",
         },
     )
 
+    @staticmethod
     def func(payload):
-        return USER_MODEL(username=payload["sub"])
+        return USER_MODEL(username=payload["entraObjectId"])
 
-    monkeypatch.setattr(
-        USER_MODEL.objects, "get_or_create_for_cognito", func, raising=False
-    )
+    if cognito_user_manager:
+        monkeypatch.setattr(
+            f"{cognito_user_manager}.get_or_create_for_cognito", func, raising=False
+        )
+    else:
+        monkeypatch.setattr(
+            USER_MODEL.objects, "get_or_create_for_cognito", func, raising=False
+        )
 
     request = rf.get("/", HTTP_X_UHD_AUTH=b"bearer %s" % token.encode("utf8"))
     auth = backend.JSONWebTokenAuthentication()
     user, auth_token = auth.authenticate(request)
     assert user
-    assert user.username == "username"
+    assert user.username == "entraOID"
     assert auth_token == token.encode("utf8")
 
 

tests/unit/metrics/api/django_cognito_jwt/test_user_manager.py

@@ -0,0 +1,50 @@
+from unittest import mock
+from django.contrib.auth import get_user_model
+from django.contrib.auth.models import User
+
+from metrics.api.django_cognito_jwt.user_manager import CognitoManager
+
+USER_MODEL = get_user_model()
+
+
+def test_get_or_create_for_cognito_get_existing_user():
+    jwt_payload = {
+        "entraObjectId": "unique_user_id",
+    }
+
+    mock_user = mock.MagicMock()
+    mock_user.username = jwt_payload["entraObjectId"]
+    mock_user.is_active = True
+
+    with mock.patch.object(USER_MODEL.objects, "get", return_value=mock_user):
+        user = CognitoManager.get_or_create_for_cognito(jwt_payload)
+        assert user
+        assert user.username == jwt_payload["entraObjectId"]
+        assert user.is_active is True
+
+
+def test_get_or_create_for_cognito_create_user():
+    jwt_payload = {
+        "entraObjectId": "unique_user_id",
+    }
+
+    def create_user_mock(*args, username, **kwargs):
+        mock_user = mock.MagicMock()
+        mock_user.username = username
+        return mock_user
+
+    with (
+        mock.patch.object(USER_MODEL.objects, "get", side_effect=User.DoesNotExist),
+        mock.patch.object(
+            USER_MODEL.objects, "create_user", side_effect=create_user_mock
+        ) as create_user,
+    ):
+        user = CognitoManager.get_or_create_for_cognito(jwt_payload)
+        assert user
+        assert user.username == jwt_payload["entraObjectId"]
+        assert user.is_active is True
+
+        create_user.assert_called_once_with(
+            username=jwt_payload["entraObjectId"],
+            password=None,
+        )