pycaptcha/services/limiter.py

# -*- coding: utf-8 -*-
"""
Created on 2022-06-15
---------
@summary:  限制器
---------
@author: Dzr
"""
from fastapi import FastAPI
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address

import setting
from services.robot import send_msg
from services.utils import current_date

LIMITER = Limiter(key_func=get_remote_address)  # 创建限制器实例对象

LM_COUNTER = 0  # 计数器
LM_DATE = None  # 当前日期
LM_FORECAST_WARNING = True  # 上限预告标识
LM_EXCEEDED_WARNING = True  # 超限告警标识
MAX_LIMIT = setting.LIMITER_MAX_LIMIT  # 最大访问次数
"""
    参考：https://limits.readthedocs.io/en/latest/quickstart.html
    Examples:
        1、10 per hour
        2、10/hour
        3、10/hour;100/day;2000 per year
        4、100/day, 500/7days
        5、5/minute
"""
LM_VALUE = f"{MAX_LIMIT}/day"  # 限制数量


def register_limiter(app: FastAPI):
    app.state.limiter = LIMITER
    app.add_exception_handler(RateLimitExceeded, limit_exceeded_handler)


def limit_exceeded_handler(*args, **kwargs):
    global LM_EXCEEDED_WARNING
    if LM_EXCEEDED_WARNING:
        msg = "今日接口调用次数已达上限！\n 继续使用请点击"
        send_msg("超级鹰", MAX_LIMIT, MAX_LIMIT, msg, allow_reset=True)
        LM_EXCEEDED_WARNING = False
    return _rate_limit_exceeded_handler(*args, **kwargs)


def flush_limiter():
    global LM_COUNTER, LM_FORECAST_WARNING, LM_EXCEEDED_WARNING
    LM_COUNTER = 0
    LM_FORECAST_WARNING = True
    LM_EXCEEDED_WARNING = True


def limiter_warring():
    global LM_DATE, LM_COUNTER, LM_FORECAST_WARNING
    curr_date = current_date()
    if curr_date != LM_DATE:
        if LM_DATE is not None:
            flush_limiter()
            LM_COUNTER += 1
        LM_DATE = curr_date
    val = '{:.2f}'.format(LM_COUNTER / MAX_LIMIT * 100)
    if float(val) > 80 and LM_FORECAST_WARNING:
        send_msg("超级鹰", MAX_LIMIT, LM_COUNTER, f"使用次数已超过{val}%")
        LM_FORECAST_WARNING = False