DRF with Cognito Authentication

...
...
...
# Application definition

INSTALLED_APPS = [
    ...
    ...
    "rest_framework",
]

LOCAL_APPS = [
    "articles",
]

INSTALLED_APPS += LOCAL_APPS

MIDDLEWARE = [
    ...
    "django.contrib.auth.middleware.AuthenticationMiddleware",
    "cognito.middleware.CognitoAuthMiddleware",
    ...
]
...
...
...
# AWS Cognito Configuration
COGNITO_CONFIG = {
    # URL consists of https://{pool-name}.auth.{region}.amazoncognito.com/oauth2/token
    # Replace it below accordingly
    "url": "https://frontenddrfcognito77b3bdee_userpool_77b3bdee-dev.auth.ap-southeast-2.amazoncognito.com/oauth2/token",
    "app_client_id": os.environ["AWS_APP_CLIENT_ID"],
    "region": "ap-southeast-2",
    "aws_user_pools_id": os.environ["AWS_USER_POOLS_ID"],
    "aws_user_pools_web_client_id": os.environ["AWS_APP_CLIENT_ID"],
}

# Rest Framework settings
REST_FRAMEWORK = {
    "DEFAULT_AUTHENTICATION_CLASSES": [
        "cognito.authentication.CognitoAuthentication",
    ],
}

# Default imports
from django.db import models

# Custom imports
# None


class Article(models.Model):
    name = models.CharField(max_length=50)

    def __str__(self) -> str:
        return self.name

# Default imports
from rest_framework import serializers

# Custom imports
from .models import Article


class ArticleSerializer(serializers.ModelSerializer):
    class Meta:
        model = Article
        fields = "__all__"

# Default imports
from rest_framework import generics

# Custom imports
from .models import Article
from .serializers import ArticleSerializer


class ListCreateArticles(generics.ListCreateAPIView):
    serializer_class = ArticleSerializer
    queryset = Article.objects.all()

# Default imports
from django.urls import path

# Custom imports
from .views import ListCreateArticles

urlpatterns = [
    path("", ListCreateArticles.as_view(), name="list_create_articles"),
]

from django.contrib import admin
from django.urls import path, include

urlpatterns = [
    path("admin/", admin.site.urls),
    path("articles/", include("articles.urls")),
]

...
# Other imports
...

class CognitoAuthMiddleware(
    CognitoAuthenticationMixin, middleware.AuthenticationMiddleware
):
    @staticmethod
    def get_auth_token(request):
        try:
            return request.META["HTTP_AUTHORIZATION"]
        except Exception:
            raise NoAuthToken()

    def process_request(self, request):
        if request.path.startswith(reverse("admin:index")):
            return None
        request.user = SimpleLazyObject(lambda: self.authenticate(request))

...
# Other imports
from . import keys
# Make sure to have the keys array in the __init__.py file
# in the cognito folder
# Standard Format:
# keys = [{...},{...}]

# In order to get the keys data, visit
# https://cognito-idp.{region}.amazonaws.com/{userPoolId}/.well-known/jwks.json
# Replace region and userPoolId with respective values
...

def get_jwt_claims(token):
    # get the kid from the headers prior to verification
    headers = jwt.get_unverified_headers(token)
    kid = headers["kid"]
    # search for the kid in the downloaded public keys
    key_index = -1
    for i in range(len(keys)):
        if kid == keys[i]["kid"]:
            key_index = i
            break
    if key_index == -1:
        print("Public key not found in jwks.json")
        return []
    # construct the public key
    public_key = jwk.construct(keys[key_index])
    # get the last two sections of the token,
    # message and signature (encoded in base64)
    message, encoded_signature = str(token).rsplit(".", 1)
    # decode the signature
    decoded_signature = base64url_decode(encoded_signature.encode("utf-8"))
    # verify the signature
    if not public_key.verify(message.encode("utf8"), decoded_signature):
        print("Signature verification failed")
        return []
    # print("Signature successfully verified")
    # since we passed the verification, we can now safely
    # use the unverified claims
    claims = jwt.get_unverified_claims(token)
    ts = claims["exp"]
    os.environ["TZ"] = "Asia/Kolkata"
    time.tzset()
    print(
        "Current Expiry of Token : {}".format(
            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ts))
        )
    )
    # Checking token expiry
    if time.time() > claims["exp"]:
        print("Token is expired")
        return []
    if claims["aud"] != settings.COGNITO_CONFIG["app_client_id"]:
        print("Token was not issued for this audience")
        return []
    return claims


class CognitoAuthenticationMixin:
    @staticmethod
    def get_auth_token(request):
        raise NotImplementedError()

    def authenticate(self, request):
        token = self.get_auth_token(request)
        try:
            claims = get_jwt_claims(token)
            if len(claims) > 0:
                user = UserModel.objects.get(email=claims["email"])
                return user
            raise NoSuchClaims()
        except UserModel.DoesNotExist:
            raise NoSuchUser()
        except Exception:
            raise InvalidAuthToken()


class CognitoAuthentication(
    CognitoAuthenticationMixin, authentication.BaseAuthentication
):
    @staticmethod
    def get_auth_token(request):
        try:
            return request.META["HTTP_AUTHORIZATION"]
        except Exception:
            raise NoAuthToken()

    def authenticate(self, request):
        user = super(CognitoAuthentication, self).authenticate(request)
        return user, None