wip: Правки по standard.welcome

src/ocab_modules/ocab_modules/standard/command_helper/main.py

@@ -1,11 +1,10 @@
 from aiogram.types import BotCommand
 
-from ocab_core.modules_system.public_api import (
+from ocab_core.modules_system.public_api import set_my_commands
-    set_my_commands, log
-)
 
 commands = dict()
 
+
 def register_command(command, description, role="USER"):
     if role not in commands:
         commands[role] = dict()
@@ -26,7 +25,7 @@ async def set_user_commands():
                 )
             )
 
-    log(bot_commands)
+    # log(bot_commands)
 
     await set_my_commands(
         bot_commands,

src/ocab_modules/ocab_modules/standard/welcome/main.py

@@ -2,46 +2,51 @@ import asyncio
 import random
 
 from aiogram import Bot, Router, types
-from aiogram.exceptions import TelegramBadRequest
+from aiogram.enums import ChatMemberStatus
-from aiogram.filters import IS_MEMBER, IS_NOT_MEMBER, ChatMemberUpdatedFilter
+from aiogram.filters import JOIN_TRANSITION, LEAVE_TRANSITION, ChatMemberUpdatedFilter
 from aiogram.types import ChatMemberUpdated
 
 from ocab_core.modules_system.public_api import log, register_router
 
-from .verifications_methods.base import VerificationCallback
+from .utils import get_plural_form
-from .verifications_methods.iamhuman import IAmHumanButton, IAmHumanInput
+from .verifications_methods.base import VerificationCallback, VerificationMethod
-from .verifications_methods.math import (
+from .verifications_methods.iamhuman import IAmHumanButton
-    MathButtonsVerification,
+from .verifications_methods.math import MathButtonsVerification
-    MathInputVerificationMethod,
+from .verifications_methods.question import QuestionButtonsVerification
-)
-from .verifications_methods.question import (
-    QuestionButtonsVerification,
-    QuestionInputVerification,
-)
 
+# По хорошему, надо вынести в конфиг, но пока оставим так.
 verification_methods = [
     IAmHumanButton(),
-    IAmHumanInput(),
+    # IAmHumanInput(),
     MathButtonsVerification(),
-    MathInputVerificationMethod(),
+    # MathInputVerificationMethod(),
     QuestionButtonsVerification(),
-    QuestionInputVerification(),
+    # QuestionInputVerification(),
 ]
 
 verification_tasks = {}
 
 
 async def new_member_handler(event: ChatMemberUpdated, bot: Bot):
-    if event.new_chat_member.status == "member":
+    # НЕ СРАБОТАЕТ, ЕСЛИ ЧЕЛОВЕК УЖЕ ОГРАНИЧЕН В ПРАВАХ (RESTRICTED)
+    if event.new_chat_member.status == ChatMemberStatus.MEMBER:
         user_id = event.from_user.id
         chat_id = event.chat.id
 
-        method = random.choice(verification_methods)  # nosec
+        method: VerificationMethod = random.choice(verification_methods)  # nosec
+
+        await method.pre_task(
+            chat_id,
+            user_id,
+            bot,
+        )
+
         task_data = await method.create_task(event, bot)
 
         task_data["user_id"] = user_id
         task_data["chat_id"] = chat_id
         task_data["method"] = method
+        task_data["timeout"] = method.timeout(task_data)
 
         task = asyncio.create_task(verify_timeout(bot, task_data))
 
@@ -51,20 +56,33 @@ async def new_member_handler(event: ChatMemberUpdated, bot: Bot):
         }
 
 
+async def left_member_handler(event: ChatMemberUpdated, bot: Bot):
+    user_id = event.from_user.id
+    chat_id = event.chat.id
+
+    if (user_id, chat_id) not in verification_tasks:
+        return
+
+    task = verification_tasks[(user_id, chat_id)]
+    task["task"].cancel()
+    verification_tasks.pop((user_id, chat_id), None)
+    task_data = task["task_data"]
+    method: VerificationMethod = task_data["method"]
+    await method.post_task(task_data, bot, success=False)
+
+
 async def verify_timeout(bot: Bot, task_data: dict):
     try:
         chat_id = task_data["chat_id"]
         user_id = task_data["user_id"]
+        method: VerificationMethod = task_data["method"]
 
-        await asyncio.sleep(30)
+        await asyncio.sleep(task_data["timeout"])
-        try:
+
-            if "message_id" in task_data:
+        await method.post_task(task_data, success=False)
-                await bot.delete_message(chat_id, task_data["message_id"])
-        except TelegramBadRequest:
-            return
 
         chat_member = await bot.get_chat_member(chat_id, user_id)
