Merge pull request #109 from sphrak/feature/add-optional-token-limit-…

…per-user

add: optional token limit per user object

docs/changes.md

@@ -1,4 +1,7 @@
-#Changelog
+# Changelog
+
+## 3.4.0
+- adds optional token limit
 
 ## 3.3.1
 - Ensure compatibility with Django 2.1 up to Python 3.7

docs/index.md

@@ -13,7 +13,8 @@ in to DRF. However, it overcomes some problems present in the default implementa
 
     Knox provides one token per call to the login view - allowing
     each client to have its own token which is deleted on the server side when the client
-    logs out.
+    logs out. Knox also provides an optional setting to limit the amount of tokens generated
+    per user.
 
     Knox also provides an option for a logged in client to remove *all* tokens
     that the server has - forcing all clients to re-authenticate.

docs/settings.md

@@ -14,6 +14,7 @@ REST_KNOX = {
   'AUTH_TOKEN_CHARACTER_LENGTH': 64,
   'TOKEN_TTL': timedelta(hours=10),
   'USER_SERIALIZER': 'knox.serializers.UserSerializer',
+  'TOKEN_LIMIT_PER_USER': None,
   'AUTO_REFRESH': FALSE,
 }
 #...snip...
@@ -54,6 +55,10 @@ Setting the TOKEN_TTL to `None` will create tokens that never expire.
 Warning: setting a 0 or negative timedelta will create tokens that instantly expire,
 the system will not prevent you setting this.
 
+## TOKEN_LIMIT_PER_USER
+This allows you to control how many tokens can be issued per user.
+By default this option is disabled and set to `None` -- thus no limit.
+
 ## USER_SERIALIZER
 This is the reference to the class used to serialize the `User` objects when
 succesfully returning from `LoginView`. The default is `knox.serializers.UserSerializer`

knox/settings.py

@@ -10,6 +10,7 @@
     'AUTH_TOKEN_CHARACTER_LENGTH': 64,
     'TOKEN_TTL': timedelta(hours=10),
     'USER_SERIALIZER': None,
+    'TOKEN_LIMIT_PER_USER': None,
     'AUTO_REFRESH': False,
     'MIN_REFRESH_INTERVAL': 60,
 }

knox/views.py

@@ -1,4 +1,7 @@
 from django.contrib.auth.signals import user_logged_in, user_logged_out
+from django.utils import timezone
+from django.conf import settings
+
 from rest_framework import status
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
@@ -15,6 +18,12 @@ class LoginView(APIView):
     permission_classes = (IsAuthenticated,)
 
     def post(self, request, format=None):
+        if knox_settings.TOKEN_LIMIT_PER_USER is not None:
+            if request.user.auth_token_set.filter(expires__gt=timezone.now()).count() >= knox_settings.TOKEN_LIMIT_PER_USER:
+                return Response(
+                    {"error": "Maximum amount of tokens allowed per user exceeded."},
+                    status=status.HTTP_403_FORBIDDEN
+                )
         token = AuthToken.objects.create(request.user)
         user_logged_in.send(sender=request.user.__class__, request=request, user=request.user)
         UserSerializer = knox_settings.USER_SERIALIZER

tests/tests.py

@@ -4,7 +4,7 @@
 from django.utils.six.moves import reload_module
 from django.contrib.auth import get_user_model
 from django.test import override_settings
-from knox import auth
+from knox import auth, views
 
 try:
     # For django >= 2.0
@@ -38,6 +38,11 @@ def get_basic_auth_header(username, password):
 auto_refresh_knox = knox_settings.defaults.copy()
 auto_refresh_knox["AUTO_REFRESH"] = True
 
+token_user_limit_knox = knox_settings.defaults.copy()
+token_user_limit_knox["TOKEN_LIMIT_PER_USER"] = 10
+
+user_serializer_knox = knox_settings.defaults.copy()
+user_serializer_knox["USER_SERIALIZER"] = UserSerializer
 
 class AuthTestCase(TestCase):
 
@@ -73,15 +78,18 @@ def test_login_returns_serialized_token(self):
         self.assertNotIn(username_field, response.data)
 
     def test_login_returns_serialized_token_and_username_field(self):
-        self.assertEqual(AuthToken.objects.count(), 0)
-        url = reverse('knox_login')
-        self.client.credentials(
-            HTTP_AUTHORIZATION=get_basic_auth_header(self.username, self.password)
-        )
-        knox_settings.USER_SERIALIZER = UserSerializer
-        response = self.client.post(url, {}, format='json')
+
+        with override_settings(REST_KNOX=user_serializer_knox):
+            reload_module(views)
+            self.assertEqual(AuthToken.objects.count(), 0)
+            url = reverse('knox_login')
+            self.client.credentials(
+                HTTP_AUTHORIZATION=get_basic_auth_header(self.username, self.password)
+            )
+            response = self.client.post(url, {}, format='json')
+            self.assertEqual(user_serializer_knox["USER_SERIALIZER"], UserSerializer)
+        reload_module(views)
         self.assertEqual(response.status_code, 200)
-        self.assertNotEqual(knox_settings.USER_SERIALIZER, None)
         self.assertIn('token', response.data)
         username_field = self.user.USERNAME_FIELD
         self.assertIn('user', response.data)
@@ -254,3 +262,37 @@ def handler(sender, username, **kwargs):
 
         self.assertTrue(self.signal_was_called)
 
+    def test_exceed_token_amount_per_user(self):
+
+        with override_settings(REST_KNOX=token_user_limit_knox):
+            reload_module(views)
+            for _ in range(10):
+                token = AuthToken.objects.create(user=self.user)
+            url = reverse('knox_login')
+            self.client.credentials(
+                HTTP_AUTHORIZATION=get_basic_auth_header(self.username, self.password)
+            )
+            response = self.client.post(url, {}, format='json')
+        reload_module(views)
+        self.assertEqual(response.status_code, 403)
+        self.assertEqual(response.data, {"error": "Maximum amount of tokens allowed per user exceeded."})
+
+    def test_does_not_exceed_on_expired_keys(self):
+
+        with override_settings(REST_KNOX=token_user_limit_knox):
+            reload_module(views)
+            for _ in range(9):
+                token = AuthToken.objects.create(user=self.user)
+            AuthToken.objects.create(user=self.user, expires=timedelta(seconds=0))
+            # now 10 keys, but only 9 valid so request should succeed.
+            url = reverse('knox_login')
+            self.client.credentials(
+                HTTP_AUTHORIZATION=get_basic_auth_header(self.username, self.password)
+            )
+            response = self.client.post(url, {}, format='json')
+            failed_response = self.client.post(url, {}, format='json')
+        reload_module(views)
+        self.assertEqual(response.status_code, 200)
+        self.assertIn('token', response.data)
+        self.assertEqual(failed_response.status_code, 403)
+        self.assertEqual(failed_response.data, {"error": "Maximum amount of tokens allowed per user exceeded."})