ThrottlingMiddleware.on_pre_process_update()   A
last analyzed

Complexity

Conditions 4

Size

Total Lines 14
Code Lines 8

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 8
dl 0
loc 14
rs 10
c 0
b 0
f 0
cc 4
nop 3
1
# -*- coding: utf-8 -*-
2
3
4
from aiogram.dispatcher.handler import CancelHandler, current_handler
5
from utils.db_api.set_ban_members import ban_members
6
from utils.db_api.unban_members import unban_member
7
from data.config import ban_time, exceeded_count
8
from aiogram.dispatcher.middlewares import BaseMiddleware
9
from aiogram.dispatcher import DEFAULT_RATE_LIMIT
10
from aiogram.utils.exceptions import Throttled
11
from aiogram import types, Dispatcher
12
import asyncio
13
14
15
"""
16
17
18
    Created on 26.09.2021
19
    
20
    @author: Nikita
21
22
23
"""
24
25
26
class ThrottlingMiddleware(BaseMiddleware):
27
28
    def __init__(self, limit=DEFAULT_RATE_LIMIT, key_prefix='antiflood_'):
29
30
        self.rate_limit = limit
31
        self.prefix = key_prefix
32
33
        super(ThrottlingMiddleware, self).__init__()
34
35
    async def on_pre_process_update(self, update: types.Update, data: dict):
36
37
        if update.message:
38
            user_id = update.message.from_user.id
39
40
        elif update.callback_query:
41
            user_id = update.callback_query.from_user.id
42
43
        else:
44
            return
45
46
        if user_id in await unban_member(user_id=user_id, ban_time=ban_time):
47
48
            raise CancelHandler()
49
50
    async def on_process_message(self, message: types.Message, data: dict):
51
52
        handler = current_handler.get()
53
        dispatcher = Dispatcher.get_current()
54
55
        if handler:
56
57
            limit = getattr(handler, "throttling_rate_limit", self.rate_limit)
58
            key = getattr(handler, "throttling_key", f"{self.prefix}_{handler.__name__}")
59
60
        else:
61
62
            limit = self.rate_limit
63
            key = f"{self.prefix}_message"
64
65
        try:
66
            await dispatcher.throttle(key, rate=limit)
67
68
        except Throttled as t:
69
70
            await self.message_throttled(message, t)
71
            raise CancelHandler()
72
73
    async def message_throttled(self, message: types.Message, throttled: Throttled):
74
75
        """
76
77
78
        Notify user only on first exceed and notify about unlocking only on last exceed
79
80
81
        """
82
83
        # Prevent flooding
84
        if throttled.exceeded_count == exceeded_count:
85
86
            await message.reply('Слишком много запросов 😤!')
87
88
            await ban_members(user_id=throttled.user,
89
                              time_out=throttled.called_at)
90
91
            await message.reply(f'Вы были заблокированы, на {ban_time} секунд!')
92
93
            await asyncio.sleep(ban_time)
94
95
            await message.reply('Вас разблокировали, теперь Вы снова можете пообщаться со мной 🧐!')
96