Возможность ограничивать права на бота

This commit is contained in:
mihalin 2021-10-01 19:48:23 +03:00
parent 5fcb5b8900
commit 1aeec0c9d8
5 changed files with 74 additions and 3 deletions

View File

@ -3,7 +3,8 @@ import argparse
from tortoise import Tortoise from tortoise import Tortoise
from olgram.router import dp 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 from server.custom import init_redis
import olgram.commands.bots # noqa: F401 import olgram.commands.bots # noqa: F401
@ -22,7 +23,8 @@ async def init_database():
async def init_olgram(): 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 from aiogram.types import BotCommand
await bot.set_my_commands( await bot.set_my_commands(
[ [

View File

@ -6,11 +6,13 @@ from aiogram import types
from aiogram.dispatcher import FSMContext from aiogram.dispatcher import FSMContext
from textwrap import dedent from textwrap import dedent
from olgram.settings import OlgramSettings from olgram.settings import OlgramSettings
from olgram.utils.permissions import public
from olgram.router import dp from olgram.router import dp
@dp.message_handler(commands=["start"], state="*") @dp.message_handler(commands=["start"], state="*")
@public()
async def start(message: types.Message, state: FSMContext): async def start(message: types.Message, state: FSMContext):
""" """
Команда /start Команда /start
@ -33,6 +35,7 @@ async def start(message: types.Message, state: FSMContext):
@dp.message_handler(commands=["help"], state="*") @dp.message_handler(commands=["help"], state="*")
@public()
async def help(message: types.Message, state: FSMContext): async def help(message: types.Message, state: FSMContext):
""" """
Команда /help Команда /help
@ -42,3 +45,12 @@ async def help(message: types.Message, state: FSMContext):
Техническая поддержка: @civsocit_feedback_bot Техническая поддержка: @civsocit_feedback_bot
Версия {OlgramSettings.version()} Версия {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)

View File

@ -32,6 +32,11 @@ class OlgramSettings(AbstractSettings):
def version(cls): def version(cls):
return "0.1.0" return "0.1.0"
@classmethod
@lru_cache
def admin_id(cls):
return cls._get_env("ADMIN_ID", True)
class ServerSettings(AbstractSettings): class ServerSettings(AbstractSettings):
@classmethod @classmethod

View 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()