Django Rest Framework custom JWT authentication
Tags: django, jwt, security, djangorestframework
Published: 2020-05-12 15:27:00

> This article is not a tutorial or a guide, it is more like a request for code review and validate the implementation from more experienced Django developers, so please don't use this code unless you are able to review and validate yourself 

## Introduction
Web application security is critical and doing it on your own in production is a huge responsibility, that is why I love Django because it takes care of most of the security vulnerabilities out there, but my problem shows up when I wanted to build my back-end as a restful API so I can connect all type of clients (SPA, mobile, desktop, ... etc) 
Django Rest Framework (DRF) comes with different builtin [authentication classes](https://www.django-rest-framework.org/tutorial/4-authentication-and-permissions/), token authentication or [JWT](https://github.com/SimpleJWT/django-rest-framework-simplejwt) are the way to go for my use case but I still have that worry of how to save the tokens in client-side
everybody says don't save the token in localstorage because of [XSS attacks](https://en.wikipedia.org/wiki/Cross-site_scripting) and better to save your token in httponly cookie, but cookies are open to [CSRF attack](https://en.wikipedia.org/wiki/Cross-site_request_forgery) too and DRF disable CSRF protection for all the APIView so what is the best practice to do this.
for a while I kept using it as everybody especially in mobile (i do react-native) we have secure storages and also if the device is not (rooted android or jailbroken ios) then each app is sandboxed and the tokens are safe in most cases
The problem still exists in web clients (SPAs), so I came into an implementation that might be useful and I wanted to document it here so I get feedback from more experienced developers, I can summarize my implementation into the following steps:
- the user sends a POST request with username and password to log in, then the server will do 3 things
  - generate an `access_token` which is a short life jwt (maybe 5 mins) and send it in the response body
  - generate a `refresh_token` which is a long life jwt (days) and send it in an httponly cookie, so it won't be accessible from the client javascript
  - send a normal cookie that contains a CSRF token

- The developer needs to be sure that all unsafe views (POST, UPDATE, PUT, Delete) are protected by the builtin Django CSRF protection because as I mentioned above DRF disables them by default.

- in the client-side the developer should take care of:
  - In the client-side, every request will contain refresh token in the cookies automatically (be sure that your client domain is whitelisted in the server cors headers settings)
  - Send the `access_token` in the `Authorization` header.
  - Send the CSRF token in the `X-CSRFTOKEN` header if he is doing POST request
  - when he needs a new `access_token`, he need to send a POST request to refresh token endpoint
 
enough talking, let us see some code

## Project Setup 
```bash
python3 -m venv .venv
source .venv/bin/activate
pip install django django-cors-headers djangorestframework PyJWT

# create Django project with the venv activates
django-admin startproject project

# create an app 
./manage.py start app accounts
```

in the project settings enable the apps and add some settings
```python
INSTALLED_APPS = [
    ...
    # 3rd party apps
    'corsheaders',
    'rest_framework',
    
    # project apps
    'accounts',
]

CORS_ALLOW_CREDENTIALS = True # to accept cookies via ajax request
CORS_ORIGIN_WHITELIST = [
    'http://localhost:3000' # the domain for front-end app(you can add more than 1) 
]

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated', # make all endpoints private
    )
}
```

before applying the first migration and as recommended in the [official Django docs](https://docs.djangoproject.com/en/3.0/topics/auth/customizing/#using-a-custom-user-model-when-starting-a-project) I like to start my project with a custom user model even if I won't use it now
in `accounts.models` define the user model 

```python
from django.contrib.auth.models import AbstractUser

class User(AbstractUser):
    pass
```
in the `project.settings.py` define the user model

```python
...
AUTH_USER_MODEL = 'accounts.User'
...
```
later you can refer to the user model by one of the following 

```python
from django.conf import settings
User = settings.AUTH_USER_MODEL

# OR

from django.contrib.auth import get_user_model
User = get_user_model()
```


and in `accounts.admin` register the new user model to the admin site

```python
from django.contrib import admin
from django.contrib.auth.admin import UserAdmin
from accounts.models import User

admin.site.register(User, UserAdmin)
```

create a super user 
```
./manage.py createsuperuser
```

and finally, an endpoint to test, as we declare in the above settings all the endpoint will require authentication by default, we can override this in certain views as we will going to do in the login later

we will create the user profile endpoint that will return current authenticated user object as JSON, for that we will need to create a user serializer

```python
# accounts.serializers
from rest_framework import serializers
from accounts.models import User


class UserSerializer(serializers.ModelSerializer):
    class Meta:
        model = User
        fields = ['id', 'username', 'email',
                  'first_name', 'last_name', 'is_active']

```

```python
# project.urls
from accounts import urls as accounts_urls

urlpatterns = [
    path('accounts/', include(accounts_urls)),
]

# accounts.urls
urlpatterns = [
    path('profile', profile, name='profile'),
]

# accounts.views
from rest_framework.decorators import api_view
from rest_framework.response import Response
from .serializers import UserSerializer


@api_view(['GET'])
def profile(request):
    user = request.user
    serialized_user = UserSerializer(user).data
    return Response({'user': serialized_user })
```
now if you tried to acceess this endpoint you will get 403 error, as described in the introduction we need to login then send the access_token in the request header

## Login View
The login endpoint will be a post request with `username` and `password` in the request body.
we will make the login view public using the permission class decorator `AllowAny`, also we will decorate the view with `@ensure_csrf_cookie` forcing Django to send the CSRF cookie in the response if the login success

if the login success, we will have:
- an `access_token` in the response body.
- a `refreshtoken` in an httponly cookie.
- a `csrftoken` in a normal cookie so we can read it from javascript and resend it when needed

```python
# accounts.views
from django.contrib.auth import get_user_model
from rest_framework.response import Response
from rest_framework import exceptions
from rest_framework.permissions import AllowAny
from rest_framework.decorators import api_view, permission_classes
from django.views.decorators.csrf import ensure_csrf_cookie
from accounts.serializers import UserSerializer
from accounts.utils import generate_access_token, generate_refresh_token


@api_view(['POST'])
@permission_classes([AllowAny])
@ensure_csrf_cookie
def login_view(request):
    User = get_user_model()
    username = request.data.get('username')
    password = request.data.get('password')
    response = Response()
    if (username is None) or (password is None):
        raise exceptions.AuthenticationFailed(
            'username and password required')

    user = User.objects.filter(username=username).first()
    if(user is None):
        raise exceptions.AuthenticationFailed('user not found')
    if (not user.check_password(password)):
        raise exceptions.AuthenticationFailed('wrong password')

    serialized_user = UserSerializer(user).data

    access_token = generate_access_token(user)
    refresh_token = generate_refresh_token(user)

    response.set_cookie(key='refreshtoken', value=refresh_token, httponly=True)
    response.data = {
        'access_token': access_token,
        'user': serialized_user,
    }

    return response

```

and here are the functions that generate the tokens, notice that I use a different secret to sign refresh token for more security 

```python
# accounts.utils
import datetime
import jwt
from django.conf import settings


def generate_access_token(user):

    access_token_payload = {
        'user_id': user.id,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=0, minutes=5),
        'iat': datetime.datetime.utcnow(),
    }
    access_token = jwt.encode(access_token_payload,
                              settings.SECRET_KEY, algorithm='HS256').decode('utf-8')
    return access_token


def generate_refresh_token(user):
    refresh_token_payload = {
        'user_id': user.id,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
        'iat': datetime.datetime.utcnow()
    }
    refresh_token = jwt.encode(
        refresh_token_payload, settings.REFRESH_TOKEN_SECRET, algorithm='HS256').decode('utf-8')

    return refresh_token

```

![Login Response Body](https://i.imgur.com/nC2HM5t.png)

![Login Response Cookies](https://i.imgur.com/kNnrFdE.png)

## Custom Authentication Class for DRF
Django Rest Framework makes it easy to create a custom authentication scheme, it described in details in the [official docs](https://www.django-rest-framework.org/api-guide/authentication/#custom-authentication)
The following code is originally taken from DRF source code then I add my changes as required.
notice that DRF enforce CSRF only in the session authentication [rest_framework/authentication.py](https://github.com/encode/django-rest-framework/blob/master/rest_framework/authentication.py)
```python
# accounts.authentication

import jwt
from rest_framework.authentication import BaseAuthentication
from django.middleware.csrf import CsrfViewMiddleware
from rest_framework import exceptions
from django.conf import settings
from django.contrib.auth import get_user_model


class CSRFCheck(CsrfViewMiddleware):
    def _reject(self, request, reason):
        # Return the failure reason instead of an HttpResponse
        return reason


class SafeJWTAuthentication(BaseAuthentication):
    '''
        custom authentication class for DRF and JWT
        https://github.com/encode/django-rest-framework/blob/master/rest_framework/authentication.py
    '''

    def authenticate(self, request):

        User = get_user_model()
        authorization_heaader = request.headers.get('Authorization')

        if not authorization_heaader:
            return None
        try:
            # header = 'Token xxxxxxxxxxxxxxxxxxxxxxxx'
            access_token = authorization_heaader.split(' ')[1]
            payload = jwt.decode(
                access_token, settings.SECRET_KEY, algorithms=['HS256'])

        except jwt.ExpiredSignatureError:
            raise exceptions.AuthenticationFailed('access_token expired')
        except IndexError:
            raise exceptions.AuthenticationFailed('Token prefix missing')

        user = User.objects.filter(id=payload['user_id']).first()
        if user is None:
            raise exceptions.AuthenticationFailed('User not found')

        if not user.is_active:
            raise exceptions.AuthenticationFailed('user is inactive')

        self.enforce_csrf(request)
        return (user, None)

    def enforce_csrf(self, request):
        """
        Enforce CSRF validation
        """
        check = CSRFCheck()
        # populates request.META['CSRF_COOKIE'], which is used in process_view()
        check.process_request(request)
        reason = check.process_view(request, None, (), {})
        print(reason)
        if reason:
            # CSRF failed, bail with explicit error message
            raise exceptions.PermissionDenied('CSRF Failed: %s' % reason)

```

once you create that class, go to the `project.settings` and activate it in the `REST_FRAMEWORK` section like this `appname.filename.classname`

```python

REST_FRAMEWORK = {
    'DEFAULT_AUTHENTICATION_CLASSES': (
        'accounts.authentication.SafeJWTAuthentication',
    ),
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated',
    )
}

```
Now as we have `access_token` and set the authentication method we can revisit the `profile` endpoint but this time we will set the `Authorization` header

![Profile endpoint response](https://i.imgur.com/rOrsafD.png)


## Refresh Token View
Whenever the token is expired or you need a new token for any reason we need a `refresh_token` endpoint.
this view needs the permission of `AlloAny` because we don't have an `access_token` but it will be protected by 2 other things
- a valid `refresh_token` that sent in the httoponly cookie.
- the CSRF token so we are sure the above cookie not compromised if the 2 conditions are satisfied then the server will generate a new valid `access_token` and send it back.

if the `refresh_token` is invalid or expired, the user will need to re-logion 

```python
import jwt
from django.conf import settings
from django.contrib.auth import get_user_model
from django.views.decorators.csrf import csrf_protect
from rest_framework import exceptions
from rest_framework.response import Response
from rest_framework.permissions import AllowAny
from rest_framework.decorators import api_view, permission_classes
from accounts.utils import generate_access_token


@api_view(['POST'])
@permission_classes([AllowAny])
@csrf_protect
def refresh_token_view(request):
    '''
    To obtain a new access_token this view expects 2 important things:
        1. a cookie that contains a valid refresh_token
        2. a header 'X-CSRFTOKEN' with a valid csrf token, client app can get it from cookies "csrftoken"
    '''
    User = get_user_model()
    refresh_token = request.COOKIES.get('refreshtoken')
    if refresh_token is None:
        raise exceptions.AuthenticationFailed(
            'Authentication credentials were not provided.')
    try:
        payload = jwt.decode(
            refresh_token, settings.REFRESH_TOKEN_SECRET, algorithms=['HS256'])
    except jwt.ExpiredSignatureError:
        raise exceptions.AuthenticationFailed(
            'expired refresh token, please login again.')

    user = User.objects.filter(id=payload.get('user_id')).first()
    if user is None:
        raise exceptions.AuthenticationFailed('User not found')

    if not user.is_active:
        raise exceptions.AuthenticationFailed('user is inactive')


    access_token = generate_access_token(user)
    return Response({'access_token': access_token})

```

## Token Revoke
The last piece of the puzzle is a way to revoke the refresh token which has a long lifetime,  you might blacklist the token or assign a uuid for the token and put it in the payload then link it to the user and save it in the database, when revoking or logout you just change that uuid in the database to not match the value in the payload, you can pick what suits your application need