fix: align captcha settings and pending verification timeout

AI-Co-Authored-By: Codex

tests/test_captcha_verification.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import asyncio
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from pyrogram.enums import ChatMemberStatus
+
+
+@pytest.fixture
+async def greetings(app_modules, monkeypatch):
+    app_modules.wbb.WELCOME_DELAY_KICK_SEC = 300
+    permissions = ModuleType("wbb.core.decorators.permissions")
+    permissions.adminsOnly = lambda _permission: lambda function: function
+    notes = ModuleType("wbb.modules.notes")
+    notes.extract_urls = lambda _markup: []
+    monkeypatch.setitem(sys.modules, permissions.__name__, permissions)
+    monkeypatch.setitem(sys.modules, notes.__name__, notes)
+    module = app_modules.load("wbb.modules.greetings")
+    await asyncio.sleep(0)
+    return module
+
+
+async def test_captcha_defaults_off_and_admin_toggle_matches_bot(app_modules, monkeypatch):
+    dbfunctions = app_modules.load("wbb.utils.dbfunctions")
+    service = app_modules.load("wbb.services.chat_management")
+    monkeypatch.setattr(service, "ensure_permission", AsyncMock())
+
+    assert await dbfunctions.is_captcha_on(-100) is False
+    assert (await service.get_automation_settings(-100))["captcha_enabled"] is False
+    await app_modules.wbb.db.managed_chat_settings.insert_one(
+        {"bot_id": "primary", "chat_id": -200, "welcome": {"enabled": True}}
+    )
+    assert await dbfunctions.is_captcha_on(-200) is False
+    assert (await service.get_automation_settings(-200))["captcha_enabled"] is False
+
+    await service.apply_automation_settings(-100, {"captcha_enabled": True})
+    assert await dbfunctions.is_captcha_on(-100) is True
+    await dbfunctions.captcha_off(-100)  # The /captcha command uses the same helper.
+    assert await dbfunctions.is_captcha_on(-100) is False
+    assert (await service.get_automation_settings(-100))["captcha_enabled"] is False
+    await dbfunctions.captcha_on(-100)
+    assert (await service.get_automation_settings(-100))["captcha_enabled"] is True
+
+
+async def test_solved_captcha_is_saved_for_each_member(app_modules):
+    dbfunctions = app_modules.load("wbb.utils.dbfunctions")
+    await dbfunctions.save_captcha_solved(-100, 11)
+    await dbfunctions.save_captcha_solved(-100, 12)
+    assert await dbfunctions.has_solved_captcha_once(-100, 11)
+    assert await dbfunctions.has_solved_captcha_once(-100, 12)
+
+
+async def test_captcha_command_updates_admin_setting(greetings):
+    reply_text = AsyncMock()
+    message = SimpleNamespace(
+        chat=SimpleNamespace(id=-100),
+        command=["captcha", "开启"],
+        text="/captcha 开启",
+        reply_text=reply_text,
+    )
+    await greetings.captcha_state(None, message)
+    assert await greetings.is_captcha_on(-100)
+    message.command = ["captcha", "关闭"]
+    message.text = "/captcha 关闭"
+    await greetings.captcha_state(None, message)
+    assert not await greetings.is_captcha_on(-100)
+
+
+async def test_passed_verification_timer_does_not_ban_later_manual_mute(
+    greetings, monkeypatch
+):
+    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
+    chat = SimpleNamespace(
+        id=-100,
+        get_member=AsyncMock(
+            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
+        ),
+        ban_member=AsyncMock(),
+    )
+    message = SimpleNamespace(chat=chat, delete=AsyncMock())
+    user = SimpleNamespace(id=11)
+    await greetings.update_captcha_cache(
+        [{"chat_id": -100, "user_id": 11, "attempt_id": "old"}]
+    )
+    await greetings.save_captcha_solved(-100, 11)
+    await greetings.kick_restricted_after_delay(300, message, user, "old")
+
+    chat.ban_member.assert_not_awaited()
+    message.delete.assert_not_awaited()
+    assert await greetings.get_captcha_cache() == []
+
+
+async def test_success_callback_invalidates_timer_before_manual_mute(
+    greetings, monkeypatch
+):
+    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
+    monkeypatch.setattr(greetings, "send_welcome_message", AsyncMock())
+    chat = SimpleNamespace(
+        id=-100,
+        unban_member=AsyncMock(),
+        get_member=AsyncMock(
+            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
+        ),
+        ban_member=AsyncMock(),
+    )
+    button_message = SimpleNamespace(chat=chat, delete=AsyncMock())
+    callback = SimpleNamespace(
+        data="pressed_button ABCD 11",
+        from_user=SimpleNamespace(id=11),
+        message=button_message,
+        answer=AsyncMock(),
+    )
+    pending = {
+        "chat_id": -100,
+        "user_id": 11,
+        "attempt_id": "attempt",
+        "answer": "ABCD",
+        "keyboard": [["answer"]],
+        "attempts": 0,
+    }
+    greetings.answers_dicc = [pending]
+    await greetings.update_captcha_cache([pending])
+
+    await greetings.callback_query_welcome_button(None, callback)
+    await greetings.kick_restricted_after_delay(
+        300, button_message, SimpleNamespace(id=11), "attempt"
+    )
+
+    chat.unban_member.assert_awaited_once_with(11)
+    chat.ban_member.assert_not_awaited()
+    assert await greetings.get_captcha_cache() == []
+
+
+async def test_timeout_only_bans_matching_pending_verification(greetings, monkeypatch):
+    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
+    first_chat = SimpleNamespace(
+        id=-100,
+        get_member=AsyncMock(
+            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
+        ),
+        ban_member=AsyncMock(),
+    )
+    second_chat = SimpleNamespace(id=-200, ban_member=AsyncMock())
+    message = SimpleNamespace(chat=first_chat, delete=AsyncMock())
+    await greetings.update_captcha_cache(
+        [
+            {"chat_id": -100, "user_id": 11, "attempt_id": "new"},
+            {"chat_id": -200, "user_id": 11, "attempt_id": "other"},
+        ]
+    )
+
+    await greetings.kick_restricted_after_delay(
+        300, message, SimpleNamespace(id=11), "old"
+    )
+    first_chat.ban_member.assert_not_awaited()
+    await greetings.kick_restricted_after_delay(
+        300, message, SimpleNamespace(id=11), "new"
+    )
+    first_chat.ban_member.assert_awaited_once()
+    second_chat.ban_member.assert_not_awaited()
+    assert await greetings.get_captcha_cache() == [
+        {"chat_id": -200, "user_id": 11, "attempt_id": "other"}
+    ]

