diff --git a/src/ocab_modules/ocab_modules/standard/command_helper/main.py b/src/ocab_modules/ocab_modules/standard/command_helper/main.py index d33d1ba..f7c417e 100644 --- a/src/ocab_modules/ocab_modules/standard/command_helper/main.py +++ b/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 ( - set_my_commands, log -) +from ocab_core.modules_system.public_api import set_my_commands 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, diff --git a/src/ocab_modules/ocab_modules/standard/welcome/main.py b/src/ocab_modules/ocab_modules/standard/welcome/main.py index 1cbb27c..8df6a5d 100644 --- a/src/ocab_modules/ocab_modules/standard/welcome/main.py +++ b/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.filters import IS_MEMBER, IS_NOT_MEMBER, ChatMemberUpdatedFilter +from aiogram.enums import ChatMemberStatus +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 .verifications_methods.iamhuman import IAmHumanButton, IAmHumanInput -from .verifications_methods.math import ( - MathButtonsVerification, - MathInputVerificationMethod, -) -from .verifications_methods.question import ( - QuestionButtonsVerification, - QuestionInputVerification, -) +from .utils import get_plural_form +from .verifications_methods.base import VerificationCallback, VerificationMethod +from .verifications_methods.iamhuman import IAmHumanButton +from .verifications_methods.math import MathButtonsVerification +from .verifications_methods.question import QuestionButtonsVerification +# По хорошему, надо вынести в конфиг, но пока оставим так. 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) - try: - if "message_id" in task_data: - await bot.delete_message(chat_id, task_data["message_id"]) - except TelegramBadRequest: - return + await asyncio.sleep(task_data["timeout"]) + + await method.post_task(task_data, success=False) 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 (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: + if callback_query.from_user.id != user_id: 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: - task = verification_tasks[(user_id, chat_id)] - task_data = task["task_data"] - method = task_data["method"] - task_data["answer"] = message.text + if (user_id, chat_id) not in verification_tasks: + return + task = verification_tasks[(user_id, chat_id)] + task_data = task["task_data"] + 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: - task["task"].cancel() - await bot.delete_message(chat_id, task_data["message_id"]) - await bot.delete_message(chat_id, message.message_id) - - pass + if result: + task["task"].cancel() + await method.post_task(task_data, bot) async def module_init(): router = Router() - router.chat_member(ChatMemberUpdatedFilter(IS_NOT_MEMBER >> IS_MEMBER))( - new_member_handler - ) + router.chat_member(ChatMemberUpdatedFilter(JOIN_TRANSITION))(new_member_handler) + + router.chat_member(ChatMemberUpdatedFilter(LEAVE_TRANSITION))(left_member_handler) router.callback_query(VerificationCallback.filter())( handle_inline_button_verification diff --git a/src/ocab_modules/ocab_modules/standard/welcome/utils.py b/src/ocab_modules/ocab_modules/standard/welcome/utils.py new file mode 100644 index 0000000..b3288b4 --- /dev/null +++ b/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}" diff --git a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/base.py b/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/base.py index aba1f27..99b7f35 100644 --- a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/base.py +++ b/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): - pass + async def post_task(self, task_data, bot: Bot, success=True, user=None): + 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): - pass + async def pre_task(self, chat_id, user_id, bot: Bot): + 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"): diff --git a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/iamhuman.py b/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/iamhuman.py index acd50da..d6c3420 100644 --- a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/iamhuman.py +++ b/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} diff --git a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/math.py b/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/math.py index 97e01ed..b2b812b 100644 --- a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/math.py +++ b/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"] diff --git a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/question.py b/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/question.py index a7b51ea..785af3f 100644 --- a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/question.py +++ b/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): diff --git a/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/utils.py b/src/ocab_modules/ocab_modules/standard/welcome/verifications_methods/utils.py new file mode 100644 index 0000000..41e9a61 --- /dev/null +++ b/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"{user.first_name}" + elif mode == ParseMode.MARKDOWN: + return f"[{user.first_name}](tg://user?id={user.id})" + else: + raise ValueError(f"Unknown parse mode {mode}")