pycaptcha/services/limiter.py

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 glovar
from services.robot import send_msg
from services.utils import current_date

LIMITER = Limiter(key_func=get_remote_address)
MAX_LIMIT = glovar.lm_max_limit


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


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


def flush_limiter():
    glovar.lm_counter = 0
    glovar.lm_forecast_warning = True
    glovar.lm_exceeded_warning = True


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