Introduce JWT based login/logout

- APIs & Tests

Author: RadoRado

config/django/base.py

@@ -170,6 +170,7 @@
 }
 
 from config.settings.cors import *  # noqa
+from config.settings.jwt import *  # noqa
 from config.settings.sessions import *  # noqa
 from config.settings.celery import *  # noqa
 from config.settings.sentry import *  # noqa

config/settings/jwt.py

@@ -0,0 +1,26 @@
+import datetime
+
+from config.env import env
+
+# For more settings
+# Read everything from here - https://styria-digital.github.io/django-rest-framework-jwt/#additional-settings
+
+# Default to 7 days
+JWT_EXPIRATION_DELTA_SECONDS = env("JWT_EXPIRATION_DELTA_SECONDS", default=60 * 60 * 24 * 7)
+JWT_AUTH_COOKIE = env("JWT_AUTH_COOKIE", default="jwt")
+JWT_AUTH_COOKIE_SAMESITE = env("JWT_AUTH_COOKIE_SAMESITE", default="Lax")
+JWT_AUTH_HEADER_PREFIX = env("JWT_AUTH_HEADER_PREFIX", default="Bearer")
+
+
+JWT_AUTH = {
+    "JWT_GET_USER_SECRET_KEY": "styleguide_example.authentication.services.auth_user_get_jwt_secret_key",
+    "JWT_RESPONSE_PAYLOAD_HANDLER": "styleguide_example.authentication.services.auth_jwt_response_payload_handler",
+    "JWT_EXPIRATION_DELTA": datetime.timedelta(seconds=JWT_EXPIRATION_DELTA_SECONDS),
+    "JWT_ALLOW_REFRESH": False,
+
+    "JWT_AUTH_COOKIE": JWT_AUTH_COOKIE,
+    "JWT_AUTH_COOKIE_SECURE": True,
+    "JWT_AUTH_COOKIE_SAMESITE": JWT_AUTH_COOKIE_SAMESITE,
+
+    "JWT_AUTH_HEADER_PREFIX": JWT_AUTH_HEADER_PREFIX
+}

styleguide_example/authentication/apis.py

@@ -1,12 +1,17 @@
 from django.contrib.auth import authenticate, login, logout
+from django.conf import settings
 
 from rest_framework.views import APIView
 from rest_framework.response import Response
 from rest_framework import serializers
 from rest_framework import status
 
+from rest_framework_jwt.views import ObtainJSONWebTokenView
+
 from styleguide_example.api.mixins import ApiAuthMixin
 
+from styleguide_example.authentication.services import auth_logout
+
 from styleguide_example.users.selectors import user_get_login_data
 
 
@@ -22,12 +27,10 @@ def post(self, request):
         serializer = self.InputSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
-        print(request.user)
         user = authenticate(request, **serializer.validated_data)
-        print(user)
 
         if user is None:
-            return Response(status=status.HTTP_401_UNAUTHORIZED)
+            return Response(status=status.HTTP_400_BAD_REQUEST)
 
         login(request, user)
 
@@ -52,6 +55,30 @@ def post(self, request):
         return Response()
 
 
+class UserJwtLoginApi(ObtainJSONWebTokenView):
+    def post(self, request, *args, **kwargs):
+        # We are redefining post so we can change the response status on success
+        # Mostly for consistency with the session-based API
+        response = super().post(request, *args, **kwargs)
+
+        if response.status_code == status.HTTP_201_CREATED:
+            response.status_code = status.HTTP_200_OK
+
+        return response
+
+
+class UserJwtLogoutApi(ApiAuthMixin, APIView):
+    def post(self, request):
+        auth_logout(request.user)
+
+        response = Response()
+
+        if settings.JWT_AUTH['JWT_AUTH_COOKIE'] is not None:
+            response.delete_cookie(settings.JWT_AUTH['JWT_AUTH_COOKIE'])
+
+        return response
+
+
 class UserMeApi(ApiAuthMixin, APIView):
     def get(self, request):
         data = user_get_login_data(user=request.user)

styleguide_example/authentication/services.py

@@ -0,0 +1,22 @@
+import uuid
+
+from styleguide_example.users.models import BaseUser
+
+
+def auth_user_get_jwt_secret_key(user: BaseUser) -> str:
+    return str(user.jwt_key)
+
+
+def auth_jwt_response_payload_handler(token, user=None, request=None, issued_at=None):
+    """
+    Default implementation. Add whatever suits you here.
+    """
+    return {"token": token}
+
+
+def auth_logout(user: BaseUser) -> BaseUser:
+    user.jwt_key = uuid.uuid4()
+    user.full_clean()
+    user.save(update_fields=["jwt_key"])
+
+    return user

styleguide_example/authentication/tests/apis/test_user_login.py

@@ -1,5 +1,6 @@
 from django.test import TestCase
 from django.urls import reverse
+from django.conf import settings
 
 from rest_framework.test import APIClient
 
@@ -25,7 +26,7 @@ def test_non_existing_user_cannot_login(self):
 
         response = self.client.post(self.session_login_url, data)
 
