From 1aeec0c9d87d607506009f79f7c5a2379d676667 Mon Sep 17 00:00:00 2001
From: mihalin <stopa3.10.1996@mail.ru>
Date: Fri, 1 Oct 2021 19:48:23 +0300
Subject: [PATCH] =?UTF-8?q?=D0=92=D0=BE=D0=B7=D0=BC=D0=BE=D0=B6=D0=BD?=
 =?UTF-8?q?=D0=BE=D1=81=D1=82=D1=8C=20=D0=BE=D0=B3=D1=80=D0=B0=D0=BD=D0=B8?=
 =?UTF-8?q?=D1=87=D0=B8=D0=B2=D0=B0=D1=82=D1=8C=20=D0=BF=D1=80=D0=B0=D0=B2?=
 =?UTF-8?q?=D0=B0=20=D0=BD=D0=B0=20=D0=B1=D0=BE=D1=82=D0=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore                  |  2 +-
 main.py                     |  6 +++--
 olgram/commands/start.py    | 12 +++++++++
 olgram/settings.py          |  5 ++++
 olgram/utils/permissions.py | 52 +++++++++++++++++++++++++++++++++++++
 5 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 olgram/utils/permissions.py

diff --git a/.gitignore b/.gitignore
index 49a2a2e..2cbd47f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,4 @@ __pycache__
 *.pyc
 config.json
 docker-compose-release.yaml
-docs/build
\ No newline at end of file
+docs/build
diff --git a/main.py b/main.py
index b2f26e4..e2ded1a 100644
--- a/main.py
+++ b/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(
         [
diff --git a/olgram/commands/start.py b/olgram/commands/start.py
index a8c971b..daf5263 100644
--- a/olgram/commands/start.py
+++ b/olgram/commands/start.py
@@ -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)
diff --git a/olgram/settings.py b/olgram/settings.py
index 629d93e..648a1c6 100644
--- a/olgram/settings.py
+++ b/olgram/settings.py
@@ -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
diff --git a/olgram/utils/permissions.py b/olgram/utils/permissions.py
new file mode 100644
index 0000000..4755ffd
--- /dev/null
+++ b/olgram/utils/permissions.py
@@ -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()