From 363391b5756c864308dde4a75ad619d69eb11a28 Mon Sep 17 00:00:00 2001
From: mihalin <stopa3.10.1996@mail.ru>
Date: Tue, 18 Jan 2022 23:28:03 +0300
Subject: [PATCH] ban\unban commands

---
 main.py                 |  2 +-
 olgram/models/models.py | 11 +++++++++++
 poetry.lock             |  8 ++++----
 server/custom.py        | 28 +++++++++++++++++++++++++---
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/main.py b/main.py
index d877754..ce91bd2 100644
--- a/main.py
+++ b/main.py
@@ -52,7 +52,7 @@ def main():
                                            "разработки)", action="store_true")
     args = parser.parse_args()
 
-    loop = asyncio.get_event_loop()
+    loop = asyncio.new_event_loop()
     loop.run_until_complete(initialization())
 
     loop.create_task(dp.start_polling())
diff --git a/olgram/models/models.py b/olgram/models/models.py
index aab4cbd..42c5ed8 100644
--- a/olgram/models/models.py
+++ b/olgram/models/models.py
@@ -72,3 +72,14 @@ class GroupChat(Model):
 
     class Meta:
         table = 'group_chat'
+
+
+class BannedUser(Model):
+    id = fields.BigIntField(pk=True)
+    telegram_id = fields.BigIntField(index=True)
+    username = fields.CharField(max_length=100, default=None, null=True)
+
+    bot = fields.ForeignKeyField("models.Bot", related_name="banned_users", on_delete=fields.relational.CASCADE)
+
+    class Meta:
+        table = "bot_banned_user"
diff --git a/poetry.lock b/poetry.lock
index ac347c6..8081d35 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,6 +1,6 @@
 [[package]]
 name = "aerich"
-version = "0.5.4"
+version = "0.5.7"
 description = "A database migrations tool for Tortoise ORM."
 category = "main"
 optional = false
@@ -393,12 +393,12 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.8"
-content-hash = "c4533872a5037ef7b3ef7f9ada913fd6c0bf761ad642803103d464c90e5bc578"
+content-hash = "20aec5e4517c3e0bc03d4410544c1c898f0469c9bcfc4f4fecf4c405b1765ded"
 
 [metadata.files]
 aerich = [
-    {file = "aerich-0.5.4-py3-none-any.whl", hash = "sha256:ce338f42af750632a240648c1d6a052b3cb1b74b8b0d914c3ed30ee3b3b17349"},
-    {file = "aerich-0.5.4.tar.gz", hash = "sha256:067cf4a07b6714a8d6964302cb483fd297f94f8055cbf5489d86abeb03f02540"},
+    {file = "aerich-0.5.7-py3-none-any.whl", hash = "sha256:0684eb3d631c7c6e14caf2b2c3b9dad1c15ce8ade5771773c015a302f54ff4f6"},
+    {file = "aerich-0.5.7.tar.gz", hash = "sha256:f9ef8796f7a13ba9965eda0aa6840033bbd42b2e4e52c24d8f0dbdb85e4a5187"},
 ]
 aiocache = [
     {file = "aiocache-0.11.1-py2.py3-none-any.whl", hash = "sha256:e55c7caaa5753794fd301c3a2e592737fa1d036db9f8d04ae154facdfb48a157"},
diff --git a/server/custom.py b/server/custom.py
index b93e9d2..c93734b 100644
--- a/server/custom.py
+++ b/server/custom.py
@@ -10,7 +10,7 @@ from aioredis import Redis
 import logging
 import typing as ty
 from olgram.settings import ServerSettings
-from olgram.models.models import Bot, GroupChat
+from olgram.models.models import Bot, GroupChat, BannedUser
 
 
 logging.basicConfig(level=logging.INFO)
@@ -34,7 +34,7 @@ async def message_handler(message: types.Message, *args, **kwargs):
     _logger.info("message handler")
     bot = db_bot_instance.get()
 
-    if message.text and message.text.startswith("/start"):
+    if message.text and message.text == "/start":
         # На команду start нужно ответить, не пересылая сообщение никуда
         return SendMessage(chat_id=message.chat.id,
                            text=bot.start_text + ServerSettings.append_text())
@@ -42,7 +42,15 @@ async def message_handler(message: types.Message, *args, **kwargs):
     super_chat_id = await bot.super_chat_id()
 
     if message.chat.id != super_chat_id:
-        # Это обычный чат: сообщение нужно переслать в супер-чат
+        # Это обычный чат
+
+        # Проверить, не забанен ли пользователь
+        banned = await bot.banned_users.filter(telegram_id=message.chat.id)
+        if banned:
+            return SendMessage(chat_id=message.chat.id,
+                               text="Вы заблокированы в этом боте")
+
+        # сообщение нужно переслать в супер-чат
         new_message = await message.forward(super_chat_id)
         await _redis.set(_message_unique_id(bot.pk, new_message.message_id), message.chat.id)
 
@@ -62,6 +70,20 @@ async def message_handler(message: types.Message, *args, **kwargs):
                                        text="<i>Невозможно переслать сообщение: автор не найден</i>",
                                        parse_mode="HTML")
             chat_id = int(chat_id)
+
+            if message.text == "/ban":
+                user, _ = await BannedUser.get_or_create(telegram_id=message.from_user.id, bot=bot)
+                await user.save()
+                return SendMessage(chat_id=message.chat.id, text="Пользователь заблокирован")
+
+            if message.text == "/unban":
+                banned_user = await bot.banned_users.filter(telegram_id=chat_id).first()
+                if not banned_user:
+                    return SendMessage(chat_id=message.chat.id, text="Пользователь не был забанен")
+                else:
+                    await banned_user.delete()
+                    return SendMessage(chat_id=message.chat.id, text="Пользователь разбанен")
+
             try:
                 await message.copy_to(chat_id)
             except (exceptions.MessageError, exceptions.BotBlocked):