Mercurial > public > sg101

annotate gpp/antispam/rate_limit.py @ 472:7c3816d76c6c

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