wbb/modules/greetings.py

@@ -23,21 +23,22 @@ SOFTWARE.
 """
 
 import asyncio
-from re import findall
 from datetime import datetime, timedelta
 from random import shuffle
+from re import findall
+from uuid import uuid4
 
 from pyrogram import filters
 from pyrogram.enums import ChatMemberStatus
+from pyrogram.enums import ChatMemberStatus as CMS
 from pyrogram.errors.exceptions.bad_request_400 import (
     ChatAdminRequired,
     UserNotParticipant,
 )
-from pyrogram.enums import ChatMemberStatus as CMS
 from pyrogram.types import (
     Chat,
-    ChatPermissions,
     ChatMemberUpdated,
+    ChatPermissions,
     InlineKeyboardButton,
     InlineKeyboardMarkup,
     Message,
@@ -98,6 +99,7 @@ button=[查看群规, https://example.com]
 """
 
 answers_dicc = []
+_captcha_lock = asyncio.Lock()
 loop = asyncio.get_running_loop()
 
 
@@ -200,13 +202,13 @@ async def handle_new_member(member, chat):
     verification_data = {
         "chat_id": chat.id,
         "user_id": member.id,
+        "attempt_id": uuid4().hex,
         "answer": captcha_answer,
         "keyboard": keyboard,
         "attempts": 0,
     }
     keyboard = InlineKeyboardMarkup(keyboard)
     # Append user info, correct answer, and
-    answers_dicc.append(verification_data)
     # keyboard for later use with callback query
     button_message = await app.send_photo(
         chat_id=chat.id,
@@ -216,11 +218,15 @@ async def handle_new_member(member, chat):
     )
 
     # Save captcha answers etc in mongodb in case the bot gets crashed or restarted.
-    await update_captcha_cache(answers_dicc)
+    async with _captcha_lock:
+        answers_dicc = await get_captcha_cache()
+        answers_dicc.append(verification_data)
+        await update_captcha_cache(answers_dicc)
 
     asyncio.create_task(
         kick_restricted_after_delay(
-            WELCOME_DELAY_KICK_SEC, button_message, member
+            WELCOME_DELAY_KICK_SEC, button_message, member,
+            verification_data["attempt_id"],
         )
     )
     await asyncio.sleep(0.5)
@@ -349,45 +355,61 @@ async def callback_query_welcome_button(_, callback_query):
             reply_markup=keyboard,
         )
 
+    chat = button_message.chat
+    # Make the old timeout ineligible before changing Telegram permissions.
+    async with _captcha_lock:
+        answers_dicc = await get_captcha_cache()
+        pending = next(
+            (
+                item for item in answers_dicc
+                if item.get("user_id") == pending_user_id
+                and item.get("chat_id") == chat.id
+                and item.get("answer") == answer
+            ),
+            None,
+        )
+        if pending is None:
+            return await callback_query.answer("验证信息已失效，请退出群组后重新加入。")
+        await save_captcha_solved(chat.id, pending_user_id)
+        answers_dicc.remove(pending)
+        await update_captcha_cache(answers_dicc)
     await callback_query.answer("验证通过，欢迎加入。")
     await button_message.chat.unban_member(pending_user_id)
     await button_message.delete()
 
