From 9bbc4a3014df5d1da3771bbd1c7170a976f66842 Mon Sep 17 00:00:00 2001
From: alichass <alichass@users.noreply.github.com>
Date: Wed, 11 Mar 2020 06:01:03 -0400
Subject: [PATCH] change standard auth stuff to work with simplejwt

---
 dj_rest_auth/serializers.py |  3 ++-
 dj_rest_auth/utils.py       | 12 +++++-------
 dj_rest_auth/views.py       | 18 +++---------------
 3 files changed, 10 insertions(+), 23 deletions(-)

diff --git a/dj_rest_auth/serializers.py b/dj_rest_auth/serializers.py
index 07783da..b34c9cc 100644
--- a/dj_rest_auth/serializers.py
+++ b/dj_rest_auth/serializers.py
@@ -135,7 +135,8 @@ class JWTSerializer(serializers.Serializer):
     """
     Serializer for JWT authentication.
     """
-    token = serializers.CharField()
+    access_token = serializers.CharField()
+    refresh_token = serializers.CharField()
     user = serializers.SerializerMethodField()
 
     def get_user(self, obj):
diff --git a/dj_rest_auth/utils.py b/dj_rest_auth/utils.py
index b4858c8..165963d 100644
--- a/dj_rest_auth/utils.py
+++ b/dj_rest_auth/utils.py
@@ -17,12 +17,10 @@ def default_create_token(token_model, user, serializer):
 
 def jwt_encode(user):
     try:
-        from rest_framework_jwt.settings import api_settings
+        from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
+        from rest_framework_simplejwt.views import TokenObtainPairView
     except ImportError:
-        raise ImportError("djangorestframework_jwt needs to be installed")
+        raise ImportError("rest-framework-simplejwt needs to be installed")
 
-    jwt_payload_handler = api_settings.JWT_PAYLOAD_HANDLER
-    jwt_encode_handler = api_settings.JWT_ENCODE_HANDLER
-
-    payload = jwt_payload_handler(user)
-    return jwt_encode_handler(payload)
+    refresh = TokenObtainPairSerializer.get_token(user)
+    return refresh.access_token, refresh
diff --git a/dj_rest_auth/views.py b/dj_rest_auth/views.py
index c5bcb1a..f30980c 100644
--- a/dj_rest_auth/views.py
+++ b/dj_rest_auth/views.py
@@ -59,7 +59,7 @@ class LoginView(GenericAPIView):
         self.user = self.serializer.validated_data['user']
 
         if getattr(settings, 'REST_USE_JWT', False):
-            self.token = jwt_encode(self.user)
+            self.access_token, self.refresh_token = jwt_encode(self.user)
         else:
             self.token = create_token(self.token_model, self.user,
                                       self.serializer)
@@ -73,7 +73,8 @@ class LoginView(GenericAPIView):
         if getattr(settings, 'REST_USE_JWT', False):
             data = {
                 'user': self.user,
-                'token': self.token
+                'access_token': self.access_token,
+                'refresh_token': self.refresh_token
             }
             serializer = serializer_class(instance=data,
                                           context={'request': self.request})
@@ -82,15 +83,6 @@ class LoginView(GenericAPIView):
                                           context={'request': self.request})
 
         response = Response(serializer.data, status=status.HTTP_200_OK)
-        if getattr(settings, 'REST_USE_JWT', False):
-            from rest_framework_jwt.settings import api_settings as jwt_settings
-            if jwt_settings.JWT_AUTH_COOKIE:
-                from datetime import datetime
-                expiration = (datetime.utcnow() + jwt_settings.JWT_EXPIRATION_DELTA)
-                response.set_cookie(jwt_settings.JWT_AUTH_COOKIE,
-                                    self.token,
-                                    expires=expiration,
-                                    httponly=True)
         return response
 
     def post(self, request, *args, **kwargs):
@@ -133,10 +125,6 @@ class LogoutView(APIView):
 
         response = Response({"detail": _("Successfully logged out.")},
                             status=status.HTTP_200_OK)
-        if getattr(settings, 'REST_USE_JWT', False):
-            from rest_framework_jwt.settings import api_settings as jwt_settings
-            if jwt_settings.JWT_AUTH_COOKIE:
-                response.delete_cookie(jwt_settings.JWT_AUTH_COOKIE)
         return response