simonw / ratelimitcache

A memcached backed rate limiting decorator for Django.

This URL has Read+Write access

ratelimitcache / ratelimitcache.py
raw normal view history
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 1 from django.http import HttpResponseForbidden
2 from django.core.cache import cache
3 from datetime import datetime, timedelta
b8950893 » simonw 2009-01-07 Use sha hash of the POST fi... 4 import functools, sha
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 5
6 class ratelimit(object):
7 "Instances of this class can be used as decorators"
8 # This class is designed to be sub-classed
165bfbae » simonw 2009-01-07 Made the defaults more inte... 9 minutes = 2 # The time period
10 requests = 20 # Number of allowed requests in that time period
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 11
12 prefix = 'rl-' # Prefix for memcache key
13
14 def __init__(self, **options):
bb3d8d75 » simonw 2009-01-07 Fixed bug in constructor sh... 15 for key, value in options.items():
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 16 setattr(self, key, value)
17
18 def __call__(self, fn):
19 def wrapper(request, *args, **kwargs):
20 return self.view_wrapper(request, fn, *args, **kwargs)
21 functools.update_wrapper(wrapper, fn)
22 return wrapper
23
24 def view_wrapper(self, request, fn, *args, **kwargs):
25 if not self.should_ratelimit(request):
26 return fn(request, *args, **kwargs)
27
28 counts = self.get_counters(request).values()
29
30 # Increment rate limiting counter
31 self.cache_incr(self.current_key(request))
32
33 # Have they failed?
34 if sum(counts) >= self.requests:
35 return self.disallowed(request)
36
37 return fn(request, *args, **kwargs)
38
39 def cache_get_many(self, keys):
40 return cache.get_many(keys)
41
42 def cache_incr(self, key):
43 # memcache is only backend that can increment atomically
44 try:
55292676 » simonw 2009-01-07 Call .add() before .incr(),... 45 # add first, to ensure the key exists
591db0eb » simonw 2009-09-24 expire_after is now a metho... 46 cache._cache.add(key, '0', time=self.expire_after())
55292676 » simonw 2009-01-07 Call .add() before .incr(),... 47 cache._cache.incr(key)
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 48 except AttributeError:
591db0eb » simonw 2009-09-24 expire_after is now a metho... 49 cache.set(key, cache.get(key, 0) + 1, self.expire_after())
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 50
51 def should_ratelimit(self, request):
52 return True
53
54 def get_counters(self, request):
55 return self.cache_get_many(self.keys_to_check(request))
56
57 def keys_to_check(self, request):
58 extra = self.key_extra(request)
59 now = datetime.now()
60 return [
61 '%s%s-%s' % (
62 self.prefix,
63 extra,
64 (now - timedelta(minutes = minute)).strftime('%Y%m%d%H%M')
65 ) for minute in range(self.minutes + 1)
66 ]
67
68 def current_key(self, request):
69 return '%s%s-%s' % (
70 self.prefix,
71 self.key_extra(request),
72 datetime.now().strftime('%Y%m%d%H%M')
73 )
74
75 def key_extra(self, request):
76 # By default, their IP address is used
77 return request.META.get('REMOTE_ADDR', '')
78
79 def disallowed(self, request):
80 "Over-ride this method if you want to log incidents"
81 return HttpResponseForbidden('Rate limit exceeded')
591db0eb » simonw 2009-09-24 expire_after is now a metho... 82
83 def expire_after(self):
84 "Used for setting the memcached cache expiry"
85 return (self.minutes + 1) * 60
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 86
08ef7b85 » simonw 2009-01-07 Fixed a couple of bugs in r... 87 class ratelimit_post(ratelimit):
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 88 "Rate limit POSTs - can be used to protect a login form"
89 key_field = None # If provided, this POST var will affect the rate limit
90
91 def should_ratelimit(self, request):
92 return request.method == 'POST'
93
94 def key_extra(self, request):
95 # IP address and key_field (if it is set)
08ef7b85 » simonw 2009-01-07 Fixed a couple of bugs in r... 96 extra = super(ratelimit_post, self).key_extra(request)
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 97 if self.key_field:
b8950893 » simonw 2009-01-07 Use sha hash of the POST fi... 98 value = sha.new(request.POST.get(self.key_field, '')).hexdigest()
99 extra += '-' + value
6a95b45d » simonw 2009-01-07 Implementation of rate limiter 100 return extra
101