-        self.assertEqual(401, response.status_code)
+        self.assertEqual(400, response.status_code)
 
     def test_existing_user_can_login_and_access_apis(self):
         """
@@ -94,3 +95,100 @@ def test_existing_user_can_logout(self):
 
         response = self.client.get(self.me_url)
         self.assertEqual(403, response.status_code)
+
+
+class UserJwtLoginTests(TestCase):
+    def setUp(self):
+        self.client = APIClient()
+
+        self.jwt_login_url = reverse('api:authentication:jwt:login')
+        self.jwt_logout_url = reverse('api:authentication:jwt:logout')
+        self.me_url = reverse('api:authentication:me')
+
+    def test_non_existing_user_cannot_login(self):
+        self.assertEqual(0, BaseUser.objects.count())
+
+        data = {
+            'email': '[email protected]',
+            'password': 'hacksoft'
+        }
+
+        response = self.client.post(self.jwt_login_url, data)
+
+        self.assertEqual(400, response.status_code)
+
+    def test_existing_user_can_login_and_access_apis(self):
+        """
+        1. Create user
+        2. Assert login is OK
+        3. Call /api/auth/me
+        4. Assert valid response
+        """
+        credentials = {
+            "email": "[email protected]",
+            "password": "password"
+        }
+
+        user_create(
+            **credentials
+        )
+
+        response = self.client.post(self.jwt_login_url, credentials)
+
+        self.assertEqual(200, response.status_code)
+
+        data = response.data
+        self.assertIn("token", data)
+        token = data["token"]
+
+        jwt_cookie = response.cookies.get(settings.JWT_AUTH["JWT_AUTH_COOKIE"])
+
+        self.assertEqual(token, jwt_cookie.value)
+
+        response = self.client.get(self.me_url)
+        self.assertEqual(200, response.status_code)
+
+        # Now, try without session attached to the client
+        client = APIClient()
+
+        response = client.get(self.me_url)
+        self.assertEqual(403, response.status_code)
+
+        auth_headers = {
+            "HTTP_AUTHORIZATION": f"{settings.JWT_AUTH['JWT_AUTH_HEADER_PREFIX']} {token}"
+        }
+        response = client.get(self.me_url, **auth_headers)
+        self.assertEqual(200, response.status_code)
+
+    def test_existing_user_can_logout(self):
+        """
+        1. Create user
+        2. Login, can access APIs
+        3. Logout, cannot access APIs
+        """
+        credentials = {
+            "email": "[email protected]",
+            "password": "password"
+        }
+
+        user = user_create(
+            **credentials
+        )
+
+        key_before_logout = user.jwt_key
+
+        response = self.client.post(self.jwt_login_url, credentials)
+        self.assertEqual(200, response.status_code)
+
+        response = self.client.get(self.me_url)
+        self.assertEqual(200, response.status_code)
+
+        self.client.post(self.jwt_logout_url)
+
+        response = self.client.get(self.me_url)
+        self.assertEqual(403, response.status_code)
+
+        user.refresh_from_db()
+        key_after_logout = user.jwt_key
+
+        self.assertNotEqual(key_before_logout, key_after_logout)

styleguide_example/authentication/urls.py

@@ -3,6 +3,10 @@
 from .apis import (
     UserSessionLoginApi,
     UserSessionLogoutApi,
+
+    UserJwtLoginApi,
+    UserJwtLogoutApi,
+
     UserMeApi,
 )
 
@@ -23,6 +27,21 @@
 
         ], "session"))
     ),
+    path(
+        'jwt/',
+        include(([
+            path(
+                "login/",
+                UserJwtLoginApi.as_view(),
+                name="login"
+            ),
+            path(
+                "logout/",
+                UserJwtLogoutApi.as_view(),
+                name="logout"
+            )
+        ], "jwt"))
+    ),
     path(
         'me/',
         UserMeApi.as_view(),

styleguide_example/users/admin.py

@@ -14,9 +14,25 @@ class BaseUserAdmin(admin.ModelAdmin):
     list_filter = ('is_active', 'is_admin', 'is_superuser')
 
     fieldsets = (
-        (None, {'fields': ('email',)}),
+        (
+            None, {
+                'fields': ('email',)
+            }
+        ),
+        (
+            "Booleans", {
+                "fields": ("is_active", "is_admin", "is_superuser")
+            }
+        ),
+        (
+            "Timestamps", {
+                "fields": ("created_at", "updated_at")
+            }
+        )
     )
 
+    readonly_fields = ("created_at", "updated_at", )
+
     def save_model(self, request, obj, form, change):
         if change:
             return super().save_model(request, obj, form, change)

styleguide_example/users/migrations/0003_baseuser_jwt_key.py

@@ -0,0 +1,19 @@
+# Generated by Django 3.2.10 on 2021-12-20 14:24
+
+from django.db import migrations, models
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('users', '0002_alter_baseuser_id'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='baseuser',
+            name='jwt_key',
+            field=models.UUIDField(default=uuid.uuid4),
+        ),
+    ]

styleguide_example/users/models.py

@@ -1,3 +1,5 @@
+import uuid
+
 from django.db import models
 from django.contrib.auth.models import (
     BaseUserManager as BUM,
@@ -58,6 +60,9 @@ class BaseUser(BaseModel, AbstractBaseUser, PermissionsMixin):
     is_active = models.BooleanField(default=True)
     is_admin = models.BooleanField(default=False)
 
+    # This should potentially be an encrypted field
+    jwt_key = models.UUIDField(default=uuid.uuid4)
+
     objects = BaseUserManager()
 
     USERNAME_FIELD = 'email'