tgcaptchabot/main.py

import os
import time
import random
import string
import asyncio
from datetime import datetime, timedelta

import aiohttp
from dotenv import load_dotenv
from telethon import TelegramClient, events
from telethon.utils import get_peer_id, get_display_name
from telethon.tl.functions.channels import EditBannedRequest
from telethon.tl.types import ChatBannedRights
from captcha.image import ImageCaptcha


def get_user_name(user):
    full_name = get_display_name(user)
    if not full_name:
        full_name = user.username

    if not full_name:
        full_name = "user"

    return full_name


def get_link_to_user(user):
    full_name = get_user_name(user)

    if user.username:
        return f"[{full_name}](@{user.username})"

    return f"[{full_name}](tg://user?id={user.id})"


load_dotenv()

API_ID = int(os.getenv("API_ID"))
API_HASH = os.getenv("API_HASH")
BOT_TOKEN = os.getenv("BOT_TOKEN")
ALLOWED_GROUPS = set(
    map(lambda cid: int("-100" + cid), os.getenv("ALLOWED_GROUPS").split(","))
)
CAPTCHA_WAIT_TIME = int(os.getenv("CAPTCHA_WAIT_TIME", 80))
MAX_CAPTCHA_ATTEMPTS = int(os.getenv("MAX_CAPTCHA_ATTEMPTS", 8))
USE_LOLS_API = os.getenv("USE_LOLS_API", "yes") == "yes"

bot = TelegramClient("captcha_bot", API_ID, API_HASH).start(bot_token=BOT_TOKEN)
pending = {peer_id: {} for peer_id in ALLOWED_GROUPS}
whitelist = {}


async def lols_check(user_id):
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(
                f"https://api.lols.bot/account?id={user_id}"
            ) as response:
                data = await response.json()

                banned = data.get("banned", False)
                scammer = data.get("scammer", False)

                return not banned and not scammer
    except Exception:
        pass

    return True


class PendingUser:
    def __init__(self, peer_id, user_id):
        self.peer_id = peer_id
        self.user_id = user_id

        self.ban_flag = False

        self.text = None
        self.message_id = None

        self.captcha = None
        self.captcha_im = None

        self.join_ts = time.time()
        self.retries = 0

    @property
    def expired(self):
        return time.time() - self.join_ts >= CAPTCHA_WAIT_TIME

    async def start(self, event):
        if USE_LOLS_API:
            if not await lols_check(self.user_id):
                await self.kick(forever=True)

                return

        await self.update_captcha()

        if self.user_id not in pending[self.peer_id]:
            return

        self.text = f"Welcome, {get_link_to_user(event.user)}! Please write a message with the **four latin letters (case insensitive)** that appear in this image to verify that you are a human. If you don't solve the captcha in **{CAPTCHA_WAIT_TIME}** seconds, you will be automatically kicked out of the group."

        message = await bot.send_message(
            event.chat_id,
            self.text,
            file=self.captcha_im,
        )

        self.message_id = message

    async def update_captcha(self, only_text=False):
        if not only_text:
            self.captcha = "".join(random.sample(string.ascii_uppercase, 4))
            self.captcha_im = ImageCaptcha().generate(self.captcha)

            self.captcha_im.name = "captcha.png"

        if self.message_id is not None:
            try:
                attempts = MAX_CAPTCHA_ATTEMPTS - self.retries

                message = await bot.edit_message(
                    self.peer_id,
                    self.message_id,
                    text=self.text
                    + f"\n\nYou have **{attempts}** attempt{'' if attempts == 1 else 's'} left.",
                    file=None if only_text else self.captcha_im,
                )

                self.message_id = message.id
            except Exception:
                pass

    async def check_captcha(self, captcha):
        if self.ban_flag:
            await self.kick()

            return False

        captcha = captcha.replace(" ", "").strip().upper()

        if captcha == self.captcha:
            del pending[self.peer_id][self.user_id]

            try:
                await bot.delete_messages(self.peer_id, self.message_id)
            except Exception:
                pass

            return True

        self.retries += 1

        if self.retries > MAX_CAPTCHA_ATTEMPTS:
            await self.kick()
        else:
            await self.update_captcha(
                only_text=not (
                    len(captcha) == 4
                    and all(map(lambda c: c in string.ascii_uppercase, captcha))
                )
            )

        return False

    async def remove(self):
        del pending[self.peer_id][self.user_id]

        if self.message_id is not None:
            try:
                await bot.delete_messages(self.peer_id, self.message_id)
            except Exception:
                pass

    async def kick(self, forever=False):
        if self.user_id not in pending[self.peer_id]:
            return

        await self.remove()

        try:
            await bot(
                EditBannedRequest(
                    self.peer_id,
                    self.user_id,
                    ChatBannedRights(
                        until_date=None
                        if forever
                        else datetime.now() + timedelta(minutes=5),
                        view_messages=True,
                    ),
                )
            )
        except Exception:
            pass


@bot.on(events.ChatAction)
async def handler(event):
    if event.chat_id not in ALLOWED_GROUPS:
        return

    peer_id = event.chat_id

    if (event.user_left or event.user_kicked) and event.user_id in pending[peer_id]:
        await pending[peer_id][event.user_id].remove()

        return

    if not event.user_joined:
        return

    if event.user_id in whitelist:
        return

    pending[peer_id][event.user_id] = PendingUser(
        peer_id=event.chat_id, user_id=event.user_id
    )

    await pending[peer_id][event.user_id].start(event)


@bot.on(events.NewMessage)
async def handler(event):
    peer_id = get_peer_id(event.peer_id)
    if peer_id not in ALLOWED_GROUPS:
        return

    if not event.sender:
        return

    pending_user = pending[peer_id].get(event.sender.id)
    if not pending_user:
        return

    try:
        await event.delete()
    except Exception:
        pass

    if time.time() - pending_user.join_ts <= 1.0:
        pending_user.ban_flag = True

    if pending_user.message_id is None:
        return

    text = getattr(event, "text", "")

    if await pending_user.check_captcha(text):
        whitelist[event.sender.id] = time.time()

        for peer_id in pending:
            if event.sender.id in pending[peer_id]:
                await pending[peer_id][event.sender.id].remove()


async def check_pending():
    while True:
        for peer_id in dict(pending):
            for user_id in pending[peer_id]:
                pending_user = pending[peer_id][user_id]

                if pending_user.expired:
                    await pending_user.kick()

        await asyncio.sleep(1)


async def cleanup_whitelist():
    while True:
        for user_id in dict(whitelist):
            ts = whitelist[user_id]

            if time.time() - ts >= 30 * 60:
                del whitelist[user_id]

        await asyncio.sleep(60)


if __name__ == "__main__":
    bot.loop.create_task(check_pending())
    bot.loop.create_task(cleanup_whitelist())
    bot.start()
    bot.run_until_disconnected()