telegram-antispam-bot/bot/handlers.py

from __future__ import annotations

import asyncio
import time

from aiogram import Bot, Dispatcher, F, Router
from aiogram.exceptions import TelegramBadRequest, TelegramForbiddenError
from aiogram.filters import Command, CommandStart
from aiogram.types import (
    CallbackQuery,
    ChatMemberAdministrator,
    ChatMemberOwner,
    ChatMemberUpdated,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    Message,
)

from .config import Config
from .db import Database


router = Router(name="antispam")

_member_status_cache: dict[tuple[int, int], tuple[str, float]] = {}
_MEMBER_STATUS_CACHE_TTL_SECONDS = 5.0


def _is_allowed_chat(cfg: Config, chat_id: int) -> bool:
    return cfg.chat_id is None or cfg.chat_id == chat_id


async def _is_admin(bot: Bot, chat_id: int, user_id: int) -> bool:
    member = await bot.get_chat_member(chat_id, user_id)
    return isinstance(member, (ChatMemberAdministrator, ChatMemberOwner))


async def _get_member_status_cached(bot: Bot, chat_id: int, user_id: int) -> str | None:
    """
    Returns member.status string, or None if lookup failed.
    Cached briefly to avoid per-message API calls during bursts.
    """
    now_mono = time.monotonic()
    key = (chat_id, user_id)
    cached = _member_status_cache.get(key)
    if cached is not None:
        status, valid_until = cached
        if now_mono <= valid_until:
            return status

    try:
        member = await bot.get_chat_member(chat_id, user_id)
    except TelegramBadRequest:
        return None

    status = getattr(member, "status", None)
    if status is None:
        return None

    _member_status_cache[key] = (status, now_mono + _MEMBER_STATUS_CACHE_TTL_SECONDS)
    return status


async def _kick_user(bot: Bot, chat_id: int, user_id: int) -> None:
    # "Kick" in modern Bot API is typically ban + unban.
    try:
        await bot.ban_chat_member(chat_id, user_id)
    except TelegramForbiddenError:
        return
    except TelegramBadRequest:
        return
    try:
        await bot.unban_chat_member(chat_id, user_id, only_if_banned=True)
    except TelegramBadRequest:
        pass


def _challenge_kb(chat_id: int, user_id: int) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup(
        inline_keyboard=[
            [
                InlineKeyboardButton(
                    text="Я человек ✅",
                    callback_data=f"verify:{chat_id}:{user_id}",
                )
            ]
        ]
    )


async def _on_user_joined(
    *,
    bot: Bot,
    cfg: Config,
    db: Database,
    chat_id: int,
    user_id: int,
    user_full_name: str,
) -> None:
    try:
        if await _is_admin(bot, chat_id, user_id):
            return
    except TelegramBadRequest:
        return

    await db.delete_verified(chat_id=chat_id, user_id=user_id)
    await _ensure_challenge(
        bot=bot,
        cfg=cfg,
        db=db,
        chat_id=chat_id,
        user_id=user_id,
        user_full_name=user_full_name,
    )


async def _ensure_challenge(
    *,
    bot: Bot,
    cfg: Config,
    db: Database,
    chat_id: int,
    user_id: int,
    user_full_name: str,
) -> None:
    """
    Ensures there's an active pending verification for the user in this chat.
    Used on join and also on first message (in case bot missed join event).
    """
    now = int(time.time())
    pending = await db.get_pending(chat_id=chat_id, user_id=user_id)
    if pending is not None and pending.expires_at >= now:
        return

    expires_at = now + cfg.verify_window_seconds
    text = (
        f"Привет, {user_full_name}!\n\n"
        "Чтобы писать в чате, нажми кнопку ниже в течение 1 минуты."
    )
    sent = await bot.send_message(chat_id, text, reply_markup=_challenge_kb(chat_id, user_id))
    await db.mark_pending(
        chat_id=chat_id,
        user_id=user_id,
        challenge_message_id=sent.message_id,
        expires_at=expires_at,
    )


async def register_handlers(dp: Dispatcher, *, cfg: Config, db: Database) -> None:
    dp.include_router(router)
    dp["cfg"] = cfg
    dp["db"] = db


@router.message(Command("ping"))
async def ping(message: Message, cfg: Config) -> None:
    if not _is_allowed_chat(cfg, message.chat.id):
        return
    await message.answer("pong")


@router.message(CommandStart())
async def start(message: Message) -> None:
    if message.chat.type != "private":
        return

    await message.answer(
        "Привет! Это личная переписка с ботом — здесь нет никаких настроек.\n\n"
        "Чтобы я работал как антиспам, добавьте меня администратором в группу и выдайте права:\n\n"
        "- удалять сообщения\n"
        "- банить (кикать) пользователей\n\n"
        "Спасибо!"
    )


@router.message(Command("chatid"))
async def chatid(message: Message, cfg: Config) -> None:
    if not _is_allowed_chat(cfg, message.chat.id):
        return
    await message.answer(f"chat_id: {message.chat.id}")


