Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 546:0d007c8c4851

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Correct the hg pull command in the fabfile.
author Brian Neal <bgneal@gmail.com>
date Fri, 13 Jan 2012 19:09:08 -0600
parents 6f5fff924877
children 6a265b5768ca
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@508 11 from core.services import get_redis_connection
bgneal@508 12
bgneal@472 13
bgneal@472 14 logger = logging.getLogger(__name__)
bgneal@472 15
bgneal@472 16
bgneal@479 17 # This exception is thrown upon any Redis error. This insulates client code from
bgneal@479 18 # knowing that we are using Redis and will allow us to use something else in the
bgneal@479 19 # future.
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@508 37 conn = get_redis_connection()
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@479 65 try:
bgneal@479 66 conn.setex(key, count, _to_seconds(interval))
bgneal@479 67 except redis.RedisError, e:
bgneal@479 68 logger.error("rate limit (block_ip): %s" % e)
bgneal@479 69 raise RateLimiterUnavailable
bgneal@479 70
bgneal@472 71 logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval)
bgneal@473 72
bgneal@473 73
bgneal@473 74 class RateLimiter(object):
bgneal@473 75 """
bgneal@473 76 This class encapsulates the rate limiting logic for a given IP address.
bgneal@473 77
bgneal@473 78 """
bgneal@473 79 def __init__(self, ip, set_point, interval, lockout):
bgneal@473 80 self.ip = ip
bgneal@473 81 self.set_point = set_point
bgneal@473 82 self.interval = interval
bgneal@473 83 self.lockout = lockout
bgneal@473 84 self.key = _make_key(ip)
bgneal@473 85 self.conn = _get_connection()
bgneal@473 86
bgneal@473 87 def is_blocked(self):
bgneal@473 88 """
bgneal@473 89 Return True if the IP is blocked, and false otherwise.
bgneal@473 90
bgneal@473 91 """
bgneal@479 92 try:
bgneal@479 93 val = self.conn.get(self.key)
bgneal@479 94 except redis.RedisError, e:
bgneal@479 95 logger.error("RateLimiter (is_blocked): %s" % e)
bgneal@479 96 raise RateLimiterUnavailable
bgneal@479 97
bgneal@473 98 try:
bgneal@473 99 val = int(val) if val else 0
bgneal@473 100 except ValueError:
bgneal@473 101 return False
bgneal@473 102
bgneal@473 103 blocked = val >= self.set_point
bgneal@473 104 if blocked:
bgneal@473 105 logger.info("Rate limiter blocking %s", self.ip)
bgneal@473 106
bgneal@473 107 return blocked
bgneal@473 108
bgneal@473 109 def incr(self):
bgneal@473 110 """
bgneal@473 111 One is added to a counter associated with the IP address. If the
bgneal@473 112 counter exceeds set_point per interval, True is returned, and False
bgneal@473 113 otherwise. If the set_point is exceeded for the first time, the counter
bgneal@473 114 associated with the IP is set to expire according to the lockout
bgneal@473 115 parameter.
bgneal@473 116
bgneal@473 117 """
bgneal@479 118 try:
bgneal@479 119 val = self.conn.incr(self.key)
bgneal@473 120
bgneal@479 121 # Set expire time, if necessary.
bgneal@479 122 # If this is the first time, set it according to interval.
bgneal@479 123 # If the set_point has just been exceeded, set it according to lockout.
bgneal@479 124 if val == 1:
bgneal@479 125 self.conn.expire(self.key, _to_seconds(self.interval))
bgneal@479 126 elif val == self.set_point:
bgneal@479 127 self.conn.expire(self.key, _to_seconds(self.lockout))
bgneal@473 128
bgneal@479 129 tripped = val >= self.set_point
bgneal@473 130
bgneal@479 131 if tripped:
bgneal@479 132 logger.info("Rate limiter tripped for %s; counter = %d", self.ip, val)
bgneal@479 133 return tripped
bgneal@473 134
bgneal@479 135 except redis.RedisError, e:
bgneal@479 136 logger.error("RateLimiter (incr): %s" % e)
bgneal@479 137 raise RateLimiterUnavailable