Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 479:32cec6cd8808

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Refactor RateLimiter so that if Redis is not running, everything still runs normally (minus the rate limiting protection). My assumption that creating a Redis connection would throw an exception if Redis wasn't running was wrong. The exceptions actually occur when you issue a command. This is for #224.
author Brian Neal <bgneal@gmail.com>
date Sun, 25 Sep 2011 00:49:05 +0000
parents 5e826e232932
children 6f5fff924877
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@479 20 # This exception is thrown upon any Redis error. This insulates client code from
bgneal@479 21 # knowing that we are using Redis and will allow us to use something else in the
bgneal@479 22 # future.
bgneal@473 23 class RateLimiterUnavailable(Exception):
bgneal@473 24 pass
bgneal@473 25
bgneal@473 26
bgneal@472 27 def _make_key(ip):
bgneal@472 28 """
bgneal@472 29 Creates and returns a key string from a given IP address.
bgneal@472 30
bgneal@472 31 """
bgneal@472 32 return 'rate-limit-' + ip
bgneal@472 33
bgneal@472 34
bgneal@472 35 def _get_connection():
bgneal@472 36 """
bgneal@472 37 Create and return a Redis connection. Returns None on failure.
bgneal@472 38 """
bgneal@472 39 try:
bgneal@472 40 conn = redis.Redis(host=HOST, port=PORT, db=DB)
bgneal@472 41 except redis.RedisError, e:
bgneal@472 42 logger.error("rate limit: %s" % e)
bgneal@473 43 raise RateLimiterUnavailable
bgneal@472 44
bgneal@473 45 return conn
bgneal@472 46
bgneal@472 47
bgneal@472 48 def _to_seconds(interval):
bgneal@472 49 """
bgneal@472 50 Converts the timedelta interval object into a count of seconds.
bgneal@472 51
bgneal@472 52 """
bgneal@472 53 return interval.days * 24 * 3600 + interval.seconds
bgneal@472 54
bgneal@472 55
bgneal@472 56 def block_ip(ip, count=1000000, interval=datetime.timedelta(weeks=2)):
bgneal@472 57 """
bgneal@472 58 This function jams the rate limit record for the given IP so that the IP is
bgneal@472 59 blocked for the given interval. If the record doesn't exist, it is created.
bgneal@472 60 This is useful for manually blocking an IP after detecting suspicious
bgneal@472 61 behavior.
bgneal@473 62 This function may throw RateLimiterUnavailable.
bgneal@472 63
bgneal@472 64 """
bgneal@472 65 key = _make_key(ip)
bgneal@472 66 conn = _get_connection()
bgneal@472 67
bgneal@479 68 try:
bgneal@479 69 conn.setex(key, count, _to_seconds(interval))
bgneal@479 70 except redis.RedisError, e:
bgneal@479 71 logger.error("rate limit (block_ip): %s" % e)
bgneal@479 72 raise RateLimiterUnavailable
bgneal@479 73
bgneal@472 74 logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval)
bgneal@473 75
bgneal@473 76
bgneal@473 77 class RateLimiter(object):
bgneal@473 78 """
bgneal@473 79 This class encapsulates the rate limiting logic for a given IP address.
bgneal@473 80
bgneal@473 81 """
bgneal@473 82 def __init__(self, ip, set_point, interval, lockout):
bgneal@473 83 self.ip = ip
bgneal@473 84 self.set_point = set_point
bgneal@473 85 self.interval = interval
bgneal@473 86 self.lockout = lockout
bgneal@473 87 self.key = _make_key(ip)
bgneal@473 88 self.conn = _get_connection()
bgneal@473 89
bgneal@473 90 def is_blocked(self):
bgneal@473 91 """
bgneal@473 92 Return True if the IP is blocked, and false otherwise.
bgneal@473 93
bgneal@473 94 """
bgneal@479 95 try:
bgneal@479 96 val = self.conn.get(self.key)
bgneal@479 97 except redis.RedisError, e:
bgneal@479 98 logger.error("RateLimiter (is_blocked): %s" % e)
bgneal@479 99 raise RateLimiterUnavailable
bgneal@479 100
bgneal@473 101 try:
bgneal@473 102 val = int(val) if val else 0
bgneal@473 103 except ValueError:
bgneal@473 104 return False
bgneal@473 105
bgneal@473 106 blocked = val >= self.set_point
bgneal@473 107 if blocked:
bgneal@473 108 logger.info("Rate limiter blocking %s", self.ip)
bgneal@473 109
bgneal@473 110 return blocked
bgneal@473 111
bgneal@473 112 def incr(self):
bgneal@473 113 """
bgneal@473 114 One is added to a counter associated with the IP address. If the
bgneal@473 115 counter exceeds set_point per interval, True is returned, and False
bgneal@473 116 otherwise. If the set_point is exceeded for the first time, the counter
bgneal@473 117 associated with the IP is set to expire according to the lockout
bgneal@473 118 parameter.
bgneal@473 119
bgneal@473 120 """
bgneal@479 121 try:
bgneal@479 122 val = self.conn.incr(self.key)
bgneal@473 123
bgneal@479 124 # Set expire time, if necessary.
bgneal@479 125 # If this is the first time, set it according to interval.
bgneal@479 126 # If the set_point has just been exceeded, set it according to lockout.
bgneal@479 127 if val == 1:
bgneal@479 128 self.conn.expire(self.key, _to_seconds(self.interval))
bgneal@479 129 elif val == self.set_point:
bgneal@479 130 self.conn.expire(self.key, _to_seconds(self.lockout))
bgneal@473 131
bgneal@479 132 tripped = val >= self.set_point
bgneal@473 133
bgneal@479 134 if tripped:
bgneal@479 135 logger.info("Rate limiter tripped for %s; counter = %d", self.ip, val)
bgneal@479 136 return tripped
bgneal@473 137
bgneal@479 138 except redis.RedisError, e:
bgneal@479 139 logger.error("RateLimiter (incr): %s" % e)
bgneal@479 140 raise RateLimiterUnavailable