Merge pull request #99 from jazzband/alichass-jwt-custom-claims

Adds support for JWT Custom Claims Serializer

dj_rest_auth/tests/test_api.py

@@ -18,6 +18,19 @@ try:
 except ImportError:
     from django.core.urlresolvers import reverse
 
+from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
+from jwt import decode as decode_jwt
+
+class TESTTokenObtainPairSerializer(TokenObtainPairSerializer):
+    @classmethod
+    def get_token(cls, user):
+        token = super().get_token(user)
+        # Add custom claims
+        token['name'] = user.username
+        token['email'] = user.email
+
+        return token
+
 
 @override_settings(ROOT_URLCONF="tests.urls")
 class APIBasicTests(TestsMixin, TestCase):
@@ -605,3 +618,54 @@ class APIBasicTests(TestsMixin, TestCase):
         # test other TokenError, AttributeError, TypeError (invalid format)
         resp = self.post(self.logout_url, status=200, data=json.dumps({'refresh': token}))
         self.assertEqual(resp.status_code, 500)
+
+
+    @override_settings(REST_USE_JWT=True)
+    @override_settings(JWT_AUTH_COOKIE=None)
+    @override_settings(REST_FRAMEWORK=dict(
+        DEFAULT_AUTHENTICATION_CLASSES=[
+            'dj_rest_auth.utils.JWTCookieAuthentication'
+        ]
+    ))
+    @override_settings(REST_SESSION_LOGIN=False)
+    @override_settings(JWT_TOKEN_CLAIMS_SERIALIZER = 'tests.test_api.TESTTokenObtainPairSerializer')
+    def test_custom_jwt_claims(self):
+        payload = {
+            "username": self.USERNAME,
+            "password": self.PASS
+        }
+        get_user_model().objects.create_user(self.USERNAME, self.EMAIL, self.PASS)
+
+        self.post(self.login_url, data=payload, status_code=200)
+        self.assertEqual('access_token' in self.response.json.keys(), True)
+        self.token = self.response.json['access_token']
+        claims = decode_jwt(self.token, settings.SECRET_KEY, algorithms='HS256')
+        self.assertEquals(claims['user_id'], 1)
+        self.assertEquals(claims['name'], 'person')
+        self.assertEquals(claims['email'], 'person1@world.com')
+
+
+    @override_settings(REST_USE_JWT=True)
+    @override_settings(JWT_AUTH_COOKIE='jwt-auth')
+    @override_settings(REST_FRAMEWORK=dict(
+        DEFAULT_AUTHENTICATION_CLASSES=[
+            'dj_rest_auth.utils.JWTCookieAuthentication'
+        ]
+    ))
+    @override_settings(REST_SESSION_LOGIN=False)
+    @override_settings(JWT_TOKEN_CLAIMS_SERIALIZER = 'tests.test_api.TESTTokenObtainPairSerializer')
+    def test_custom_jwt_claims_cookie_w_authentication(self):
+        payload = {
+            "username": self.USERNAME,
+            "password": self.PASS
+        }
+        get_user_model().objects.create_user(self.USERNAME, self.EMAIL, self.PASS)
+        resp = self.post(self.login_url, data=payload, status_code=200)
+        self.assertEqual(['jwt-auth'], list(resp.cookies.keys()))
+        token = resp.cookies.get('jwt-auth').value
+        claims = decode_jwt(token, settings.SECRET_KEY, algorithms='HS256')
+        self.assertEquals(claims['user_id'], 1)
+        self.assertEquals(claims['name'], 'person')
+        self.assertEquals(claims['email'], 'person1@world.com')
+        resp = self.get('/protected-view/')
+        self.assertEquals(resp.status_code, 200)

dj_rest_auth/utils.py

@@ -15,18 +15,15 @@ def default_create_token(token_model, user, serializer):
     return token
 
 
-def jwt_encode(user):
-    try:
-        from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
-    except ImportError:
-        raise ImportError("rest-framework-simplejwt needs to be installed")
-
-    refresh = TokenObtainPairSerializer.get_token(user)
-    return refresh.access_token, refresh
-
-
 try:
+    from django.conf import settings
     from rest_framework_simplejwt.authentication import JWTAuthentication
+    from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
+
+    def jwt_encode(user):
+        TOPS = import_callable(getattr(settings, 'JWT_TOKEN_CLAIMS_SERIALIZER', TokenObtainPairSerializer))
+        refresh = TOPS.get_token(user)
+        return refresh.access_token, refresh
 
     class JWTCookieAuthentication(JWTAuthentication):
         """
@@ -35,7 +32,6 @@ try:
         preference to the header).
         """
         def authenticate(self, request):
-            from django.conf import settings
             cookie_name = getattr(settings, 'JWT_AUTH_COOKIE', None)
             header = self.get_header(request)
             if header is None:
@@ -53,4 +49,4 @@ try:
             return self.get_user(validated_token), validated_token
 
 except ImportError:
-    pass
+    raise ImportError("rest-framework-simplejwt needs to be installed")

docs/configuration.rst

@@ -12,6 +12,8 @@ Configuration
 
         - JWT_SERIALIZER - (Using REST_USE_JWT=True) response for successful authentication in ``dj_rest_auth.views.LoginView``, default value ``dj_rest_auth.serializers.JWTSerializer``
 
+        - JWT_TOKEN_CLAIMS_SERIALIZER - A custom JWT Claim serializer. Default is `rest_framework_simplejwt.serializers.TokenObtainPairSerializer`
+
         - USER_DETAILS_SERIALIZER - serializer class in ``dj_rest_auth.views.UserDetailsView``, default value ``dj_rest_auth.serializers.UserDetailsSerializer``
 
         - PASSWORD_RESET_SERIALIZER - serializer class in ``dj_rest_auth.views.PasswordResetView``, default value ``dj_rest_auth.serializers.PasswordResetSerializer``