Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 477:737a0c261a77

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add a link to John Blair's 50 years of surf music article on the banner.
author Brian Neal <bgneal@gmail.com>
date Fri, 09 Sep 2011 22:16:40 +0000
parents 5e826e232932
children 32cec6cd8808
rev   line source
bgneal@472 1 """
bgneal@472 2 This module contains the rate limiting functionality.
bgneal@472 3
bgneal@472 4 """
bgneal@472 5 import datetime
bgneal@472 6 import logging
bgneal@472 7
bgneal@472 8 import redis
bgneal@472 9 from django.conf import settings
bgneal@472 10
bgneal@472 11
bgneal@472 12 logger = logging.getLogger(__name__)
bgneal@472 13
bgneal@472 14 # Redis connection and database settings
bgneal@472 15 HOST = getattr(settings, 'RATE_LIMIT_REDIS_HOST', 'localhost')
bgneal@472 16 PORT = getattr(settings, 'RATE_LIMIT_REDIS_PORT', 6379)
bgneal@472 17 DB = getattr(settings, 'RATE_LIMIT_REDIS_DB', 0)
bgneal@472 18
bgneal@472 19
bgneal@473 20 class RateLimiterUnavailable(Exception):
bgneal@473 21 pass
bgneal@473 22
bgneal@473 23
bgneal@472 24 def _make_key(ip):
bgneal@472 25 """
bgneal@472 26 Creates and returns a key string from a given IP address.
bgneal@472 27
bgneal@472 28 """
bgneal@472 29 return 'rate-limit-' + ip
bgneal@472 30
bgneal@472 31
bgneal@472 32 def _get_connection():
bgneal@472 33 """
bgneal@472 34 Create and return a Redis connection. Returns None on failure.
bgneal@472 35 """
bgneal@472 36 try:
bgneal@472 37 conn = redis.Redis(host=HOST, port=PORT, db=DB)
bgneal@472 38 except redis.RedisError, e:
bgneal@472 39 logger.error("rate limit: %s" % e)
bgneal@473 40 raise RateLimiterUnavailable
bgneal@472 41
bgneal@473 42 return conn
bgneal@472 43
bgneal@472 44
bgneal@472 45 def _to_seconds(interval):
bgneal@472 46 """
bgneal@472 47 Converts the timedelta interval object into a count of seconds.
bgneal@472 48
bgneal@472 49 """
bgneal@472 50 return interval.days * 24 * 3600 + interval.seconds
bgneal@472 51
bgneal@472 52
bgneal@472 53 def block_ip(ip, count=1000000, interval=datetime.timedelta(weeks=2)):
bgneal@472 54 """
bgneal@472 55 This function jams the rate limit record for the given IP so that the IP is
bgneal@472 56 blocked for the given interval. If the record doesn't exist, it is created.
bgneal@472 57 This is useful for manually blocking an IP after detecting suspicious
bgneal@472 58 behavior.
bgneal@473 59 This function may throw RateLimiterUnavailable.
bgneal@472 60
bgneal@472 61 """
bgneal@472 62 key = _make_key(ip)
bgneal@472 63 conn = _get_connection()
bgneal@472 64
bgneal@472 65 conn.setex(key, count, _to_seconds(interval))
bgneal@472 66 logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval)
bgneal@473 67
bgneal@473 68
bgneal@473 69 class RateLimiter(object):
bgneal@473 70 """
bgneal@473 71 This class encapsulates the rate limiting logic for a given IP address.
bgneal@473 72
bgneal@473 73 """
bgneal@473 74 def __init__(self, ip, set_point, interval, lockout):
bgneal@473 75 self.ip = ip
bgneal@473 76 self.set_point = set_point
bgneal@473 77 self.interval = interval
bgneal@473 78 self.lockout = lockout
bgneal@473 79 self.key = _make_key(ip)
bgneal@473 80 self.conn = _get_connection()
bgneal@473 81
bgneal@473 82 def is_blocked(self):
bgneal@473 83 """
bgneal@473 84 Return True if the IP is blocked, and false otherwise.
bgneal@473 85
bgneal@473 86 """
bgneal@473 87 val = self.conn.get(self.key)
bgneal@473 88 try:
bgneal@473 89 val = int(val) if val else 0
bgneal@473 90 except ValueError:
bgneal@473 91 return False
bgneal@473 92
bgneal@473 93 blocked = val >= self.set_point
bgneal@473 94 if blocked:
bgneal@473 95 logger.info("Rate limiter blocking %s", self.ip)
bgneal@473 96
bgneal@473 97 return blocked
bgneal@473 98
bgneal@473 99 def incr(self):
bgneal@473 100 """
bgneal@473 101 One is added to a counter associated with the IP address. If the
bgneal@473 102 counter exceeds set_point per interval, True is returned, and False
bgneal@473 103 otherwise. If the set_point is exceeded for the first time, the counter
bgneal@473 104 associated with the IP is set to expire according to the lockout
bgneal@473 105 parameter.
bgneal@473 106
bgneal@473 107 """
bgneal@473 108 val = self.conn.incr(self.key)
bgneal@473 109
bgneal@473 110 # Set expire time, if necessary.
bgneal@473 111 # If this is the first time, set it according to interval.
bgneal@473 112 # If the set_point has just been exceeded, set it according to lockout.
bgneal@473 113 if val == 1:
bgneal@473 114 self.conn.expire(self.key, _to_seconds(self.interval))
bgneal@473 115 elif val == self.set_point:
bgneal@473 116 self.conn.expire(self.key, _to_seconds(self.lockout))
bgneal@473 117
bgneal@473 118 tripped = val >= self.set_point
bgneal@473 119
bgneal@473 120 if tripped:
bgneal@473 121 logger.info("Rate limiter tripped for %s; counter = %d", self.ip, val)
bgneal@473 122
bgneal@473 123 return tripped