From 2be96132cda844697a66d2cccc8de96b72ecafae Mon Sep 17 00:00:00 2001
From: Ian Foote <ian@incuna.com>
Date: Fri, 20 Jun 2014 12:02:35 +0100
Subject: [PATCH] Add default throttle rates

* Limit GetToken requests to 10/hour
* Limit PasswordResetEmail requests to 3/hour
---
 README.md                               |  9 +++++--
 user_management/api/tests/test_views.py | 36 +++++++++++++++++++++++++
 user_management/api/throttling.py       | 17 ++++++++++++
 user_management/api/views.py            |  7 +++--
 user_management/tests/run.py            |  4 ---
 5 files changed, 63 insertions(+), 10 deletions(-)
 create mode 100644 user_management/api/throttling.py

diff --git a/README.md b/README.md
index 57a4c37..cb3ea72 100644
--- a/README.md
+++ b/README.md
@@ -138,8 +138,13 @@ with a selection from
 The `/auth/` and `/auth/password_reset/` URLs are protected against throttling
 using the built-in [DRF throttle module](http://www.django-rest-framework.org/api-guide/throttling).
 
-You need to set the throttling rates in `DEFAULT_THROTTLE_RATES` setting for
-`REST_FRAMEWORK` in your `settings.py`:
+The default throttle rates are:
+    
+    'logins': '10/hour'
+    'passwords': '3/hour'
+
+You can customise the throttling rates by setting `DEFAULT_THROTTLE_RATES`
+in your `settings.py`:
 
     DEFAULT_THROTTLE_RATES = {
         'logins': '100/day',
diff --git a/user_management/api/tests/test_views.py b/user_management/api/tests/test_views.py
index d99c5b7..0109785 100644
--- a/user_management/api/tests/test_views.py
+++ b/user_management/api/tests/test_views.py
@@ -65,6 +65,24 @@ def test_delete_no_token(self):
         response = self.view_class.as_view()(request)
         self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
 
+    def test_default_user_auth_throttle(self):
+        default_rate = 10
+        auth_url = reverse('user_management_api:auth')
+        expected_status = status.HTTP_429_TOO_MANY_REQUESTS
+
+        request = APIRequestFactory().get(auth_url)
+        view = self.view_class.as_view()
+
+        # make all but one of our allowed requests
+        for i in range(default_rate - 1):
+            view(request)
+
+        response = view(request)  # our last allowed request
+        self.assertNotEqual(response.status_code, expected_status)
+
+        response = view(request)  # our throttled request
+        self.assertEqual(response.status_code, expected_status)
+
     @patch('rest_framework.throttling.ScopedRateThrottle.THROTTLE_RATES', new={
         'logins': '1/minute',
     })
@@ -252,6 +270,24 @@ def test_options(self):
             expected_post_options,
         )
 
+    def test_default_user_password_reset_throttle(self):
+        default_rate = 3
+        auth_url = reverse('user_management_api:password_reset')
+        expected_status = status.HTTP_429_TOO_MANY_REQUESTS
+
+        request = APIRequestFactory().get(auth_url)
+        view = self.view_class.as_view()
+
+        # make all but one of our allowed requests
+        for i in range(default_rate - 1):
+            view(request)
+
+        response = view(request)  # our last allowed request
+        self.assertNotEqual(response.status_code, expected_status)
+
+        response = view(request)  # our throttled request
+        self.assertEqual(response.status_code, expected_status)
+
     @patch('rest_framework.throttling.ScopedRateThrottle.THROTTLE_RATES', new={
         'passwords': '1/minute',
     })
diff --git a/user_management/api/throttling.py b/user_management/api/throttling.py
new file mode 100644
index 0000000..9137d78
--- /dev/null
+++ b/user_management/api/throttling.py
@@ -0,0 +1,17 @@
+from rest_framework.throttling import ScopedRateThrottle
+
+
+class DefaultRateMixin(object):
+    def get_rate(self):
+        try:
+            return self.THROTTLE_RATES[self.scope]
+        except KeyError:
+            return self.default_rate
+
+
+class LoginRateThrottle(DefaultRateMixin, ScopedRateThrottle):
+    default_rate = '10/hour'
+
+
+class PasswordResetRateThrottle(DefaultRateMixin, ScopedRateThrottle):
+    default_rate = '3/hour'
diff --git a/user_management/api/views.py b/user_management/api/views.py
index efe7a52..703989a 100644
--- a/user_management/api/views.py
+++ b/user_management/api/views.py
@@ -10,9 +10,8 @@
 from rest_framework.authtoken.models import Token
 from rest_framework.authtoken.views import ObtainAuthToken
 from rest_framework.permissions import AllowAny, IsAuthenticated
-from rest_framework.throttling import ScopedRateThrottle
 
-from . import serializers, permissions
+from . import permissions, serializers, throttling
 
 
 User = get_user_model()
@@ -20,7 +19,7 @@
 
 class GetToken(ObtainAuthToken):
     renderer_classes = (renderers.JSONRenderer, renderers.BrowsableAPIRenderer)
-    throttle_classes = [ScopedRateThrottle]
+    throttle_classes = [throttling.LoginRateThrottle]
     throttle_scope = 'logins'
 
     def delete(self, request, *args, **kwargs):
@@ -70,7 +69,7 @@ class PasswordResetEmail(generics.GenericAPIView):
     permission_classes = [permissions.IsNotAuthenticated]
     template_name = 'user_management/password_reset_email.html'
     serializer_class = serializers.PasswordResetEmailSerializer
-    throttle_classes = [ScopedRateThrottle]
+    throttle_classes = [throttling.PasswordResetRateThrottle]
     throttle_scope = 'passwords'
 
     def email_context(self, site, user):
diff --git a/user_management/tests/run.py b/user_management/tests/run.py
index 1d59976..935e9c0 100755
--- a/user_management/tests/run.py
+++ b/user_management/tests/run.py
@@ -39,10 +39,6 @@
         'DEFAULT_AUTHENTICATION_CLASSES': (
             'rest_framework.authentication.TokenAuthentication',
         ),
-        'DEFAULT_THROTTLE_RATES': {
-            'logins': '100/day',
-            'passwords': '100/day',
-        }
     },
 )