-        if chat_member.status == "member":
+        if chat_member.status in [ChatMemberStatus.MEMBER, ChatMemberStatus.RESTRICTED]:
             await bot.ban_chat_member(chat_id, user_id)
 
     except Exception as e:
@@ -79,51 +97,80 @@ async def handle_inline_button_verification(
     user_id = callback_data.user_id
     chat_id = callback_data.chat_id
 
-    if callback_query.from_user.id == user_id:
+    if callback_query.from_user.id != user_id:
-        if (user_id, chat_id) in verification_tasks:
-            task = verification_tasks[(user_id, chat_id)]
-            task_data = task["task_data"]
-            method = task_data["method"]
-            task_data["answer"] = callback_data.answer
-
-            result = await method.verify(task_data)
-
-            if result:
-                task["task"].cancel()
-                await bot.delete_message(chat_id, callback_query.message.message_id)
-            else:
-                await callback_query.answer("Неправильный ответ!", show_alert=True)
-        pass
-    else:
         await callback_query.answer("Эта кнопка не для вас!", show_alert=True)
+        return
+
+    if (user_id, chat_id) not in verification_tasks:
+        await callback_query.answer()
+        return
+
+    verification_task = verification_tasks[(user_id, chat_id)]
+    task_data = verification_task["task_data"]
+    method: VerificationMethod = task_data["method"]
+    task_data["answer"] = callback_data.answer
+
+    result = await method.verify(task_data)
+
+    if result:
+        verification_task["task"].cancel()
+        await method.post_task(task_data, bot)
+        return
+
+    if "attempts_count" not in task_data:
+        await callback_query.answer("Неправильный ответ!", show_alert=True)
+        return
+
+    attempts_count = task_data["attempts_count"]
+
+    verification_task["task"].cancel()
+    attempts_count = attempts_count - 1
+
+    if attempts_count > 0:
+        await callback_query.answer(
+            "Неправильный ответ! "
+            + "У вас еще "
+            + get_plural_form(attempts_count, "попытка", "попытки", "попыток"),
+            show_alert=True,
+        )
+        task_data["timeout"] = method.timeout(task_data)
+        task_data["attempts_count"] = attempts_count
+    else:
+        task_data["timeout"] = 0
+        await callback_query.answer()
+
+    task = asyncio.create_task(verify_timeout(bot, task_data))
+    verification_task["task"] = task
+
+    # Эта строчка нужна, т.к. во время cancel происходит pop
+    verification_tasks[(user_id, chat_id)] = verification_task
 
 
 async def handle_input_verification(message: types.Message, bot: Bot):
     user_id = message.from_user.id
     chat_id = message.chat.id
 
-    if (user_id, chat_id) in verification_tasks:
+    if (user_id, chat_id) not in verification_tasks:
-        task = verification_tasks[(user_id, chat_id)]
+        return
-        task_data = task["task_data"]
+    task = verification_tasks[(user_id, chat_id)]
-        method = task_data["method"]
+    task_data = task["task_data"]
-        task_data["answer"] = message.text
+    method: VerificationMethod = task_data["method"]
+    task_data["answer"] = message.text
+    task_data["answer_message_id"] = message.message_id
 
-        result = await method.verify(task_data)
+    result = await method.verify(task_data)
 
-        if result:
+    if result:
-            task["task"].cancel()
+        task["task"].cancel()
-            await bot.delete_message(chat_id, task_data["message_id"])
+        await method.post_task(task_data, bot)
-            await bot.delete_message(chat_id, message.message_id)
-
-        pass
 
 
 async def module_init():
     router = Router()
 
-    router.chat_member(ChatMemberUpdatedFilter(IS_NOT_MEMBER >> IS_MEMBER))(
+    router.chat_member(ChatMemberUpdatedFilter(JOIN_TRANSITION))(new_member_handler)
-        new_member_handler
+
-    )
+    router.chat_member(ChatMemberUpdatedFilter(LEAVE_TRANSITION))(left_member_handler)
 
     router.callback_query(VerificationCallback.filter())(
         handle_inline_button_verification

src/ocab_modules/ocab_modules/standard/welcome/utils.py

@@ -0,0 +1,9 @@
+def get_plural_form(number, singular, genitive_singular, plural):
+    if 11 <= number % 100 <= 19:
+        return f"{number} {plural}"
+    elif number % 10 == 1:
+        return f"{number} {singular}"
+    elif 2 <= number % 10 <= 4:
+        return f"{number} {genitive_singular}"
+    else:
+        return f"{number} {plural}"

src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/base.py

@@ -1,18 +1,133 @@
+import time
+
+from aiogram import Bot
+from aiogram.exceptions import TelegramBadRequest
 from aiogram.filters.callback_data import CallbackData
+from aiogram.types import ChatMemberUpdated, ChatPermissions
+
+from .utils import user_mention
+
+
+async def mute_user(chat_id, user_id, until, bot: Bot):
+    end_time = until + int(time.time())
+    await bot.restrict_chat_member(
+        chat_id,
+        user_id,
+        until_date=end_time,
+        use_independent_chat_permissions=True,
+        permissions=ChatPermissions(
+            can_send_messages=False,
+            can_send_audios=False,
+            can_send_documents=False,
+            can_send_photos=False,
+            can_send_videos=False,
+            can_send_video_notes=False,
+            can_send_voice_notes=False,
+            can_send_polls=False,
+            can_send_other_messages=False,
+            can_add_web_page_previews=False,
+            can_change_info=False,
+            can_invite_users=False,
+            can_pin_messages=False,
+            can_manage_topics=False,
+        ),
+    )
+
+
+async def unmute_user(chat_id, user_id, bot: Bot):
+    await bot.restrict_chat_member(
+        chat_id,
+        user_id,
+        use_independent_chat_permissions=True,
+        permissions=ChatPermissions(
+            can_send_messages=True,
+            can_send_audios=True,
+            can_send_documents=True,
+            can_send_photos=True,
+            can_send_videos=True,
+            can_send_video_notes=True,
+            can_send_voice_notes=True,
+            can_send_polls=True,
+            can_send_other_messages=True,
+            can_add_web_page_previews=True,
+            can_change_info=True,
+            can_invite_users=True,
+            can_pin_messages=True,
+            can_manage_topics=True,
+        ),
+    )
 
 
 class VerificationMethod:
-    pass
+
+    def timeout(self, task_data=None) -> int:
+        """
+        Время ожидания
+        """
+        return 30
+
+    def method_name(self):
+        pass
+
+    async def pre_task(self, chat_id, user_id, bot: Bot):
+        pass
+
+    async def create_task(self, event: ChatMemberUpdated, bot: Bot):
+        pass
+
+    async def post_task(self, task_data, bot: Bot, success=True):
+        pass
+
+    async def verify(self, task_data):
+        pass
 
 
 class InputVerificationMethod(VerificationMethod):
-    def verify(input_value: str, task_data):
+    async def post_task(self, task_data, bot: Bot, success=True, user=None):
-        pass
+        chat_id = task_data["chat_id"]
+        message_id = task_data["message_id"]
+        answer_message_id = task_data["answer_message_id"]
+
+        await bot.delete_message(chat_id, message_id)
+        await bot.delete_message(chat_id, answer_message_id)
+
+        if not success or user is None:
+            return
+
+        await bot.send_message(
+            chat_id,
+            f"{user_mention(user)}, успешно прошли проверку! "
+            "Пожалуйста, соблюдайте правила группы.",
+        )
 
 
 class InlineButtonVerificationMethod(VerificationMethod):
-    def verify(input_value: str, task_data):
+    async def pre_task(self, chat_id, user_id, bot: Bot):
-        pass
+        try:
+            await mute_user(chat_id, user_id, 0, bot)
+        except TelegramBadRequest:
+            pass
+
+    async def post_task(self, task_data, bot: Bot, success=True, user=None):
+        user_id = task_data["user_id"]
+        chat_id = task_data["chat_id"]
+        message_id = task_data["message_id"]
+
+        await bot.delete_message(chat_id, message_id)
+
+        try:
+            await unmute_user(chat_id, user_id, bot)
+        except TelegramBadRequest:
+            pass
+
+        if not success or user is None:
+            return
+
+        await bot.send_message(
+            chat_id,
+            f"{user_mention(user)}, успешно прошли проверку! "
+            "Пожалуйста, соблюдайте правила группы.",
+        )
 
 
 class VerificationCallback(CallbackData, prefix="verify"):

src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/iamhuman.py

@@ -1,6 +1,7 @@
 import random
 
 from aiogram import Bot
+from aiogram.enums import ParseMode
 from aiogram.types import ChatMemberUpdated, InlineKeyboardButton, InlineKeyboardMarkup
 
 from .base import (
@@ -8,6 +9,7 @@ from .base import (
     InputVerificationMethod,
     VerificationCallback,
 )
+from .utils import user_mention
 
 
 class IAmHumanButton(InlineButtonVerificationMethod):
@@ -36,9 +38,10 @@ class IAmHumanButton(InlineButtonVerificationMethod):
 
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! ",
             "Нажмите кнопку, чтобы подтвердить, что вы не робот.",
             reply_markup=keyboard,
+            parse_mode=ParseMode.HTML,
         )
 
         return {"message_id": message.message_id}
@@ -62,8 +65,9 @@ class IAmHumanInput(InputVerificationMethod):
         text = self.get_text()
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! "
             f'Напишите "{text}", чтобы подтвердить, что вы не робот.',
+            parse_mode=ParseMode.HTML,
         )
         return {"message_id": message.message_id, "correct": text}
 

src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/math.py

@@ -1,6 +1,7 @@
 import random
 
 from aiogram import Bot
+from aiogram.enums import ParseMode
 from aiogram.types import ChatMemberUpdated, InlineKeyboardButton, InlineKeyboardMarkup
 
 from .base import (
@@ -8,6 +9,7 @@ from .base import (
     InputVerificationMethod,
     VerificationCallback,
 )
+from .utils import user_mention
 
 
 class MathInputVerificationMethod(InputVerificationMethod):
@@ -34,9 +36,11 @@ class MathInputVerificationMethod(InputVerificationMethod):
         problem, answer = self.generate_math_problem()
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! "
             "Решите простую математическую задачу, "
-            f"чтобы подтвердить, что вы не робот: {problem} = ?",
+            "чтобы подтвердить, что вы не робот:\\n"
+            f"{problem} = ?",
+            parse_mode=ParseMode.HTML,
         )
         return {"message_id": message.message_id, "correct": answer}
 
@@ -74,7 +78,7 @@ class MathButtonsVerification(InlineButtonVerificationMethod):
         options = [correct_answer]
         while len(options) < 4:
             wrong_answer = random.randint(
-                max(1, correct_answer - 5), correct_answer + 5
+                correct_answer - 5, correct_answer + 5
             )  # nosec
             if wrong_answer not in options:
                 options.append(wrong_answer)
@@ -96,13 +100,19 @@ class MathButtonsVerification(InlineButtonVerificationMethod):
 
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! "
             "Решите простую математическую задачу, "
-            f"чтобы подтвердить, что вы не робот: {problem} = ?",
+            "чтобы подтвердить, что вы не робот:\\n"
+            f"{problem} = ?",
             reply_markup=keyboard,
+            parse_mode=ParseMode.HTML,
         )
 
