bgneal@472: """
bgneal@472: This module contains the rate limiting functionality.
bgneal@472: 
bgneal@472: """
bgneal@472: import datetime
bgneal@472: import logging
bgneal@472: 
bgneal@472: import redis
bgneal@472: from django.conf import settings
bgneal@472: 
bgneal@472: 
bgneal@472: logger = logging.getLogger(__name__)
bgneal@472: 
bgneal@472: # Redis connection and database settings
bgneal@472: HOST = getattr(settings, 'RATE_LIMIT_REDIS_HOST', 'localhost')
bgneal@472: PORT = getattr(settings, 'RATE_LIMIT_REDIS_PORT', 6379)
bgneal@472: DB = getattr(settings, 'RATE_LIMIT_REDIS_DB', 0)
bgneal@472: 
bgneal@472: 
bgneal@472: def _make_key(ip):
bgneal@472:     """
bgneal@472:     Creates and returns a key string from a given IP address.
bgneal@472: 
bgneal@472:     """
bgneal@472:     return 'rate-limit-' + ip
bgneal@472: 
bgneal@472: 
bgneal@472: def _get_connection():
bgneal@472:     """
bgneal@472:     Create and return a Redis connection. Returns None on failure.
bgneal@472:     """
bgneal@472:     try:
bgneal@472:         conn = redis.Redis(host=HOST, port=PORT, db=DB)
bgneal@472:         return conn
bgneal@472:     except redis.RedisError, e:
bgneal@472:         logger.error("rate limit: %s" % e)
bgneal@472: 
bgneal@472:     return None
bgneal@472: 
bgneal@472: 
bgneal@472: def _to_seconds(interval):
bgneal@472:     """
bgneal@472:     Converts the timedelta interval object into a count of seconds.
bgneal@472: 
bgneal@472:     """
bgneal@472:     return interval.days * 24 * 3600 + interval.seconds
bgneal@472: 
bgneal@472: 
bgneal@472: def rate_check(ip, set_point, interval, lockout):
bgneal@472:     """
bgneal@472:     This function performs a rate limit check.
bgneal@472:     One is added to a counter associated with the IP address. If the
bgneal@472:     counter exceeds set_point per interval, True is returned, and False
bgneal@472:     otherwise. If the set_point is exceeded for the first time, the counter
bgneal@472:     associated with the IP is set to expire according to the lockout parameter.
bgneal@472:     This locks the IP address as this function will then return True for the
bgneal@472:     period specified by lockout.
bgneal@472: 
bgneal@472:     """
bgneal@472:     if not ip:
bgneal@472:         logger.error("rate_limit.rate_check could not get IP")
bgneal@472:         return False
bgneal@472:     key = _make_key(ip)
bgneal@472: 
bgneal@472:     conn = _get_connection()
bgneal@472:     if not conn:
bgneal@472:         return False
bgneal@472: 
bgneal@472:     val = conn.incr(key)
bgneal@472: 
bgneal@472:     # Set expire time, if necessary.
bgneal@472:     # If this is the first time, set it according to interval.
bgneal@472:     # If the set_point has just been exceeded, set it according to lockout.
bgneal@472:     if val == 1:
bgneal@472:         conn.expire(key, _to_seconds(interval))
bgneal@472:     elif val == set_point:
bgneal@472:         conn.expire(key, _to_seconds(lockout))
bgneal@472: 
bgneal@472:     tripped = val >= set_point
bgneal@472: 
bgneal@472:     if tripped:
bgneal@472:         logger.info("Rate limiter tripped for %s; counter = %d", ip, val)
bgneal@472: 
bgneal@472:     return tripped
bgneal@472: 
bgneal@472: 
bgneal@472: def block_ip(ip, count=1000000, interval=datetime.timedelta(weeks=2)):
bgneal@472:     """
bgneal@472:     This function jams the rate limit record for the given IP so that the IP is
bgneal@472:     blocked for the given interval. If the record doesn't exist, it is created.
bgneal@472:     This is useful for manually blocking an IP after detecting suspicious
bgneal@472:     behavior.
bgneal@472: 
bgneal@472:     """
bgneal@472:     if not ip:
bgneal@472:         logger.error("rate_limit.block_ip could not get IP")
bgneal@472:         return
bgneal@472: 
bgneal@472:     key = _make_key(ip)
bgneal@472:     conn = _get_connection()
bgneal@472:     if not conn:
bgneal@472:         return
bgneal@472: 
bgneal@472:     conn.setex(key, count, _to_seconds(interval))
bgneal@472:     logger.info("Rate limiter blocked IP %s; %d / %s", ip, count, interval)