Allow passing request/view context to JWT token generation

dj_rest_auth/app_settings.py

@@ -10,7 +10,7 @@
     'TOKEN_SERIALIZER': 'dj_rest_auth.serializers.TokenSerializer',
     'JWT_SERIALIZER': 'dj_rest_auth.serializers.JWTSerializer',
     'JWT_SERIALIZER_WITH_EXPIRATION': 'dj_rest_auth.serializers.JWTSerializerWithExpiration',
-    'JWT_TOKEN_CLAIMS_SERIALIZER': 'rest_framework_simplejwt.serializers.TokenObtainPairSerializer',
+    'JWT_TOKEN_CLAIMS_SERIALIZER': 'dj_rest_auth.serializers.TokenWithCustomClaimsSerializer',
     'USER_DETAILS_SERIALIZER': 'dj_rest_auth.serializers.UserDetailsSerializer',
     'PASSWORD_RESET_SERIALIZER': 'dj_rest_auth.serializers.PasswordResetSerializer',
     'PASSWORD_RESET_CONFIRM_SERIALIZER': 'dj_rest_auth.serializers.PasswordResetConfirmSerializer',

dj_rest_auth/serializers.py

@@ -6,6 +6,7 @@
 from django.utils.translation import gettext_lazy as _
 from rest_framework import exceptions, serializers
 from rest_framework.exceptions import ValidationError
+from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
 
 from .app_settings import api_settings
 
@@ -359,3 +360,14 @@ def save(self):
         if not self.logout_on_password_change:
             from django.contrib.auth import update_session_auth_hash
             update_session_auth_hash(self.request, self.user)
+
+
+class TokenWithCustomClaimsSerializer(TokenObtainPairSerializer):
+
+    @classmethod
+    def get_token(cls, user, custom_claims={}):
+        token = super().get_token(user)
+        if custom_claims:
+            for key, value in custom_claims.items():
+                token[key] = value
+        return token

dj_rest_auth/utils.py

@@ -6,12 +6,12 @@ def default_create_token(token_model, user, serializer):
     return token
 
 
-def jwt_encode(user):
+def jwt_encode(user, **kwargs):
     from dj_rest_auth.app_settings import api_settings
 
     JWTTokenClaimsSerializer = api_settings.JWT_TOKEN_CLAIMS_SERIALIZER
-
-    refresh = JWTTokenClaimsSerializer.get_token(user)
+    custom_claims = kwargs.pop('custom_claims', {})
+    refresh = JWTTokenClaimsSerializer.get_token(user, custom_claims)
     return refresh.access_token, refresh
 
 

dj_rest_auth/views.py

@@ -62,12 +62,12 @@ def get_response_serializer(self):
             response_serializer = api_settings.TOKEN_SERIALIZER
         return response_serializer
 
-    def login(self):
+    def login(self, **kwargs):
         self.user = self.serializer.validated_data['user']
         token_model = get_token_model()
 
         if api_settings.USE_JWT:
-            self.access_token, self.refresh_token = jwt_encode(self.user)
+            self.access_token, self.refresh_token = jwt_encode(self.user, **kwargs)
         elif token_model:
             self.token = api_settings.TOKEN_CREATOR(token_model, self.user, self.serializer)
 
@@ -123,10 +123,13 @@ def post(self, request, *args, **kwargs):
         self.request = request
         self.serializer = self.get_serializer(data=self.request.data)
         self.serializer.is_valid(raise_exception=True)
-
-        self.login()
+        kwargs['custom_claims'] = self.get_non_user_custom_claims(request, *args, **kwargs)
+        self.login(**kwargs)
         return self.get_response()
 
+    def get_non_user_custom_claims(self, request, *args, **kwargs):
+        return {}
+
 
 class LogoutView(APIView):
     """