Add codeowners file

.gitignore

@@ -4,3 +4,4 @@ dist
 *.egg-info
 docs/_build/*
 .tox/
+.idea/

CODEOWNERS

@@ -0,0 +1 @@
+.fabric/ @unacademy/sre

ratelimit/__init__.py

@@ -1,4 +1,4 @@
-VERSION = (1, 1, 0)
+VERSION = (1, 3, 7)
 __version__ = '.'.join(map(str, VERSION))
 
 ALL = (None,)  # Sentinel value for all HTTP methods.

ratelimit/decorators.py

@@ -12,7 +12,7 @@
 __all__ = ['ratelimit']
 
 
-def ratelimit(group=None, key=None, rate=None, method=ALL, block=False):
+def ratelimit(group=None, key=None, rate=None, method=ALL, block=False, reset=None):
     def decorator(fn):
         @wraps(fn)
         def _wrapped(*args, **kw):
@@ -24,9 +24,9 @@ def _wrapped(*args, **kw):
             request.limited = getattr(request, 'limited', False)
             ratelimited = is_ratelimited(request=request, group=group, fn=fn,
                                          key=key, rate=rate, method=method,
-                                         increment=True)
+                                         increment=True, reset=reset)
             if ratelimited and block:
-                raise Ratelimited()
+                raise Ratelimited("Too many requests", 429)
             return fn(*args, **kw)
         return _wrapped
     return decorator

ratelimit/exceptions.py

@@ -3,3 +3,20 @@
 
 class Ratelimited(PermissionDenied):
     pass
+
+
+class CustomException(Exception):
+    def __init__(self, message):
+        super(CustomException, self).__init__(message)
+
+
+class DatastoreConnectionError(CustomException):
+    def __init__(self):
+        message = 'Could not establish connection to data store'
+        super(DatastoreConnectionError, self).__init__(message)
+
+
+class RateLimited(CustomException):
+    def __init__(self):
+        message = 'Rate limited'
+        super(RateLimited, self).__init__(message)

ratelimit/redis_rate_limit.py

@@ -0,0 +1,141 @@
+import time
+import redis
+
+from functools import wraps
+from django.conf import settings
+from rediscluster import StrictRedisCluster
+from .exceptions import RateLimited, DatastoreConnectionError
+
+__author__ = 'vikaschahal'
+
+
+class RateLimiter(object):
+    """Base class for Rate Limiting"""
+
+    def __init__(self, limit, window, connection, key):
+        """
+        :param limit: number of requests allowed
+        :type limit: int
+        :param window: window in secs in which :limit number requests allowed
+        :type window: int
+        """
+        self._connection = connection
+        self._key = key
+        self._limit = limit
+        self._window = window
+
+    def is_allowed(self, log_current_request=True):
+        """
+        :param log_current_request: Consider the call for rate limiting
+        :type log_current_request: bool
+        :return: Whether a requests is allowed or rate limited.
+        :rtype: bool
+        """
+        raise NotImplementedError
+
+    @property
+    def remaining_requests(self):
+        raise NotImplementedError
+
+    def limit(self, func):
+        """Decorator to check the rate limit."""
+
+        @wraps(func)
+        def decorated(*args, **kwargs):
+            if self.is_allowed():
+                return func(*args, **kwargs)
+            raise RateLimited
+
+        return decorated
+
+
+class RedisRateLimiterConnection(object):
+    def __init__(self, host=None, port=None, db=0, connection=None):
+        self.connection = None
+        if host:
+            if settings.REDIS_HOST_INTERNAL_NEW_IS_CLUSTER:
+                connection = StrictRedisCluster(startup_nodes=[{'host': host, 'port': 6379}], skip_full_coverage_check=True)
+            else:
+                connection = redis.StrictRedis(host, port, db)
+            if not connection.ping():
+                raise DatastoreConnectionError
+            self.connection = connection
+        elif connection:
+            if not connection.ping():
+                raise DatastoreConnectionError
+            self.connection = connection
+        else:
+            raise DatastoreConnectionError
+
+
+class RedisRateLimiter(RateLimiter):
+    def __init__(self, limit, window, connection, key):
+        super(RedisRateLimiter, self).__init__(limit, window, connection, key)
+        self._pipeline = self._connection.connection.pipeline()
+
+    def _increment_request(self):
+        key_value = int(time.time()) + self._window
+        self._pipeline.zadd(
+            self._key, key_value, key_value
+        )
+        self._pipeline.expire(self._key, self._window) # set key expiry
+        self._pipeline.execute()
+
+    def is_allowed(self, log_current_request=True):
+        if log_current_request:
+            self._increment_request()
+        current_time = time.time()
+        self._pipeline.zremrangebyscore(self._key, '-inf', current_time)
+        self._pipeline.zcount(self._key, '-inf', '+inf')
+        result = self._pipeline.execute()
+        return result[-1] <= self._limit
+
+    def count(self, log_current_request=True):
+        if log_current_request:
+            self._increment_request()
+        current_time = time.time()
+        self._pipeline.zremrangebyscore(self._key, '-inf', current_time)
+        self._pipeline.zcount(self._key, '-inf', '+inf')
+        result = self._pipeline.execute()
+        return result[-1]
+
+
+class IpRateLimiter(RateLimiter):
+    def __init__(self, limit, window, connection, key):
+        super(IpRateLimiter, self).__init__(limit, window, connection, key)
+        self._pipeline = self._connection.connection.pipeline()
+
+    def add(self, value):
+        self._pipeline.exists(self._key)
+        results = self._pipeline.execute()
+        self._pipeline.sadd(self._key, value)
+        if not results[0]:
+            self._pipeline.expire(self._key, self._window)
+        self._pipeline.execute()
+
+    def count(self):
+        self._pipeline.scard(self._key)
+        result = self._pipeline.execute()
+        return result[0]
+
+    def delete(self):
+        self._pipeline.delete(self._key)
+        self._pipeline.execute()
+
+    def is_allowed(self):
+        self._pipeline.exists(self._key)
+        self._pipeline.scard(self._key)
+        results = self._pipeline.execute()
+        if not results[0]:
+            return True
+        return results[1] < self._limit
+
+    @property
+    def remaining_requests(self):
+        return self._limit - self.count()
+
+
+if settings.REDIS_HOST_INTERNAL_NEW_IS_CLUSTER:
+    redis_connection = RedisRateLimiterConnection(host=settings.REDIS_HOST_INTERNAL_NEW)
+else:
+    redis_connection = RedisRateLimiterConnection(host=settings.REDIS_HOST_INTERNAL, port=6379, db=0)