-    if len(answers_dicc) != 0:
-        for ii in answers_dicc:
-            if (
-                ii["user_id"] == pending_user_id
-                and ii["chat_id"] == button_message.chat.id
-            ):
-                answers_dicc.remove(ii)
-                await update_captcha_cache(answers_dicc)
-
-    chat = callback_query.message.chat
-
-    # Save this verification in db, so we don't have to
-    # send captcha to this user when he joins again.
-    await save_captcha_solved(chat.id, pending_user_id)
-
     return await send_welcome_message(chat, pending_user_id, True)
 
 
 async def kick_restricted_after_delay(
-    delay, button_message: Message, user: User
+    delay, button_message: Message, user: User, attempt_id: str
 ):
-    """If the new member is still restricted after the delay, delete
-    button message and join message and then kick him
-    """
+    """Only expire the still-pending verification that started this timer."""
     global answers_dicc
     await asyncio.sleep(delay)
     group_chat = button_message.chat
     user_id = user.id
-    await button_message.delete()
-    if len(answers_dicc) != 0:
-        for i in answers_dicc:
-            if i["user_id"] == user_id:
-                answers_dicc.remove(i)
-                await update_captcha_cache(answers_dicc)
-    await _ban_restricted_user_until_date(group_chat, user_id, duration=delay)
+    async with _captcha_lock:
+        answers_dicc = await get_captcha_cache()
+        pending = next(
+            (
+                item for item in answers_dicc
+                if item.get("chat_id") == group_chat.id
+                and item.get("user_id") == user_id
+                and item.get("attempt_id") == attempt_id
+            ),
+            None,
+        )
+        if pending is None:
+            return
+        answers_dicc.remove(pending)
+        await update_captcha_cache(answers_dicc)
+        if await has_solved_captcha_once(group_chat.id, user_id):
+            return
+        try:
+            await button_message.delete()
+        except Exception:
+            pass  # The message may already have been deleted by another handler.
+        await _ban_restricted_user_until_date(group_chat, user_id, duration=delay)
 
 
 async def _ban_restricted_user_until_date(

wbb/utils/dbadmin.py

@@ -361,7 +361,10 @@ async def get_managed_chat_settings(chat_id: int) -> dict[str, Any]:
     doc = await managed_chat_settingsdb.find_one(
         {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)}
     )
-    return doc or {
+    if doc is not None:
+        doc.setdefault("captcha_enabled", False)
+        return doc
+    return {
         "chat_id": int(chat_id),
         "auto_replies": [],
         "blacklist_words": [],

wbb/utils/dbfunctions.py

@@ -27,6 +27,7 @@ from string import ascii_lowercase
 from typing import Dict, List, Union
 
 from wbb import db
+from wbb.utils.dbadmin import get_managed_chat_settings, update_managed_chat_settings
 
 # SOME THINGS ARE FUCKED UP HERE, LIKE TOGGLEABLES HAVE THEIR OWN COLLECTION
 # (SHOULD FIX IT WITH SOMETHING LIKE TOGGLEDB), MOST OF THE CODE IS BAD AF
@@ -446,24 +447,20 @@ async def save_couple(chat_id: int, date: str, couple: dict):
 
 
 async def is_captcha_on(chat_id: int) -> bool:
-    chat = await captchadb.find_one({"chat_id": chat_id})
-    if not chat:
-        return True
-    return False
+    settings = await get_managed_chat_settings(chat_id)
+    return bool(settings.get("captcha_enabled", False))
 
 
 async def captcha_on(chat_id: int):
-    is_captcha = await is_captcha_on(chat_id)
-    if is_captcha:
-        return
-    return await captchadb.delete_one({"chat_id": chat_id})
+    await update_managed_chat_settings(chat_id, {"captcha_enabled": True})
+    return await captchadb.delete_many({"chat_id": chat_id})
 
 
 async def captcha_off(chat_id: int):
-    is_captcha = await is_captcha_on(chat_id)
-    if not is_captcha:
-        return
-    return await captchadb.insert_one({"chat_id": chat_id})
+    await update_managed_chat_settings(chat_id, {"captcha_enabled": False})
+    return await captchadb.update_one(
+        {"chat_id": chat_id}, {"$set": {"chat_id": chat_id}}, upsert=True
+    )
 
 
 async def has_solved_captcha_once(chat_id: int, user_id: int):
@@ -475,8 +472,8 @@ async def has_solved_captcha_once(chat_id: int, user_id: int):
 
 async def save_captcha_solved(chat_id: int, user_id: int):
     return await solved_captcha_db.update_one(
-        {"chat_id": chat_id},
-        {"$set": {"user_id": user_id}},
+        {"chat_id": chat_id, "user_id": user_id},
+        {"$set": {"chat_id": chat_id, "user_id": user_id}},
         upsert=True,
     )