@router.message(F.new_chat_members)
async def on_new_members(message: Message, bot: Bot, cfg: Config, db: Database) -> None:
    chat_id = message.chat.id
    if not _is_allowed_chat(cfg, chat_id):
        return

    for user in message.new_chat_members:
        if user.is_bot:
            continue
        await _on_user_joined(
            bot=bot,
            cfg=cfg,
            db=db,
            chat_id=chat_id,
            user_id=user.id,
            user_full_name=user.full_name,
        )


@router.chat_member()
async def on_chat_member_updated(
    event: ChatMemberUpdated,
    bot: Bot,
    cfg: Config,
    db: Database,
) -> None:
    """
    Supergroups (especially linked discussion groups) often hide join service messages.
    In that case Telegram sends chat_member updates instead of new_chat_members.
    """
    chat_id = event.chat.id
    if not _is_allowed_chat(cfg, chat_id):
        return

    user = event.new_chat_member.user
    if user.is_bot:
        return

    old_status = event.old_chat_member.status
    new_status = event.new_chat_member.status
    if old_status not in ("left", "kicked") or new_status not in ("member", "restricted"):
        return

    await _on_user_joined(
        bot=bot,
        cfg=cfg,
        db=db,
        chat_id=chat_id,
        user_id=user.id,
        user_full_name=user.full_name,
    )


@router.callback_query(F.data.startswith("verify:"))
async def on_verify(
    callback: CallbackQuery,
    bot: Bot,
    cfg: Config,
    db: Database,
) -> None:
    parts = (callback.data or "").split(":")
    if len(parts) != 3:
        await callback.answer("Ошибка кнопки.", show_alert=True)
        return

    chat_id = int(parts[1])
    target_user_id = int(parts[2])

    if not _is_allowed_chat(cfg, chat_id):
        await callback.answer()
        return

    if callback.message is None or callback.message.chat.id != chat_id:
        await callback.answer("Кнопка не из этого чата.", show_alert=True)
        return

    if callback.from_user.id != target_user_id:
        await callback.answer("Кнопка только для нового участника.", show_alert=True)
        return

    try:
        if await _is_admin(bot, chat_id, target_user_id):
            await callback.answer("Админам проверка не нужна.")
            return
    except TelegramBadRequest:
        pass

    pending = await db.get_pending(chat_id=chat_id, user_id=target_user_id)
    now = int(time.time())
    if pending is None:
        await callback.answer("Проверка уже неактуальна.")
        return
    if pending.expires_at < now:
        await callback.answer("Время истекло. Зайдите в чат заново.", show_alert=True)
        return

    await db.mark_verified(chat_id=chat_id, user_id=target_user_id)

    # Delete challenge message
    try:
        await bot.delete_message(chat_id, pending.challenge_message_id)
    except TelegramBadRequest:
        pass

    thanks = await bot.send_message(chat_id, "Спасибо! Проверка пройдена.")
    await callback.answer("Готово.")

    async def _delete_thanks() -> None:
        await asyncio.sleep(cfg.thanks_ttl_seconds)
        try:
            await bot.delete_message(chat_id, thanks.message_id)
        except TelegramBadRequest:
            pass

    asyncio.create_task(_delete_thanks())


@router.message()
async def anti_spam_all_messages(message: Message, bot: Bot, cfg: Config, db: Database) -> None:
    if message.chat is None:
        return
    chat_id = message.chat.id
    if not _is_allowed_chat(cfg, chat_id):
        return

    # Discussion groups can receive automatic forwards from a linked channel.
    # These messages are not user spam and must never be deleted.
    if getattr(message, "is_automatic_forward", False):
        return
    if message.sender_chat is not None and getattr(message.sender_chat, "type", None) == "channel":
        return

    user = message.from_user
    if user is None or user.is_bot:
        return

    # Don't touch admins/owner.
    try:
        if await _is_admin(bot, chat_id, user.id):
            return
    except TelegramBadRequest:
        return

    # Key rule for discuss comments:
    # If user is not a member of the group -> delete their message.
    status = await _get_member_status_cached(bot, chat_id, user.id)
    if status in ("left", "kicked"):
        try:
            await message.delete()
        except TelegramBadRequest:
            pass
        return
    if status is None:
        return

    # For new joins: delete messages until verified.
    if await db.is_verified(chat_id=chat_id, user_id=user.id):
        return

    pending = await db.get_pending(chat_id=chat_id, user_id=user.id)
    if pending is None:
        # Not in pending list => do not moderate.
        # This prevents touching existing members in large chats,
        # and only targets users caught at join time.
        return

    try:
        await message.delete()
    except TelegramBadRequest:
        pass


async def sweeper_loop(*, bot: Bot, cfg: Config, db: Database) -> None:
    while True:
        await asyncio.sleep(2)
        now = int(time.time())
        expired = await db.get_expired_pending(now=now)
        for p in expired:
            if not _is_allowed_chat(cfg, p.chat_id):
                await db.delete_pending(chat_id=p.chat_id, user_id=p.user_id)
                continue

            # Delete the challenge message (if still exists)
            try:
                await bot.delete_message(p.chat_id, p.challenge_message_id)
            except TelegramBadRequest:
                pass

            # Kick user for a new attempt.
            await _kick_user(bot, p.chat_id, p.user_id)

            await db.delete_pending(chat_id=p.chat_id, user_id=p.user_id)