feat: Ignore Rate Limit for Whitelisted Clients

redis_rate_limit/__init__.py

@@ -43,21 +43,23 @@ class RateLimit(object):
     This class offers an abstraction of a Rate Limit algorithm implemented on
     top of Redis >= 2.6.0.
     """
-    def __init__(self, resource, client, max_requests, expire=None, redis_pool=REDIS_POOL):
+    def __init__(self, resource, client, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL):
         """
         Class initialization method checks if the Rate Limit algorithm is
         actually supported by the installed Redis version and sets some
         useful properties.
-
         If Rate Limit is not supported, it raises an Exception.
-
         :param resource: resource identifier string (i.e. ‘user_pictures’)
         :param client: client identifier string (i.e. ‘192.168.0.10’)
         :param max_requests: integer (i.e. ‘10’)
+        :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1'])
         :param expire: seconds to wait before resetting counters (i.e. ‘60’)
         :param redis_pool: instance of redis.ConnectionPool.
                Default: ConnectionPool(host='127.0.0.1', port=6379, db=0)
         """
+        self.client = client
+        self.whitelisted_clients = whitelisted_clients
+
         self._redis = Redis(connection_pool=redis_pool)
         if not self._is_rate_limit_supported():
             raise RedisVersionNotSupported()
@@ -88,7 +90,6 @@ def get_usage(self):
         """
         Returns actual resource usage by client. Note that it could be greater
         than the maximum number of requests set.
-
         :return: integer: current usage
         """
         return int(self._redis.get(self._rate_limit_key) or 0)
@@ -97,7 +98,6 @@ def get_wait_time(self):
         """
         Returns estimated optimal wait time for subsequent requests.
         If limit has already been reached, return wait time until it gets reset.
-
         :return: float: wait time in seconds
         """
         expire = self._redis.pttl(self._rate_limit_key)
@@ -111,22 +111,18 @@ def get_wait_time(self):
     def has_been_reached(self):
         """
         Checks if Rate Limit has been reached.
-
         :return: bool: True if limit has been reached or False otherwise
         """
         return self.get_usage() >= self._max_requests
 
     def increment_usage(self, increment_by=1):
         """
         Calls a LUA script that should increment the resource usage by client.
-
         If the resource limit overflows the maximum number of requests, this
         method raises an Exception.
-
         :param increment_by: The count to increment the rate limiter by.
         This is by default 1, but higher values are provided for more flexible
         rate-limiting schemes.
-
         :return: integer: current usage
         """
         if increment_by > self._max_requests:
@@ -140,8 +136,11 @@ def increment_usage(self, increment_by=1):
                              .format(increment_by=increment_by))
 
         try:
-            current_usage = self._redis.evalsha(
-                INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by)
+            if self.client not in self.whitelisted_clients:
+                current_usage = self._redis.evalsha(
+                    INCREMENT_SCRIPT_HASH, 1, self._rate_limit_key, self._expire, increment_by)
+            else:
+                current_usage = 0
         except NoScriptError:
             current_usage = self._redis.eval(
                 INCREMENT_SCRIPT, 1, self._rate_limit_key, self._expire, increment_by)
@@ -155,7 +154,6 @@ def _is_rate_limit_supported(self):
         """
         Checks if Rate Limit is supported which can basically be found by
         looking at Redis database version that should be 2.6.0 or greater.
-
         :return: bool
         """
         redis_version = self._redis.info()['redis_version']
@@ -172,17 +170,19 @@ def _reset(self):
 
 
 class RateLimiter(object):
-    def __init__(self, resource, max_requests, expire=None, redis_pool=REDIS_POOL):
+    def __init__(self, resource, max_requests, whitelisted_clients=[], expire=None, redis_pool=REDIS_POOL):
         """
         Rate limit factory. Checks if RateLimit is supported when limit is called.
         :param resource: resource identifier string (i.e. ‘user_pictures’)
         :param max_requests: integer (i.e. ‘10’)
+        :param whitelisted_clients: list of ip addresses (i.e. ['127.0.0.1'])
         :param expire: seconds to wait before resetting counters (i.e. ‘60’)
         :param redis_pool: instance of redis.ConnectionPool.
                Default: ConnectionPool(host='127.0.0.1', port=6379, db=0)
        """
         self.resource = resource
         self.max_requests = max_requests
+        self.whitelisted_clients = whitelisted_clients
         self.expire = expire
         self.redis_pool = redis_pool
 
@@ -194,6 +194,7 @@ def limit(self, client):
             resource=self.resource,
             client=client,
             max_requests=self.max_requests,
+            whitelisted_clients=self.whitelisted_clients,
             expire=self.expire,
             redis_pool=self.redis_pool,
-        )
+        )

tests/rate_limit_test.py

@@ -41,6 +41,19 @@ def test_limit_10_max_request(self):
         self.assertEqual(self.rate_limit.get_usage(), 11)
         self.assertEqual(self.rate_limit.has_been_reached(), True)
 
+    def test_whitelisted_clients(self):
+        """
+        Should not increment counter if client is part of whitelisted_clients list.
+        """
+        self.rate_limit = RateLimit(resource='test', client='localhost', whitelisted_clients=['localhost'],
+                                    max_requests=10, expire=2)
+        self.assertEqual(self.rate_limit.get_usage(), 0)
+        self.assertEqual(self.rate_limit.has_been_reached(), False)
+
+        self._make_10_requests()
+        self.assertEqual(self.rate_limit.get_usage(), 0)
+        self.assertEqual(self.rate_limit.has_been_reached(), False)
+
     def test_expire(self):
         """
         Should not raise TooManyRequests Exception when trying to increment for