EmojiCaptchaBot/plugins/handlers/user_join.py

import asyncio
import contextlib
import logging
from datetime import datetime, timedelta

from pyrogram import filters
from pyrogram.types import (
    ChatPermissions,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    Message,
)

from classes.pyroclient import PyroClient
from classes.pyrogroup import PyroGroup
from classes.pyrouser import PyroUser
from modules.database import col_schedule
from modules.kicker import kick_unstarted

logger = logging.getLogger(__name__)


@PyroClient.on_message(
    filters.new_chat_members & filters.group & ~filters.me & ~filters.bot  # type: ignore
)
async def handler_user_join(app: PyroClient, message: Message):
    group = await PyroGroup.create_if_not_exists(message.chat.id, None, True)
    locale = group.select_locale(app, message.from_user)

    if (
        app.config["whitelist"]["enabled"]
        and group.id not in app.config["whitelist"]["groups"]
    ):
        logger.info(
            "User %s has joined the group %s, but it's not whitelisted, ignoring.",
            message.from_user.id,
            group.id,
        )
        return

    logger.info("User %s has joined the group %s", message.from_user.id, group.id)

    await message.delete()

    # If user has already failed the test and joined once more
    with contextlib.suppress(KeyError):
        user = await app.find_user(message.from_user, group.id)
        if user.failed is True:
            logger.info(
                "User %s has previously failed the captcha, kicking and banning him",
                user.id,
            )
            banned = await app.ban_chat_member(group.id, user.id)

            if isinstance(banned, Message):
                await banned.delete()

            return

    await app.restrict_chat_member(
        chat_id=group.id,
        user_id=message.from_user.id,
        permissions=ChatPermissions(can_send_messages=False),
    )

    user = await PyroUser.create_if_not_exists(message.from_user.id, group.id)

    if user.mistakes > 0 or user.score > 0:
        await user.set_score(0)
        await user.set_mistakes(0)

    if len(user.selected) > 0:
        await user.set_selected([])

    await asyncio.sleep(2)
    verification_request = await app.send_message(
        chat_id=group.id,
        text=app._("welcome", "messages", locale=locale).format(
            mention=message.from_user.mention
        ),
        reply_markup=InlineKeyboardMarkup(
            [
                [
                    InlineKeyboardButton(
                        app._("verify", "buttons", locale=locale),
                        callback_data=f"verify:{user.id}",
                    )
                ],
                [
                    InlineKeyboardButton(
                        app._("ban", "buttons", locale=locale),
                        callback_data=f"ban:{user.id}",
                    )
                ],
            ],
        ),
    )

    if app.scheduler is not None:
        job = app.scheduler.add_job(
            kick_unstarted,
            "date",
            [app, user.id, verification_request.chat.id, verification_request.id],
            run_date=datetime.now() + timedelta(seconds=app.config["timeouts"]["join"]),
        )
        await col_schedule.insert_one(
            {"user": user.id, "group": group.id, "job_id": job.id}
        )