Mercurial > public > sg101
comparison gpp/antispam/rate_limit.py @ 472:7c3816d76c6c
Implement rate limiting on registration and login for #224.
| author | Brian Neal <bgneal@gmail.com> |
|---|---|
| date | Thu, 25 Aug 2011 02:23:55 +0000 |
| parents | |
| children | 5e826e232932 |
comparison
equal
deleted
inserted
replaced
| 471:d83296cac940 | 472:7c3816d76c6c |
|---|---|
| 1 """ | |
| 2 This module contains the rate limiting functionality. | |
| 3 | |
| 4 """ | |
| 5 import datetime | |
| 6 import logging | |
| 7 | |
| 8 import redis | |
| 9 from django.conf import settings | |
| 10 | |
| 11 | |
| 12 logger = logging.getLogger(__name__) | |
| 13 | |
| 14 # Redis connection and database settings | |
| 15 HOST = getattr(settings, 'RATE_LIMIT_REDIS_HOST', 'localhost') | |
| 16 PORT = getattr(settings, 'RATE_LIMIT_REDIS_PORT', 6379) | |
| 17 DB = getattr(settings, 'RATE_LIMIT_REDIS_DB', 0) | |
| 18 | |
| 19 | |
| 20 def _make_key(ip): | |
| 21 """ | |
| 22 Creates and returns a key string from a given IP address. | |
| 23 | |
| 24 """ | |
| 25 return 'rate-limit-' + ip | |
| 26 | |
| 27 | |
| 28 def _get_connection(): | |
| 29 """ | |
| 30 Create and return a Redis connection. Returns None on failure. | |
| 31 """ | |
| 32 try: | |
| 33 conn = redis.Redis(host=HOST, port=PORT, db=DB) | |
| 34 return conn | |
| 35 except redis.RedisError, e: | |
| 36 logger.error("rate limit: %s" % e) | |
| 37 | |
| 38 return None | |
| 39 | |
| 40 | |
| 41 def _to_seconds(interval): | |
| 42 """ | |
| 43 Converts the timedelta interval object into a count of seconds. | |
| 44 | |
| 45 """ | |
| 46 return interval.days * 24 * 3600 + interval.seconds | |
| 47 | |
| 48 | |
| 49 def rate_check(ip, set_point, interval, lockout): | |
| 50 """ | |
| 51 This function performs a rate limit check. | |
| 52 One is added to a counter associated with the IP address. If the | |
| 53 counter exceeds set_point per interval, True is returned, and False | |
| 54 otherwise. If the set_point is exceeded for the first time, the counter | |
| 55 associated with the IP is set to expire according to the lockout parameter. | |
| 56 This locks the IP address as this function will then return True for the | |
| 57 period specified by lockout. | |
| 58 | |
| 59 """ | |
| 60 if not ip: | |
| 61 logger.error("rate_limit.rate_check could not get IP") | |
| 62 return False | |
| 63 key = _make_key(ip) | |
| 64 | |
| 65 conn = _get_connection() | |
| 66 if not conn: | |
| 67 return False | |
| 68 | |
| 69 val = conn.incr(key) | |
| 70 | |
| 71 # Set expire time, if necessary. | |
| 72 # If this is the first time, set it according to interval. | |
| 73 # If the set_point has just been exceeded, set it according to lockout. | |
| 74 if val == 1: | |
| 75 conn.expire(key, _to_seconds(interval)) | |
| 76 elif val == set_point: | |
| 77 conn.expire(key, _to_seconds(lockout)) | |
| 78 | |
| 79 tripped = val >= set_point | |
| 80 | |
| 81 if tripped: | |
| 82 logger.info("Rate limiter tripped for %s; counter = %d", ip, val) | |
| 83 | |
| 84 return tripped | |
| 85 | |
| 86 | |
| 87 def block_ip(ip, count=1000000, interval=datetime.timedelta(weeks=2)): | |
| 88 """ | |
| 89 This function jams the rate limit record for the given IP so that the IP is | |
| 90 blocked for the given interval. If the record doesn't exist, it is created. | |
| 91 This is useful for manually blocking an IP after detecting suspicious | |
| 92 behavior. | |
| 93 | |
| 94 """ | |
| 95 if not ip: | |
| 96 logger.error("rate_limit.block_ip could not get IP") | |
| 97 return | |
| 98 | |
| 99 key = _make_key(ip) | |
| 100 conn = _get_connection() | |
| 101 if not conn: | |
| 102 return | |
| 103 | |
| 104 conn.setex(key, count, _to_seconds(interval)) | |
| 105 logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval) |