From f73f3af1d3698336d0ad15dceb0add289cb2188c Mon Sep 17 00:00:00 2001
From: alichass <alichass@users.noreply.github.com>
Date: Thu, 19 Mar 2020 17:09:20 -0400
Subject: [PATCH] hopefully this should work?

---
 dj_rest_auth/app_settings.py   |  3 +++
 dj_rest_auth/tests/settings.py |  2 +-
 dj_rest_auth/utils.py          | 30 +++++++++++++++++++++++++++++-
 dj_rest_auth/views.py          | 22 +++++++++++-----------
 docs/installation.rst          |  2 +-
 5 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/dj_rest_auth/app_settings.py b/dj_rest_auth/app_settings.py
index d89202d..39e44ef 100644
--- a/dj_rest_auth/app_settings.py
+++ b/dj_rest_auth/app_settings.py
@@ -34,3 +34,6 @@ PasswordResetConfirmSerializer = serializers.get(
 )
 
 PasswordChangeSerializer = import_callable(serializers.get('PASSWORD_CHANGE_SERIALIZER', DefaultPasswordChangeSerializer))
+
+JWT_AUTH_COOKIE = getattr(settings, 'JWT_AUTH_COOKIE', None)
+
diff --git a/dj_rest_auth/tests/settings.py b/dj_rest_auth/tests/settings.py
index 7279c21..3d02309 100644
--- a/dj_rest_auth/tests/settings.py
+++ b/dj_rest_auth/tests/settings.py
@@ -68,7 +68,7 @@ TEMPLATES = [
 REST_FRAMEWORK = {
     'DEFAULT_AUTHENTICATION_CLASSES': (
         'rest_framework.authentication.SessionAuthentication',
-        'rest_framework_simplejwt.authentication.JWTAuthentication',
+        'dj_rest_auth.utils.JWTCookieAuthentication',
     )
 }
 
diff --git a/dj_rest_auth/utils.py b/dj_rest_auth/utils.py
index 85912a6..2b8094a 100644
--- a/dj_rest_auth/utils.py
+++ b/dj_rest_auth/utils.py
@@ -1,5 +1,5 @@
 from importlib import import_module
-
+from .app_settings import JWT_AUTH_COOKIE
 
 def import_callable(path_or_callable):
     if hasattr(path_or_callable, '__call__'):
@@ -23,3 +23,31 @@ def jwt_encode(user):
 
     refresh = TokenObtainPairSerializer.get_token(user)
     return refresh.access_token, refresh
+
+
+try:
+    from rest_framework_simplejwt.authentication import JWTAuthentication
+
+    class JWTCookieAuthentication(JWTAuthentication):
+        """
+        An authentication plugin that hopefully authenticates requests through a JSON web
+        token provided in a request cookie (and through the header as normal, with a preference to the header).
+        """
+        def authenticate(self, request):
+            header = self.get_header(request)
+            if header is None:
+                if JWT_AUTH_COOKIE: # or settings.JWT_AUTH_COOKIE
+                    raw_token = request.COOKIES.get(JWT_AUTH_COOKIE)  # or settings.jwt_auth_cookie
+                else:
+                    return None
+            else:
+                raw_token = self.get_raw_token(header)
+                
+            if raw_token is None:
+                return None
+
+            validated_token = self.get_validated_token(raw_token)
+
+            return self.get_user(validated_token), validated_token
+except ImportError as I:
+    pass
diff --git a/dj_rest_auth/views.py b/dj_rest_auth/views.py
index 37e4880..2896209 100644
--- a/dj_rest_auth/views.py
+++ b/dj_rest_auth/views.py
@@ -16,7 +16,7 @@ from .app_settings import (JWTSerializer, LoginSerializer,
                            PasswordChangeSerializer,
                            PasswordResetConfirmSerializer,
                            PasswordResetSerializer, TokenSerializer,
-                           UserDetailsSerializer, create_token)
+                           UserDetailsSerializer, create_token, JWT_AUTH_COOKIE)
 from .models import TokenModel
 from .utils import jwt_encode
 
@@ -85,13 +85,13 @@ class LoginView(GenericAPIView):
         response = Response(serializer.data, status=status.HTTP_200_OK)
         if getattr(settings, 'REST_USE_JWT', False):
             from rest_framework_simplejwt.settings import api_settings as jwt_settings
-            #if jwt_settings.JWT_AUTH_COOKIE        #this needs to be added to simplejwt
-            from datetime import datetime
-            expiration = (datetime.utcnow() + jwt_settings.ACCESS_TOKEN_LIFETIME)
-            response.set_cookie('somestring',     #replace with jwt_settings.JWT_AUTH_COOKIE
-                                self.access_token,
-                                expires=expiration,
-                                httponly=True)
+            if JWT_AUTH_COOKIE:        #this needs to be added to simplejwt
+                from datetime import datetime
+                expiration = (datetime.utcnow() + jwt_settings.ACCESS_TOKEN_LIFETIME)
+                response.set_cookie(JWT_AUTH_COOKIE,     #this needs to be added to simplejwt
+                                    self.access_token,
+                                    expires=expiration,
+                                    httponly=True)
         return response
 
     def post(self, request, *args, **kwargs):
@@ -135,9 +135,9 @@ class LogoutView(APIView):
         response = Response({"detail": _("Successfully logged out.")},
                             status=status.HTTP_200_OK)
         if getattr(settings, 'REST_USE_JWT', False):
-            from rest_framework_simplejwt.settings import api_settings as jwt_settings
-            #if jwt_settings.JWT_AUTH_COOKIE        #this needs to be added to simplejwt
-            response.delete_cookie('somestring') #replace with jwt_settings.JWT_AUTH_COOKIE
+            # from rest_framework_simplejwt.settings import api_settings as jwt_settings
+            if JWT_AUTH_COOKIE:        #this needs to be added to simplejwt
+                response.delete_cookie(JWT_AUTH_COOKIE)    #this needs to be added to simplejwt
         return response
 
 
diff --git a/docs/installation.rst b/docs/installation.rst
index 5004752..fa4345d 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -259,7 +259,7 @@ By default ``dj-rest-auth`` uses Django's Token-based authentication. If you wan
         ...
         'DEFAULT_AUTHENTICATION_CLASSES': (
             ...
-            'rest_framework_simplejwt.authentication.JWTAuthentication',
+            'dj_rest_auth.utils.JWTCookieAuthentication',
         )
         ...
     }