Allowed usage of custom AuthToken based on knox.AbstractAuthToken (#275)

* Allowed usage of custom AuthToken based on knox.AbstractAuthToken

* moved get_token_model() to models, updated the docs

Co-authored-by: Khalidm98 <[email protected]>

.gitignore

@@ -57,3 +57,6 @@ docs/_build/
 target/
 db.sqlite3
 site/
+
+# PyCharm Project
+.idea

docs/settings.md

@@ -10,18 +10,27 @@ Example `settings.py`
 # These are the default values if none are set
 from datetime import timedelta
 from rest_framework.settings import api_settings
+
+KNOX_TOKEN_MODEL = 'knox.AuthToken'
+
 REST_KNOX = {
   'SECURE_HASH_ALGORITHM': 'hashlib.sha512',
   'AUTH_TOKEN_CHARACTER_LENGTH': 64,
   'TOKEN_TTL': timedelta(hours=10),
   'USER_SERIALIZER': 'knox.serializers.UserSerializer',
   'TOKEN_LIMIT_PER_USER': None,
   'AUTO_REFRESH': False,
+  'MIN_REFRESH_INTERVAL': 60,
+  'AUTH_HEADER_PREFIX': 'Token',
   'EXPIRY_DATETIME_FORMAT': api_settings.DATETIME_FORMAT,
+  'TOKEN_MODEL': 'knox.AuthToken',
 }
 #...snip...
 ```
 
+## KNOX_TOKEN_MODEL
+This is the variable used in the swappable dependency of the `AuthToken` model
+
 ## SECURE_HASH_ALGORITHM
 This is a reference to the class used to provide the hashing algorithm for
 token storage.
@@ -81,6 +90,11 @@ This is the expiry datetime format returned in the login view. The default is th
 [DATETIME_FORMAT][DATETIME_FORMAT] of Django REST framework. May be any of `None`, `iso-8601`
 or a Python [strftime format][strftime format] string.
 
+## TOKEN_MODEL
+This is the reference to the model used as `AuthToken`. We can define a custom `AuthToken`
+model in our project that extends `knox.AbstractAuthToken` and add our business logic to it.
+The default is `knox.AuthToken`
+
 [DATETIME_FORMAT]: https://www.django-rest-framework.org/api-guide/settings/#date-and-time-formatting
 [strftime format]: https://docs.python.org/3/library/time.html#time.strftime
 

docs/views.md

@@ -21,6 +21,7 @@ helper methods:
 - `get_user_serializer_class(self)`, to change the class used for serializing the user
 - `get_expiry_datetime_format(self)`, to change the datetime format used for expiry
 - `format_expiry_datetime(self, expiry)`, to format the expiry `datetime` object at your convenience
+- `create_token(self)`, to create the `AuthToken` instance at your convenience
 
 Finally, if none of these helper methods are sufficient, you can also override `get_post_response_data`
 to return a fully customized payload.

knox/auth.py

@@ -14,7 +14,7 @@ def compare_digest(a, b):
 )
 
 from knox.crypto import hash_token
-from knox.models import AuthToken
+from knox.models import get_token_model
 from knox.settings import CONSTANTS, knox_settings
 from knox.signals import token_expired
 
@@ -31,7 +31,6 @@ class TokenAuthentication(BaseAuthentication):
     - `request.user` will be a django `User` instance
     - `request.auth` will be an `AuthToken` instance
     '''
-    model = AuthToken
 
     def authenticate(self, request):
         auth = get_authorization_header(request).split()
@@ -62,7 +61,7 @@ def authenticate_credentials(self, token):
         '''
         msg = _('Invalid token.')
         token = token.decode("utf-8")
-        for auth_token in AuthToken.objects.filter(
+        for auth_token in get_token_model().objects.filter(
                 token_key=token[:CONSTANTS.TOKEN_KEY_LENGTH]):
             if self._cleanup_token(auth_token):
                 continue

knox/migrations/0001_initial.py

@@ -9,6 +9,7 @@ class Migration(migrations.Migration):
 
     dependencies = [
         migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        migrations.swappable_dependency(settings.KNOX_TOKEN_MODEL),
     ]
 
     operations = [

knox/models.py

@@ -1,4 +1,6 @@
+from django.apps import apps
 from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
 from django.db import models
 from django.utils import timezone
 
@@ -9,7 +11,7 @@
 
 
 class AuthTokenManager(models.Manager):
-    def create(self, user, expiry=knox_settings.TOKEN_TTL):
+    def create(self, user, expiry=knox_settings.TOKEN_TTL, **kwargs):
         token = crypto.create_token_string()
         digest = crypto.hash_token(token)
 
@@ -18,11 +20,11 @@ def create(self, user, expiry=knox_settings.TOKEN_TTL):
 
         instance = super(AuthTokenManager, self).create(
             token_key=token[:CONSTANTS.TOKEN_KEY_LENGTH], digest=digest,
-            user=user, expiry=expiry)
+            user=user, expiry=expiry, **kwargs)
         return instance, token
 
 
-class AuthToken(models.Model):
+class AbstractAuthToken(models.Model):
 
     objects = AuthTokenManager()
 
@@ -35,5 +37,31 @@ class AuthToken(models.Model):
     created = models.DateTimeField(auto_now_add=True)
     expiry = models.DateTimeField(null=True, blank=True)
 
+    class Meta:
+        abstract = True
+
     def __str__(self):
         return '%s : %s' % (self.digest, self.user)
+
+
+class AuthToken(AbstractAuthToken):
+    class Meta:
+        swappable = 'KNOX_TOKEN_MODEL'
+
+
+def get_token_model():
+    """
+    Return the AuthToken model that is active in this project.
+    """
+
+    try:
+        return apps.get_model(knox_settings.TOKEN_MODEL, require_ready=False)
+    except ValueError:
+        raise ImproperlyConfigured(
+            "TOKEN_MODEL must be of the form 'app_label.model_name'"
+        )
+    except LookupError:
+        raise ImproperlyConfigured(
+            "TOKEN_MODEL refers to model '%s' that has not been installed"
+            % knox_settings.TOKEN_MODEL
+        )

knox/settings.py

@@ -16,6 +16,7 @@
     'MIN_REFRESH_INTERVAL': 60,
     'AUTH_HEADER_PREFIX': 'Token',
     'EXPIRY_DATETIME_FORMAT': api_settings.DATETIME_FORMAT,
+    'TOKEN_MODEL': getattr(settings, 'KNOX_TOKEN_MODEL', 'knox.AuthToken'),
 }
 
 IMPORT_STRINGS = {

knox/views.py

@@ -8,7 +8,7 @@
 from rest_framework.views import APIView
 
 from knox.auth import TokenAuthentication
-from knox.models import AuthToken
+from knox.models import get_token_model
 from knox.settings import knox_settings
 
 
@@ -35,6 +35,11 @@ def format_expiry_datetime(self, expiry):
         datetime_format = self.get_expiry_datetime_format()
         return DateTimeField(format=datetime_format).to_representation(expiry)
 
+    def create_token(self):
+        return get_token_model().objects.create(
+            user=self.request.user, expiry=self.get_token_ttl()
+        )
+
     def get_post_response_data(self, request, token, instance):
         UserSerializer = self.get_user_serializer_class()
 
@@ -59,8 +64,7 @@ def post(self, request, format=None):
                     {"error": "Maximum amount of tokens allowed per user exceeded."},
                     status=status.HTTP_403_FORBIDDEN
                 )
-        token_ttl = self.get_token_ttl()
-        instance, token = AuthToken.objects.create(request.user, token_ttl)
+        instance, token = self.create_token()
         user_logged_in.send(sender=request.user.__class__,
                             request=request, user=request.user)
         data = self.get_post_response_data(request, token, instance)

knox_project/settings.py

@@ -53,3 +53,5 @@
 USE_TZ = True
 
 STATIC_URL = '/static/'
+
+KNOX_TOKEN_MODEL = 'knox.AuthToken'