Token view extended with throttle for usernames

user_management/api/tests/test_throttling.py

@@ -1,6 +1,8 @@
 from django.core.cache import cache
+from django.core.urlresolvers import reverse_lazy
 from mock import patch
 from rest_framework import status
+from rest_framework.test import APIRequestFactory
 
 from user_management.api import views
 from user_management.models.tests.factories import UserFactory
@@ -10,6 +12,8 @@
 
 
 class GetTokenTest(APIRequestTestCase):
+    auth_url = reverse_lazy('user_management_api:auth')
+    throttle_expected_status = status.HTTP_429_TOO_MANY_REQUESTS
     view_class = views.GetToken
 
     def tearDown(self):
@@ -31,7 +35,39 @@ def test_user_auth_throttle(self):
 
         # request should be throttled now
         response = view(request)
-        self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
+        self.assertEqual(response.status_code, self.throttle_expected_status)
+
+    @patch(THROTTLE_RATE_PATH, new={'logins': '1/minute'})
+    def test_user_auth_throttle_ip(self):
+        """Ensure user gets throttled from a single IP address."""
+        data = {}
+
+        request = APIRequestFactory().post(self.auth_url, data, REMOTE_ADDR='127.0.0.1')
+        response = self.view_class.as_view()(request)
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+        request = APIRequestFactory().post(self.auth_url, data)
+        response = self.view_class.as_view()(request)
+        self.assertEqual(response.status_code, self.throttle_expected_status)
+
+    @patch(THROTTLE_RATE_PATH, new={'logins': '1/minute'})
+    def test_user_auth_throttle_username(self):
+        """Ensure username is throttled no matter what IP the user connects on."""
+        data = {'username': 'jimmy'}
+        request = APIRequestFactory().post(self.auth_url, data, REMOTE_ADDR='127.0.0.1')
+        response = self.view_class.as_view()(request)
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+        # Different user name to different ip is not throttled
+        data2 = {'username': 'another_jimmy_here'}
+        request = APIRequestFactory().post(self.auth_url, data2, REMOTE_ADDR='127.0.0.2')
+        response = self.view_class.as_view()(request)
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+        # Same username to different ip is throttled
+        request = APIRequestFactory().post(self.auth_url, data, REMOTE_ADDR='127.0.0.3')
+        response = self.view_class.as_view()(request)
+        self.assertEqual(response.status_code, self.throttle_expected_status)
 
 
 class TestPasswordResetEmail(APIRequestTestCase):

user_management/api/throttling.py

@@ -27,6 +27,21 @@ class LoginRateThrottle(
     default_rate = '10/hour'
 
 
+class UsernameLoginRateThrottle(LoginRateThrottle):
+    def get_cache_key(self, request, view):
+        if request.user.is_authenticated():
+            return None  # Only throttle unauthenticated requests
+
+        ident = request.POST.get('username')
+        if ident is None:
+            return None  # Only throttle username requests
+
+        return self.cache_format % {
+            'scope': self.scope,
+            'ident': ident.strip().lower(),
+        }
+
+
 class PasswordResetRateThrottle(
         DefaultRateMixin,
         PostRequestThrottleMixin,

user_management/api/views.py

@@ -19,7 +19,10 @@
 
 class GetToken(ObtainAuthToken):
     renderer_classes = (renderers.JSONRenderer, renderers.BrowsableAPIRenderer)
-    throttle_classes = [throttling.LoginRateThrottle]
+    throttle_classes = [
+        throttling.UsernameLoginRateThrottle,
+        throttling.LoginRateThrottle,
+    ]
     throttle_scope = 'logins'
 
     def delete(self, request, *args, **kwargs):