Django 自定义中间件做接口访问频率限制

依赖

django
redis

思路

将用户或者ip对每个接口的访问频次记录在redis，django自定义中间件拦截请求，如果超出访问限制则屏蔽掉请求。

实现

简单介绍django 自定义中间件
Django中间件必须是一个类，不需要继承任何类，并提供四个接口：
1、 process_request(self, request)该方法在请求到来的时候调用。
2、 process_view(self ,request, fnc , arg ,kwarg)在本次将要执行的View函数被调用前调用本函数。
3、 process_response(self,request,response)在执行完View函数准备将响应发到客户端前被执行。
4、 process_exception(self,request, exception)View函数在抛出异常时该函数被调用，得到的exception参数是实际上抛出的异常实例。通过此方法可以进行很好的错误控制，提供友好的用户界面。

IGNORE_PREVENT = []

class PreventMiddleware:
    """PreventMiddleware

    控制接口调用频率, 防止恶意调用和抓取数据

    """

    def process_request(self, request):
        pass

    def process_view(self, request, view_func, view_args, view_kwargs):
        """打印request进入时的log"""
        user_id = request.user_id # 自定义中间件获得用户id
        method = request.method
        ip = get_client_ip(request)
        path = request.path

        # 过滤
        for ignore_url in IGNORE_PREVENT:
            if ignore_url in path:
                return

        # 运营小号放行
        # if user_id:
        #     user = userservice.get_user_by_id(user_id)
        #     if user and user.get('staff_flag'):
        #         return

        write_success = True
        # 对‘写’接口, 做更严格的频率控制. 首先
        if method in ['POST', 'PUT', 'DELETE']:
            write_success = check_prevent(user_id, ip, path)

        if not write_success:
            # todo

def check_prevent(user_id, ip, path):
    if user_id:
        success = prevent_client.check('user_write_limit', user_id, path)
    else:
        success = prevent_client.check('ip_write_limit', ip, path)
    return success


def get_client_ip(request):
    x_forwarded_for = request.META.get('HTTP_X_FORWARDED_FOR')
    if x_forwarded_for:
        ip = x_forwarded_for.split(',')[0]
    else:
        ip = request.META.get('REMOTE_ADDR')
    return ip

在看看prevent_client.check()这个函数的实现

class FrequencyLimitService:
    def __init__(self, config, redis_host, redis_port):
        self.config = config
        self.redis = MyRedis({'host': redis_host, 'port': redis_port})

    def check(self, code, biz_key, ext_info=''):
        """
        触发一次指定code的规则调用统计, 并返回当前是否已达到策略上限
        :param code: 业务场景编码, 请在配置文件里配置
        :param biz_key: 业务唯一码, 各接入业务模块负责组装, 并保证一个场景code下该值唯一
        :param ext_info: 扩展信息, 用于打印
        :return: 如果触发上限值, 则返回False
        """
        try:
            policy = self.config[code]
            max_timers, period = policy['formula']
        except KeyError:
            logger.error('policy code: %s not config' % code)
            return False

        store_key = 'fre_%s_%s' % (code, biz_key)
        timers = self.redis.get(store_key)
        timers = int(timers) if timers else 0
        if timers:
            if timers <= max_timers:
                timers = self.redis.incr(store_key)
                if timers == 1:  # 并发控制, 如果刚好过期, 则要重新设置过期时间
                    self.redis.expire(store_key, period)
        else:
            self.redis.set(store_key, 1, ex=period)

        success = timers <= max_timers
        if success:
            logger.debug('fre check success, code: %s, key: %s, context: %s'
                        % (code, biz_key, ext_info))
        else:
            logger.info('fre check fail, code: %s, key: %s, context: %s'
                        % (code, biz_key, ext_info))
        return success

最后记得把自定义的中间件加到setting文件MIDDLEWARE_CLASSES配置里