Mercurial > public > sg101

annotate antispam/rate_limit.py @ 627:a4300639c6e7

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Wiki integration. Create task to delete old cookie records. Rework logic upon logout, as session will not be available. Set an attribute on the request instead.
author Brian Neal <bgneal@gmail.com>
date Mon, 12 Nov 2012 15:10:52 -0600
parents ee87ea74d46b
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