
import time
import redis

from django.http import HttpResponse

from . import settings

# Connection pool
POOL = redis.ConnectionPool(**settings.REDIS)

class RateLimitBackend(object):

    def source_for_request(self, request):
        '''Return a source identifier for a request'''
        if getattr(settings, 'USE_X_FORWARDED_FOR', False):
            x_forwarded_for = request.META.get('HTTP_X_FORWARDED_FOR')
            if x_forwarded_for:
                return x_forwarded_for.split(',')[0]

        return request.META.get('REMOTE_ADDR')

    def make_limit_response(self, realm):
        conf = settings.REALMS.get(realm, {})

        return HttpResponse(conf.get('message', settings.RESPONSE_MESSAGE),
            status=conf.get('code', settings.RESPONSE_CODE)
        )

    def check_realm(self, source, realm):
        '''
        Check if `source` has reached their limit in `realm`.

        Returns True if limit is reached.
        '''
        conf = settings.REALMS.get(realm, {})

        # Check against Realm whitelist
        if source in conf.get('whitelist', settings.DEFAULT_WHITELIST):
            return None

        key = 'rated:%s:%s' % (realm, source,)
        now = time.time()

        client = redis.Redis(connection_pool=POOL)
        # Do commands at once for speed
        # We don't need these to operate in a transaction, as none of the values
        # we send are dependant on values in the DB
        pipe = client.pipeline(transaction=False)
        # Add our timestamp to the range
        pipe.zadd(key, now, now)
        # Update to not expire for another DURATION
        pipe.expireat(key, int(now + conf.get('duration', settings.DEFAULT_DURATION)))
        # Remove old values
        pipe.zremrangebyscore(key, '-inf', now - settings.DEFAULT_DURATION)
        # Test count
        pipe.zcard(key)
        size = pipe.execute()[-1]
        return size > conf.get('limit', settings.DEFAULT_LIMIT)

BACKEND = RateLimitBackend()


@@ -147,16 +147,17 @@
 REDIS)%0A%0A
+%0A
 class Ra
@@ -709,16 +709,32 @@
 SSAGE),%0A
+                
         
@@ -784,25 +784,16 @@
 SE_CODE)
-%0A        
 )%0A%0A    d
@@ -1380,23 +1380,16 @@
 e of the
- values
 %0A       
@@ -1390,16 +1390,23 @@
        #
+ values
  we send
@@ -1452,14 +1452,12 @@
     
-pipe =
+with
  cli
@@ -1487,25 +1487,37 @@
 n=False)
-%0A
+ as pipe:%0A   
         
 # Add ou
@@ -1508,16 +1508,17 @@
         
+ 
 # Add ou
@@ -1546,24 +1546,28 @@
 nge%0A        
+    
 pipe.zadd(ke
@@ -1579,16 +1579,20 @@
 w, now)%0A
+    
         
@@ -1639,24 +1639,28 @@
 ION%0A        
+    
 pipe.expirea
@@ -1726,16 +1726,20 @@
 TION)))%0A
+    
         
@@ -1762,24 +1762,28 @@
 ues%0A        
+    
 pipe.zremran
@@ -1846,16 +1846,20 @@
         
+    
 # Test c
@@ -1867,24 +1867,28 @@
 unt%0A        
+  
+  
 pipe.zcard(k
@@ -1891,16 +1891,20 @@
 rd(key)%0A
+    
         
@@ -1994,16 +1994,16 @@
 LIMIT)%0A%0A
+
 BACKEND 
@@ -2023,9 +2023,8 @@
 ckend()%0A
-%0A

Use 'with pipeline' to ensure it's cleaned up properly