realms-wiki/realms/ratelimit.py

46 lines
1.5 KiB
Python
Raw Normal View History

2013-09-26 16:51:15 +03:00
import time
from functools import update_wrapper
from flask import request, g
2013-10-05 00:42:45 +03:00
from services import cache
2013-09-26 16:51:15 +03:00
class RateLimit(object):
expiration_window = 10
def __init__(self, key_prefix, limit, per, send_x_headers):
self.reset = (int(time.time()) // per) * per + per
self.key = key_prefix + str(self.reset)
self.limit = limit
self.per = per
self.send_x_headers = send_x_headers
p = cache.pipeline()
p.incr(self.key)
p.expireat(self.key, self.reset + self.expiration_window)
self.current = min(p.execute()[0], limit)
remaining = property(lambda x: x.limit - x.current)
over_limit = property(lambda x: x.current >= x.limit)
def get_view_rate_limit():
return getattr(g, '_view_rate_limit', None)
def on_over_limit(limit):
2013-10-05 00:42:45 +03:00
return 'Slow it down', 400
2013-09-26 16:51:15 +03:00
2013-10-05 00:42:45 +03:00
def ratelimiter(limit, per=300, send_x_headers=True,
2013-09-26 16:51:15 +03:00
over_limit=on_over_limit,
scope_func=lambda: request.remote_addr,
key_func=lambda: request.endpoint):
def decorator(f):
def rate_limited(*args, **kwargs):
key = 'rate-limit/%s/%s/' % (key_func(), scope_func())
rlimit = RateLimit(key, limit, per, send_x_headers)
g._view_rate_limit = rlimit
if over_limit is not None and rlimit.over_limit:
return over_limit(rlimit)
return f(*args, **kwargs)
return update_wrapper(rate_limited, f)
2013-10-05 00:42:45 +03:00
return decorator