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