Whitelist Middleware

whitelist.py

from typing import Any, Dict, List, Set

from aiogram import types
from aiogram.dispatcher.handler import CancelHandler
from aiogram.dispatcher.middlewares import BaseMiddleware
from loguru import logger


class WhitelistMiddleware(BaseMiddleware):
    def __init__(self, users: Set[int]):
        super().__init__()
        self.users = users

    def check_user_access(self, user: types.User) -> bool:
        check = user.id in self.users
        if not check:
            logger.bind(user=user.id).warning("Access denied for user {user_id}.", user_id=user.id)
            raise CancelHandler()
        return True

    async def on_pre_process_message(self, message: types.Message, data: Dict[str, Any]):
        return self.check_user_access(message.from_user)

    async def on_pre_process_edited_message(
        self, edited_message: types.Message, data: Dict[str, Any]
    ):
        return self.check_user_access(edited_message.from_user)

    async def on_pre_process_channel_post(self, channel_post: types.Message, data: Dict[str, Any]):
        return self.check_user_access(channel_post.from_user)

    async def on_pre_process_edited_channel_post(
        self, edited_channel_post: types.Message, data: Dict[str, Any]
    ):
        return self.check_user_access(edited_channel_post.from_user)

    async def on_pre_process_inline_query(
        self, inline_query: types.InlineQuery, data: Dict[str, Any]
    ):
        return self.check_user_access(inline_query.from_user)

    async def on_pre_process_chosen_inline_result(
        self, chosen_inline_result: types.ChosenInlineResult, data: Dict[str, Any]
    ):
        return self.check_user_access(chosen_inline_result.from_user)

    async def on_pre_process_callback_query(
        self, callback_query: types.CallbackQuery, data: Dict[str, Any]
    ):
        return self.check_user_access(callback_query.from_user)

    async def on_pre_process_shipping_query(
        self, shipping_query: types.ShippingQuery, data: Dict[str, Any]
    ):
        return self.check_user_access(shipping_query.from_user)

    async def on_pre_process_pre_checkout_query(
        self, pre_checkout_query: types.PreCheckoutQuery, data: Dict[str, Any]
    ):
        return self.check_user_access(pre_checkout_query.from_user)