Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 496:61b81b97f02f

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Backout last change. Stay xhtml for now.
author Brian Neal <bgneal@gmail.com>
date Sun, 06 Nov 2011 01:14:26 +0000
parents 32cec6cd8808
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