From 5691d061ccd72cbda4d2a73000cb54fc3ff78c60 Mon Sep 17 00:00:00 2001
From: Iman Karimi <imankarimi.mail@gmail.com>
Date: Wed, 1 Jun 2022 15:09:54 +0430
Subject: [PATCH] add refresh token view and serializer

---
 README.md                      | 11 +++++++++++
 auth_protection/serializers.py | 29 +++++++++++++++++++++++++++--
 auth_protection/views.py       | 12 ++++++++++--
 setup.py                       |  2 +-
 4 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 968946d..b3b5927 100644
--- a/README.md
+++ b/README.md
@@ -39,4 +39,15 @@ from auth_protection.authentications import JWTAuthProtection
 
 class SampleView(TARGET_VIEW):
     authentication_classes = [JWTAuthProtection]
+```
+
+- Change your `TokenRefreshView` view to `ProtectTokenRefreshView` (EX: urls.py):
+```python
+from auth_protection.views import ProtectTokenRefreshView
+
+urlpatterns = [
+    # ...
+    path('YOUR_PATH/refresh/', ProtectTokenRefreshView.as_view(), name='URL_NAME'),
+    # ...
+]
 ```
\ No newline at end of file
diff --git a/auth_protection/serializers.py b/auth_protection/serializers.py
index 5f066c7..15ed33e 100644
--- a/auth_protection/serializers.py
+++ b/auth_protection/serializers.py
@@ -1,6 +1,11 @@
-from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
-
+from rest_framework_simplejwt.serializers import TokenObtainPairSerializer, TokenRefreshSerializer
+from rest_framework_simplejwt.tokens import RefreshToken
 from auth_protection.utils import get_protect_key
+from django.contrib.auth import get_user_model
+from rest_framework_simplejwt.exceptions import InvalidToken
+from django.utils.translation import gettext_lazy as _
+
+User = get_user_model()
 
 
 class ProtectTokenObtainPairSerializer(TokenObtainPairSerializer):
@@ -14,3 +19,23 @@ def get_token(cls, user):
         token = super().get_token(user)
         token['protect_key'] = get_protect_key(user)
         return token
+
+
+class ProtectTokenRefreshSerializer(TokenRefreshSerializer):
+
+    def validate(self, attr):
+        refresh = RefreshToken(attr['refresh'])
+
+        user = self.get_user(user_id=refresh.get('user_id'))
+
+        if not refresh.get('protect_key') or (refresh.get('protect_key') != get_protect_key(user)):
+            raise InvalidToken(_('Token contained no recognizable user identification'))
+
+        return super(ProtectTokenRefreshSerializer, self).validate(attr)
+
+    def get_user(self, user_id):
+        try:
+            user = User.objects.get(id=user_id)
+        except User.DoesNotExist:
+            raise InvalidToken(_('Token contained no recognizable user identification'))
+        return user
diff --git a/auth_protection/views.py b/auth_protection/views.py
index 91ea44a..3b7b1e4 100644
--- a/auth_protection/views.py
+++ b/auth_protection/views.py
@@ -1,3 +1,11 @@
-from django.shortcuts import render
+from rest_framework_simplejwt.views import TokenRefreshView
 
-# Create your views here.
+from auth_protection.serializers import ProtectTokenRefreshSerializer
+
+
+class ProtectTokenRefreshView(TokenRefreshView):
+    """
+        Takes a refresh type JSON web token and returns an access type JSON web
+        token if the refresh token is valid.
+        """
+    serializer_class = ProtectTokenRefreshSerializer
diff --git a/setup.py b/setup.py
index d7ba73b..47c3a2e 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
 
 setup(
     name='django-auth-protection',
-    version='1.0.2',
+    version='1.0.3',
     zip_safe=False,
     packages=find_packages(),
     include_package_data=True,