-        return {"message_id": message.message_id, "correct": str(correct_answer)}
+        return {
+            "message_id": message.message_id,
+            "correct": str(correct_answer),
+            "attempts_count": 2,
+        }
 
     async def verify(self, task_data):
         correct: str = task_data["correct"]

src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/question.py

@@ -1,8 +1,11 @@
 import random
 
 from aiogram import Bot
+from aiogram.enums import ParseMode
 from aiogram.types import ChatMemberUpdated, InlineKeyboardButton, InlineKeyboardMarkup
 
+from ocab_modules.standard.welcome.verifications_methods.utils import user_mention
+
 from .base import (
     InlineButtonVerificationMethod,
     InputVerificationMethod,
@@ -54,9 +57,10 @@ class QuestionInputVerification(InputVerificationMethod):
         question, answer, _ = self.get_random_question()
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! "
             "Пожалуйста, ответьте на следующий вопрос, "
             f"чтобы подтвердить, что вы не робот: {question}",
+            parse_mode=ParseMode.HTML,
         )
         return {"message_id": message.message_id, "correct": answer.lower()}
 
@@ -94,23 +98,25 @@ class QuestionButtonsVerification(InlineButtonVerificationMethod):
                             user_id=user_id, chat_id=chat_id, answer=str(i)
                         ).pack(),
                     )
-                    for i, option in enumerate(options)
                 ]
+                for i, option in enumerate(options)
             ]
         )
 
         message = await bot.send_message(
             chat_id,
-            f"Привет, {event.from_user.first_name}! "
+            f"Привет, {user_mention(event.from_user)}! "
             "Пожалуйста, ответьте на следующий вопрос, "
             f"чтобы подтвердить, что вы не робот: {question}",
             reply_markup=keyboard,
+            parse_mode=ParseMode.HTML,
         )
 
         return {
             "message_id": message.message_id,
             "correct": correct_answer,
             "options": options,
+            "attempts_count": 2,
         }
 
     async def verify(self, task_data):

src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/utils.py

@@ -0,0 +1,11 @@
+from aiogram.enums import ParseMode
+from aiogram.types import User
+
+
+def user_mention(user: User, mode=ParseMode.HTML):
+    if mode == ParseMode.HTML:
+        return f"<a href='tg://user?id={user.id}'>{user.first_name}</a>"
+    elif mode == ParseMode.MARKDOWN:
+        return f"[{user.first_name}](tg://user?id={user.id})"
+    else:
+        raise ValueError(f"Unknown parse mode {mode}")