manyi/apps/web/wrapper.py

# -*- coding: utf-8 -*-
# !/usr/bin/env python

import datetime
import logging
from functools import wraps

from django.conf import settings
from typing import Any, TYPE_CHECKING

from apps import serviceCache
from apps.web.core.utils import JsonErrorResponse
from apps.web.utils import get_client_ip

if TYPE_CHECKING:
    from apps.web.core.db import RoleBaseDocument
    from django.core.handlers.wsgi import WSGIRequest
    from django.http import HttpResponse


def request_limit_by_user(operation='common', limit=50, period=24 * 60 * 60, logger=None):
    """
    http请求速率限制(按照用户ID)
    :param logger:
    :param operation: 操作
    :param limit: 限制次数
    :param period: 时长
    :return:
    """

    logger = logging.getLogger(__name__) if logger is None else logger

    def decorate(func):
        @wraps(func)
        def wrapper(request, *args, **kwargs):
            # type:(WSGIRequest, *Any, **Any)->HttpResponse

            if not settings.ENABLE_REQUEST_LIMIT:
                logger.debug('request limit is disable.')
                return func(request, *args, **kwargs)

            user = request.user  # type: RoleBaseDocument

            dayTime = datetime.datetime.now().strftime('%Y-%m-%d')
            key = '{}_{}_{}'.format(user.request_limit_key, operation, dayTime)

            value = serviceCache.get(key, 0)
            if value >= limit:
                logger.debug("{} in {} exceed max limit.".format(key, get_client_ip(request)))
                return JsonErrorResponse(description=u'操作过于频繁，休息一下喔。')
            else:
                value += 1
                serviceCache.set(key, value, period)
                return func(request, *args, **kwargs)

        return wrapper

    return decorate


def request_limit_by_ip(operation='common', limit=50, period=24 * 60 * 60, logger=None):
    """
    http请求速率限制(按照IP)
    :param logger:
    :param operation: 操作
    :param limit: 限制次数
    :param period: 时长
    :return:
    """

    logger = logging.getLogger(__name__) if logger is None else logger

    def decorate(func):
        @wraps(func)
        def wrapper(request, *args, **kwargs):
            # type:(WSGIRequest, *Any, **Any)->HttpResponse

            if not settings.ENABLE_REQUEST_LIMIT:
                logger.debug('request limit is disable.')
                return func(request, *args, **kwargs)

            ip = get_client_ip(request)

            dayTime = datetime.datetime.now().strftime('%Y-%m-%d')

            key = '{}_{}_{}'.format(ip, operation, dayTime)

            value = serviceCache.get(key, 0)
            if value >= limit:
                logger.debug("{} in {} exceed max limit.".format(operation, ip))
                return JsonErrorResponse(description=u'操作过于频繁，休息一下喔。')
            else:
                value += 1
                serviceCache.set(key, value, period)
                return func(request, *args, **kwargs)

        return wrapper

    return decorate