Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 578:a18516692273

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Correct my Redis API error in unit test.
author Brian Neal <bgneal@gmail.com>
date Sat, 05 May 2012 14:49:42 -0500
parents 6a265b5768ca
children
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
bgneal@508 10 from core.services import get_redis_connection
bgneal@508 11
bgneal@472 12
bgneal@472 13 logger = logging.getLogger(__name__)
bgneal@472 14
bgneal@472 15
bgneal@479 16 # This exception is thrown upon any Redis error. This insulates client code from
bgneal@479 17 # knowing that we are using Redis and will allow us to use something else in the
bgneal@479 18 # future.
bgneal@473 19 class RateLimiterUnavailable(Exception):
bgneal@473 20 pass
bgneal@473 21
bgneal@473 22
bgneal@472 23 def _make_key(ip):
bgneal@472 24 """
bgneal@472 25 Creates and returns a key string from a given IP address.
bgneal@472 26
bgneal@472 27 """
bgneal@472 28 return 'rate-limit-' + ip
bgneal@472 29
bgneal@472 30
bgneal@472 31 def _get_connection():
bgneal@472 32 """
bgneal@472 33 Create and return a Redis connection. Returns None on failure.
bgneal@472 34 """
bgneal@472 35 try:
bgneal@508 36 conn = get_redis_connection()
bgneal@472 37 except redis.RedisError, e:
bgneal@472 38 logger.error("rate limit: %s" % e)
bgneal@473 39 raise RateLimiterUnavailable
bgneal@472 40
bgneal@473 41 return conn
bgneal@472 42
bgneal@472 43
bgneal@472 44 def _to_seconds(interval):
bgneal@472 45 """
bgneal@472 46 Converts the timedelta interval object into a count of seconds.
bgneal@472 47
bgneal@472 48 """
bgneal@472 49 return interval.days * 24 * 3600 + interval.seconds
bgneal@472 50
bgneal@472 51
bgneal@472 52 def block_ip(ip, count=1000000, interval=datetime.timedelta(weeks=2)):
bgneal@472 53 """
bgneal@472 54 This function jams the rate limit record for the given IP so that the IP is
bgneal@472 55 blocked for the given interval. If the record doesn't exist, it is created.
bgneal@472 56 This is useful for manually blocking an IP after detecting suspicious
bgneal@472 57 behavior.
bgneal@473 58 This function may throw RateLimiterUnavailable.
bgneal@472 59
bgneal@472 60 """
bgneal@472 61 key = _make_key(ip)
bgneal@472 62 conn = _get_connection()
bgneal@472 63
bgneal@479 64 try:
bgneal@578 65 conn.setex(key, time=_to_seconds(interval), value=count)
bgneal@479 66 except redis.RedisError, e:
bgneal@479 67 logger.error("rate limit (block_ip): %s" % e)
bgneal@479 68 raise RateLimiterUnavailable
bgneal@479 69
bgneal@472 70 logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval)
bgneal@473 71
bgneal@473 72
bgneal@565 73 def unblock_ip(ip):
bgneal@565 74 """
bgneal@565 75 This function removes the block for the given IP address.
bgneal@565 76
bgneal@565 77 """
bgneal@565 78 key = _make_key(ip)
bgneal@565 79 conn = _get_connection()
bgneal@565 80 try:
bgneal@565 81 conn.delete(key)
bgneal@565 82 except redis.RedisError, e:
bgneal@565 83 logger.error("rate limit (unblock_ip): %s" % e)
bgneal@565 84 raise RateLimiterUnavailable
bgneal@565 85
bgneal@565 86 logger.info("Rate limiter unblocked IP %s", ip)
bgneal@565 87
bgneal@565 88
bgneal@473 89 class RateLimiter(object):
bgneal@473 90 """
bgneal@473 91 This class encapsulates the rate limiting logic for a given IP address.
bgneal@473 92
bgneal@473 93 """
bgneal@473 94 def __init__(self, ip, set_point, interval, lockout):
bgneal@473 95 self.ip = ip
bgneal@473 96 self.set_point = set_point
bgneal@473 97 self.interval = interval
bgneal@473 98 self.lockout = lockout
bgneal@473 99 self.key = _make_key(ip)
bgneal@473 100 self.conn = _get_connection()
bgneal@473 101
bgneal@473 102 def is_blocked(self):
bgneal@473 103 """
bgneal@473 104 Return True if the IP is blocked, and false otherwise.
bgneal@473 105
bgneal@473 106 """
bgneal@479 107 try:
bgneal@479 108 val = self.conn.get(self.key)
bgneal@479 109 except redis.RedisError, e:
bgneal@479 110 logger.error("RateLimiter (is_blocked): %s" % e)
bgneal@479 111 raise RateLimiterUnavailable
bgneal@479 112
bgneal@473 113 try:
bgneal@473 114 val = int(val) if val else 0
bgneal@473 115 except ValueError:
bgneal@473 116 return False
bgneal@473 117
bgneal@473 118 blocked = val >= self.set_point
bgneal@473 119 if blocked:
bgneal@473 120 logger.info("Rate limiter blocking %s", self.ip)
bgneal@473 121
bgneal@473 122 return blocked
bgneal@473 123
bgneal@473 124 def incr(self):
bgneal@473 125 """
bgneal@473 126 One is added to a counter associated with the IP address. If the
bgneal@473 127 counter exceeds set_point per interval, True is returned, and False
bgneal@473 128 otherwise. If the set_point is exceeded for the first time, the counter
bgneal@473 129 associated with the IP is set to expire according to the lockout
bgneal@473 130 parameter.
bgneal@473 131
bgneal@473 132 """
bgneal@479 133 try:
bgneal@479 134 val = self.conn.incr(self.key)
bgneal@473 135
bgneal@479 136 # Set expire time, if necessary.
bgneal@479 137 # If this is the first time, set it according to interval.
bgneal@479 138 # If the set_point has just been exceeded, set it according to lockout.
bgneal@479 139 if val == 1:
bgneal@479 140 self.conn.expire(self.key, _to_seconds(self.interval))
bgneal@479 141 elif val == self.set_point:
bgneal@479 142 self.conn.expire(self.key, _to_seconds(self.lockout))
bgneal@473 143
bgneal@479 144 tripped = val >= self.set_point
bgneal@473 145
bgneal@479 146 if tripped:
bgneal@479 147 logger.info("Rate limiter tripped for %s; counter = %d", self.ip, val)
bgneal@479 148 return tripped
bgneal@473 149
bgneal@479 150 except redis.RedisError, e:
bgneal@479 151 logger.error("RateLimiter (incr): %s" % e)
bgneal@479 152 raise RateLimiterUnavailable