Mercurial > public > sg101

comparison antispam/rate_limit.py @ 581:ee87ea74d46b

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