Rate Limiter для Flask

Rate Limiter

Декоратор для Python Flask

mport functools
from time import time
from flask import current_app, request, g, jsonify

_limiter = None

class MemRateLimit(object):
    """Rate limiter that uses a Python dictionary as storage."""
    def __init__(self):
        self.counters = {}

    def is_allowed(self, key, limit, period):
        """Check if the client's request should be allowed, based on the
        hit counter. Returns a 3-element tuple with a True/False result,
        the number of remaining hits in the period, and the time the
        counter resets for the next period."""
        now = int(time())
        begin_period = now // period * period
        end_period = begin_period + period

        self.cleanup(now)
        if key in self.counters:
            self.counters[key]['hits'] += 1
        else:
            self.counters[key] = {'hits': 1, 'reset': end_period}
        allow = True
        remaining = limit - self.counters[key]['hits']
        if remaining < 0:
            remaining = 0
            allow = False
        return allow, remaining, self.counters[key]['reset']

    def cleanup(self, now):
        """Eliminate expired keys."""
        for key, value in list(self.counters.items()):
            if value['reset'] < now:
                del self.counters[key]

def rate_limit(limit, period):
    """Limits the rate at which clients can send requests to 'limit' requests
    per 'period' seconds. Once a client goes over the limit all requests are
    answered with a status code 429 Too Many Requests for the remaining of
    that period."""
    def decorator(f):
        @functools.wraps(f)
        def wrapped(*args, **kwargs):
            if current_app.config['TESTING']:
                # no rate limiting for debug and testing configurations
                return f(*args, **kwargs)
            else:
                # initialize the rate limiter the first time here
                global _limiter
                if _limiter is None:
                    _limiter = MemRateLimit()

                # generate a unique key to represent the decorated function and
                # the IP address of the client. Rate limiting counters are
                # maintained on each unique key.
                key = '{0}/{1}'.format(f.__name__, request.remote_addr)
                allowed, remaining, reset = _limiter.is_allowed(key, limit,
                                                                period)

                # set the rate limit headers in g, so that they are picked up
                # by the after_request handler and attached to the response
                g.headers = {
                    'X-RateLimit-Remaining': str(remaining),
                    'X-RateLimit-Limit': str(limit),
                    'X-RateLimit-Reset': str(reset)
                }

                # if the client went over the limit respond with a 429 status
                # code, else invoke the wrapped function
                if not allowed:
                    response = jsonify(
                        {'status': 429, 'error': 'too many requests',
                         'message': 'You have exceeded your request rate'})
                    response.status_code = 429
                    return response

                # else we let the request through
                return f(*args, **kwargs)
        return wrapped
    return decorator