middleware.py

import jwt
from channels.db import database_sync_to_async
from django.conf import settings
from django.contrib.auth import get_user_model
from django.utils.translation import gettext_lazy as _
from rest_framework.exceptions import AuthenticationFailed

User = get_user_model()


class TokenAuthentication:
    """
    Simple token based authentication.

    Clients should authenticate by passing the token key in the query parameters.
    For example:

        ?token=401f7ac837da42b97f613d789819ff93537bee6a
    """

    model = None

    def get_model(self):
        if self.model is not None:
            return self.model
        from rest_framework.authtoken.models import Token

        return Token

    """
    A custom token model may be used, but must have the following properties.

    * key -- The string identifying the token
    * user -- The user to which the token belongs
    """

    def authenticate_credentials(self, key):
        model = self.get_model()
        try:
            token = model.objects.select_related("user").get(key=key)
        except model.DoesNotExist:
            raise AuthenticationFailed(_("Invalid token."))

        if not token.user.is_active:
            raise AuthenticationFailed(_("User inactive or deleted."))

        return token.user

    def authenticate(self, token):
        """
        Returns a `User` if a correct username and password have been supplied
        Args:
            token: token key

        Returns:
            User: A user instance.
        """
        try:
            user_id = jwt.decode(
                jwt=token, key=settings.SECRET_KEY, algorithms=["HS256"]
            )["user_id"]
        except jwt.exceptions.DecodeError:
            raise AuthenticationFailed(_("Invalid token."))
        except jwt.exceptions.ExpiredSignatureError:
            raise AuthenticationFailed(_("Token expired."))

        user = User.objects.get(pk=user_id)
        return user


@database_sync_to_async
def get_user(scope):
    """
    Return the user model instance associated with the given scope.
    If no user is retrieved, return an instance of `AnonymousUser`.
    """
    # postpone model import to avoid ImproperlyConfigured error before Django
    # setup is complete.
    from django.contrib.auth.models import AnonymousUser

    if "token" not in scope:
        raise ValueError(
            "Cannot find token in scope. You should wrap your consumer in "
            "TokenAuthMiddleware."
        )
    token = scope["token"]
    user = None
    try:
        auth = TokenAuthentication()
        user = auth.authenticate(token)
    except AuthenticationFailed:
        pass
    return user or AnonymousUser()


class TokenAuthMiddleware:
    """
    Custom middleware that takes a token from WebSocket subprotocols and authenticates via JWT.
    """

    SUBPROTOCOL_KEYWORD = "Bearer"

    def __init__(self, app):
        # Store the ASGI application we were passed
        self.app = app

    async def __call__(self, scope, receive, send):
        # Extract token from Sec-WebSocket-Protocol header.
        token = self._extract_token_from_subprotocol(scope.get("subprotocols", []))

        if token:
            scope["token"] = token
            scope["user"] = await get_user(scope)
        else:
            from django.contrib.auth.models import AnonymousUser

            scope["user"] = AnonymousUser()

        return await self.app(scope, receive, send)

    def _extract_token_from_subprotocol(self, subprotocols: list[str]) -> str | None:
        """
        Expect subprotocols in the form ["Bearer", "<JWT>"].
        """
        if not subprotocols:
            return None

        if len(subprotocols) >= 2 and subprotocols[0] == self.SUBPROTOCOL_KEYWORD:
            return subprotocols[1]

        return None