From 8df7df763c42371f75b30a18963cb33161823cb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jes=C3=BAs=20Espino?= <jespinog@gmail.com>
Date: Thu, 2 Mar 2017 07:47:31 +0100
Subject: [PATCH] Improvements in the throttling code

---
 taiga/base/throttling.py           | 42 ++++++++++++++++--------------
 tests/unit/test_common_throttle.py | 30 +++++++++++++++++++++
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/taiga/base/throttling.py b/taiga/base/throttling.py
index 329c7441..d4d3a361 100644
--- a/taiga/base/throttling.py
+++ b/taiga/base/throttling.py
@@ -18,6 +18,9 @@
 
 from taiga.base.api import throttling
 from django.conf import settings
+from ipware.ip import get_ip
+from netaddr import all_matching_cidrs
+from netaddr.core import AddrFormatError
 
 
 class GlobalThrottlingMixin:
@@ -26,9 +29,7 @@ class GlobalThrottlingMixin:
     logged in or not.
     """
     def get_cache_key(self, request, view):
-        ident = request.META.get("HTTP_X_FORWARDED_FOR")
-        if ident is None:
-            ident = request.META.get("REMOTE_ADDR")
+        ident = get_ip(request)
 
         return self.cache_format % {
             "scope": self.scope,
@@ -67,11 +68,24 @@ class CommonThrottle(throttling.SimpleRateThrottle):
     def has_to_finalize(self, request, response, view):
         return False
 
+    def is_whitelisted(self, ident):
+        for whitelisted in settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST']:
+            if isinstance(whitelisted, int) and whitelisted == ident:
+                return True
+            elif isinstance(whitelisted, str):
+                try:
+                    if all_matching_cidrs(ident, [whitelisted]) != []:
+                        return True
+                except(AddrFormatError, ValueError):
+                    pass
+        return False
+
     def allow_request(self, request, view):
         scope = self.get_scope(request)
         ident = self.get_ident(request)
         rates = self.get_rates(scope)
-        if ident in settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST']:
+
+        if self.is_whitelisted(ident):
             return True
 
         if rates is None or rates == []:
@@ -145,25 +159,15 @@ class CommonThrottle(throttling.SimpleRateThrottle):
         return (rate, num_requests, duration)
 
     def get_scope(self, request):
-        if request.user.is_authenticated():
-            if request.method in ["POST", "PUT", "PATCH", "DELETE"]:
-                scope = "user-write"
-            else:
-                scope = "user-read"
-        else:
-            if request.method in ["POST", "PUT", "PATCH", "DELETE"]:
-                scope = "anon-write"
-            else:
-                scope = "anon-read"
+        scope_prefix = "user" if request.user.is_authenticated() else "anon"
+        scope_sufix = "write" if request.method in ["POST", "PUT", "PATCH", "DELETE"] else "read"
+        scope = "{}-{}".format(scope_prefix, scope_sufix)
         return scope
 
     def get_ident(self, request):
         if request.user.is_authenticated():
-            ident = request.user.id
-        else:
-            ident = request.META.get("HTTP_X_FORWARDED_FOR")
-            if ident is None:
-                ident = request.META.get("REMOTE_ADDR")
+            return request.user.id
+        ident = get_ip(request)
         return ident
 
     def get_cache_key(self, ident, scope, rate):
diff --git a/tests/unit/test_common_throttle.py b/tests/unit/test_common_throttle.py
index f5586f4d..03138557 100644
--- a/tests/unit/test_common_throttle.py
+++ b/tests/unit/test_common_throttle.py
@@ -263,3 +263,33 @@ def test_not_whitelisted_anon_throttling(settings, rf):
     cache.clear()
     settings.REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['anon-read'] = None
     settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST'] = []
+
+def test_whitelisted_subnet_anon_throttling(settings, rf):
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['anon-read'] = "1/min"
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST'] = ["192.168.0.0/24"]
+    request = rf.get("/test")
+    request.user = AnonymousUser()
+    request.META["REMOTE_ADDR"] = "192.168.0.123"
+    throttling = CommonThrottle()
+    assert throttling.allow_request(request, None)
+    assert throttling.allow_request(request, None)
+    for x in range(100):
+        assert throttling.allow_request(request, None)
+    cache.clear()
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['anon-read'] = None
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST'] = []
+
+def test_not_whitelisted_subnet_anon_throttling(settings, rf):
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['anon-read'] = "1/min"
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST'] = ["192.168.0.0/24"]
+    request = rf.get("/test")
+    request.user = AnonymousUser()
+    request.META["REMOTE_ADDR"] = "192.168.1.123"
+    throttling = CommonThrottle()
+    assert throttling.allow_request(request, None)
+    assert throttling.allow_request(request, None) is False
+    for x in range(100):
+        assert throttling.allow_request(request, None) is False
+    cache.clear()
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['anon-read'] = None
+    settings.REST_FRAMEWORK['DEFAULT_THROTTLE_WHITELIST'] = []