Добавлены миддлвари на примере ThrottlingMiddleware
This commit is contained in:
parent
90791c3091
commit
c3da772339
2
bot.py
2
bot.py
|
@ -5,6 +5,7 @@ from loguru import logger
|
||||||
|
|
||||||
import filters
|
import filters
|
||||||
import handlers
|
import handlers
|
||||||
|
import middlewares
|
||||||
from data import config
|
from data import config
|
||||||
|
|
||||||
bot = Bot(token=config.BOT_TOKEN, parse_mode=types.ParseMode.HTML)
|
bot = Bot(token=config.BOT_TOKEN, parse_mode=types.ParseMode.HTML)
|
||||||
|
@ -15,6 +16,7 @@ dp = Dispatcher(bot, storage=storage)
|
||||||
# noinspection PyUnusedLocal
|
# noinspection PyUnusedLocal
|
||||||
async def on_startup(web_app: web.Application):
|
async def on_startup(web_app: web.Application):
|
||||||
filters.setup(dp)
|
filters.setup(dp)
|
||||||
|
middlewares.setup(dp)
|
||||||
handlers.user.setup(dp)
|
handlers.user.setup(dp)
|
||||||
await dp.bot.delete_webhook()
|
await dp.bot.delete_webhook()
|
||||||
await dp.bot.set_webhook(config.WEBHOOK_URL)
|
await dp.bot.set_webhook(config.WEBHOOK_URL)
|
||||||
|
|
|
@ -1,8 +1,10 @@
|
||||||
from aiogram import Dispatcher
|
from aiogram import Dispatcher
|
||||||
from aiogram.dispatcher.filters import CommandStart
|
from aiogram.dispatcher.filters import CommandStart, CommandHelp
|
||||||
|
|
||||||
|
from .help import bot_help
|
||||||
from .start import bot_start
|
from .start import bot_start
|
||||||
|
|
||||||
|
|
||||||
def setup(dp: Dispatcher):
|
def setup(dp: Dispatcher):
|
||||||
dp.register_message_handler(start, CommandStart())
|
dp.register_message_handler(bot_start, CommandStart())
|
||||||
|
dp.register_message_handler(bot_help, CommandHelp())
|
||||||
|
|
13
handlers/user/help.py
Normal file
13
handlers/user/help.py
Normal file
|
@ -0,0 +1,13 @@
|
||||||
|
from aiogram import types
|
||||||
|
|
||||||
|
from utils.misc import rate_limit
|
||||||
|
|
||||||
|
|
||||||
|
@rate_limit(5, 'help')
|
||||||
|
async def bot_help(msg: types.Message):
|
||||||
|
text = [
|
||||||
|
'Список команд: ',
|
||||||
|
'/start - Начать диалог',
|
||||||
|
'/help - Получить справку'
|
||||||
|
]
|
||||||
|
await msg.answer('\n'.join(text))
|
7
middlewares/__init__.py
Normal file
7
middlewares/__init__.py
Normal file
|
@ -0,0 +1,7 @@
|
||||||
|
from aiogram import Dispatcher
|
||||||
|
|
||||||
|
from .throttling import ThrottlingMiddleware
|
||||||
|
|
||||||
|
|
||||||
|
def setup(dp: Dispatcher):
|
||||||
|
dp.middleware.setup(ThrottlingMiddleware)
|
49
middlewares/throttling.py
Normal file
49
middlewares/throttling.py
Normal file
|
@ -0,0 +1,49 @@
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from aiogram import types, Dispatcher
|
||||||
|
from aiogram.dispatcher import DEFAULT_RATE_LIMIT
|
||||||
|
from aiogram.dispatcher.handler import CancelHandler, current_handler
|
||||||
|
from aiogram.dispatcher.middlewares import BaseMiddleware
|
||||||
|
from aiogram.utils.exceptions import Throttled
|
||||||
|
|
||||||
|
|
||||||
|
class ThrottlingMiddleware(BaseMiddleware):
|
||||||
|
"""
|
||||||
|
Simple middleware
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, limit=DEFAULT_RATE_LIMIT, key_prefix='antiflood_'):
|
||||||
|
self.rate_limit = limit
|
||||||
|
self.prefix = key_prefix
|
||||||
|
super(ThrottlingMiddleware, self).__init__()
|
||||||
|
|
||||||
|
# noinspection PyUnusedLocal
|
||||||
|
async def on_process_message(self, message: types.Message, data: dict):
|
||||||
|
handler = current_handler.get()
|
||||||
|
dispatcher = Dispatcher.get_current()
|
||||||
|
if handler:
|
||||||
|
limit = getattr(handler, 'throttling_rate_limit', self.rate_limit)
|
||||||
|
key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
|
||||||
|
else:
|
||||||
|
limit = self.rate_limit
|
||||||
|
key = f"{self.prefix}_message"
|
||||||
|
try:
|
||||||
|
await dispatcher.throttle(key, rate=limit)
|
||||||
|
except Throttled as t:
|
||||||
|
await self.message_throttled(message, t)
|
||||||
|
raise CancelHandler()
|
||||||
|
|
||||||
|
async def message_throttled(self, message: types.Message, throttled: Throttled):
|
||||||
|
handler = current_handler.get()
|
||||||
|
dispatcher = Dispatcher.get_current()
|
||||||
|
if handler:
|
||||||
|
key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
|
||||||
|
else:
|
||||||
|
key = f"{self.prefix}_message"
|
||||||
|
delta = throttled.rate - throttled.delta
|
||||||
|
if throttled.exceeded_count <= 2:
|
||||||
|
await message.reply('Too many requests! ')
|
||||||
|
await asyncio.sleep(delta)
|
||||||
|
thr = await dispatcher.check_key(key)
|
||||||
|
if thr.exceeded_count == throttled.exceeded_count:
|
||||||
|
await message.reply('Unlocked.')
|
|
@ -0,0 +1 @@
|
||||||
|
from . import user
|
|
@ -0,0 +1,3 @@
|
||||||
|
from . import db_api
|
||||||
|
from . import misc
|
||||||
|
from . import redis
|
1
utils/misc/__init__.py
Normal file
1
utils/misc/__init__.py
Normal file
|
@ -0,0 +1 @@
|
||||||
|
from .throttling import rate_limit
|
16
utils/misc/throttling.py
Normal file
16
utils/misc/throttling.py
Normal file
|
@ -0,0 +1,16 @@
|
||||||
|
def rate_limit(limit: int, key=None):
|
||||||
|
"""
|
||||||
|
Decorator for configuring rate limit and key in different functions.
|
||||||
|
|
||||||
|
:param limit:
|
||||||
|
:param key:
|
||||||
|
:return:
|
||||||
|
"""
|
||||||
|
|
||||||
|
def decorator(func):
|
||||||
|
setattr(func, 'throttling_rate_limit', limit)
|
||||||
|
if key:
|
||||||
|
setattr(func, 'throttling_key', key)
|
||||||
|
return func
|
||||||
|
|
||||||
|
return decorator
|
Loading…
Reference in New Issue
Block a user