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