mirror of
				https://github.com/civsocit/olgram.git
				synced 2023-07-22 01:29:12 +03:00 
			
		
		
		
	Возможность ограничивать права на бота
This commit is contained in:
		
							parent
							
								
									5fcb5b8900
								
							
						
					
					
						commit
						1aeec0c9d8
					
				
							
								
								
									
										6
									
								
								main.py
									
									
									
									
									
								
							
							
						
						
									
										6
									
								
								main.py
									
									
									
									
									
								
							@ -3,7 +3,8 @@ import argparse
 | 
			
		||||
from tortoise import Tortoise
 | 
			
		||||
 | 
			
		||||
from olgram.router import dp
 | 
			
		||||
from olgram.settings import TORTOISE_ORM
 | 
			
		||||
from olgram.settings import TORTOISE_ORM, OlgramSettings
 | 
			
		||||
from olgram.utils.permissions import AccessMiddleware
 | 
			
		||||
from server.custom import init_redis
 | 
			
		||||
 | 
			
		||||
import olgram.commands.bots  # noqa: F401
 | 
			
		||||
@ -22,7 +23,8 @@ async def init_database():
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
async def init_olgram():
 | 
			
		||||
    from olgram.router import bot
 | 
			
		||||
    from olgram.router import bot, dp
 | 
			
		||||
    dp.setup_middleware(AccessMiddleware(OlgramSettings.admin_id()))
 | 
			
		||||
    from aiogram.types import BotCommand
 | 
			
		||||
    await bot.set_my_commands(
 | 
			
		||||
        [
 | 
			
		||||
 | 
			
		||||
@ -6,11 +6,13 @@ from aiogram import types
 | 
			
		||||
from aiogram.dispatcher import FSMContext
 | 
			
		||||
from textwrap import dedent
 | 
			
		||||
from olgram.settings import OlgramSettings
 | 
			
		||||
from olgram.utils.permissions import public
 | 
			
		||||
 | 
			
		||||
from olgram.router import dp
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@dp.message_handler(commands=["start"], state="*")
 | 
			
		||||
@public()
 | 
			
		||||
async def start(message: types.Message, state: FSMContext):
 | 
			
		||||
    """
 | 
			
		||||
    Команда /start
 | 
			
		||||
@ -33,6 +35,7 @@ async def start(message: types.Message, state: FSMContext):
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@dp.message_handler(commands=["help"], state="*")
 | 
			
		||||
@public()
 | 
			
		||||
async def help(message: types.Message, state: FSMContext):
 | 
			
		||||
    """
 | 
			
		||||
    Команда /help
 | 
			
		||||
@ -42,3 +45,12 @@ async def help(message: types.Message, state: FSMContext):
 | 
			
		||||
    Техническая поддержка: @civsocit_feedback_bot
 | 
			
		||||
    Версия {OlgramSettings.version()}
 | 
			
		||||
    """))
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@dp.message_handler(commands=["chatid"], state="*")
 | 
			
		||||
@public()
 | 
			
		||||
async def chat_id(message: types.Message, state: FSMContext):
 | 
			
		||||
    """
 | 
			
		||||
    Команда /chatid
 | 
			
		||||
    """
 | 
			
		||||
    await message.answer(message.chat.id)
 | 
			
		||||
 | 
			
		||||
@ -32,6 +32,11 @@ class OlgramSettings(AbstractSettings):
 | 
			
		||||
    def version(cls):
 | 
			
		||||
        return "0.1.0"
 | 
			
		||||
 | 
			
		||||
    @classmethod
 | 
			
		||||
    @lru_cache
 | 
			
		||||
    def admin_id(cls):
 | 
			
		||||
        return cls._get_env("ADMIN_ID", True)
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class ServerSettings(AbstractSettings):
 | 
			
		||||
    @classmethod
 | 
			
		||||
 | 
			
		||||
							
								
								
									
										52
									
								
								olgram/utils/permissions.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										52
									
								
								olgram/utils/permissions.py
									
									
									
									
									
										Normal file
									
								
							@ -0,0 +1,52 @@
 | 
			
		||||
import aiogram.types as types
 | 
			
		||||
from aiogram.dispatcher.handler import CancelHandler, current_handler
 | 
			
		||||
from aiogram.dispatcher.middlewares import BaseMiddleware
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def public():
 | 
			
		||||
    """
 | 
			
		||||
    Хендлеры с этим декоратором будут обрабатываться даже если пользователь не является владельцем бота
 | 
			
		||||
    (например, команда /help)
 | 
			
		||||
    :return:
 | 
			
		||||
    """
 | 
			
		||||
 | 
			
		||||
    def decorator(func):
 | 
			
		||||
        setattr(func, "access_public", True)
 | 
			
		||||
        return func
 | 
			
		||||
 | 
			
		||||
    return decorator
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class AccessMiddleware(BaseMiddleware):
 | 
			
		||||
    def __init__(self, access_chat_id: int):
 | 
			
		||||
        self._access_chat_id = access_chat_id
 | 
			
		||||
        super(AccessMiddleware, self).__init__()
 | 
			
		||||
 | 
			
		||||
    @classmethod
 | 
			
		||||
    def _is_public_command(cls) -> bool:
 | 
			
		||||
        handler = current_handler.get()
 | 
			
		||||
        return handler and getattr(handler, "access_public", False)
 | 
			
		||||
 | 
			
		||||
    async def on_process_message(self, message: types.Message, data: dict):
 | 
			
		||||
        admin_id = self._access_chat_id
 | 
			
		||||
        if not admin_id:
 | 
			
		||||
            return  # Администратор бота вообще не указан
 | 
			
		||||
 | 
			
		||||
        if self._is_public_command():  # Эта команда разрешена всем пользователям
 | 
			
		||||
            return
 | 
			
		||||
 | 
			
		||||
        if message.chat.id != admin_id:
 | 
			
		||||
            await message.answer("Владелец бота ограничил доступ к этому функционалу 😞")
 | 
			
		||||
            raise CancelHandler()
 | 
			
		||||
 | 
			
		||||
    async def on_process_callback_query(self, call: types.CallbackQuery, data: dict):
 | 
			
		||||
        admin_id = self._access_chat_id
 | 
			
		||||
        if not admin_id:
 | 
			
		||||
            return  # Администратор бота вообще не указан
 | 
			
		||||
 | 
			
		||||
        if self._is_public_command():  # Эта команда разрешена всем пользователям
 | 
			
		||||
            return
 | 
			
		||||
 | 
			
		||||
        if call.message.chat.id != admin_id:
 | 
			
		||||
            await call.answer("Владелец бота ограничил доступ к этому функционалу😞")
 | 
			
		||||
            raise CancelHandler()
 | 
			
		||||
		Loading…
	
		Reference in New Issue
	
	Block a user