From 9f33150781c41a1de4fbdb7566fe0d120994162f Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Sun, 2 Aug 2026 17:12:40 +0700 Subject: [PATCH 1/7] feat: block non-whitelisted guest bot messages with progressive restriction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Telegram Bot API 10.0 (May 2026) introduced Guest Mode, allowing any user to @mention a bot in any group chat without the bot being a member. The bot posts a reply directly in the chat. This feature adds moderation for those guest bot messages. New handler (src/bot/handlers/guest_bot.py): - GuestBotFilter: custom MessageFilter matching only guest bot messages (guest_bot_caller_user or guest_bot_caller_chat set) - is_guest_bot_whitelisted: case-insensitive username matching against per-group whitelist (strips @, lowercases) - handle_guest_bot_message: deletes non-whitelisted guest bot messages and progressively restricts the invoking user (not the bot) using the existing UserWarning state machine with a new warning_kind=guest_bot discriminator — 1st violation: warning, 2nd to (N-1): silent increment, Nth: restrict + notification. Admins/trusted users are exempt. Channel callers are delete-only (no human to restrict). Already-restricted callers skip new warning cycles. Plugin wiring: - New plugin guest_bot_block registered at group=1, before inline_keyboard_spam, with GuestBotFilter for precise matching - Gated by guard_plugin(guest_bot_block) for per-group toggle control - Blocking handler (no block=False) so ApplicationHandlerStop works Config: - guest_bot_whitelist field in Settings (NoDecode annotation for comma-separated env parsing) and GroupConfig (normalization validator) - GUEST_BOT_WHITELIST env var and groups.json support - Plugin toggle via plugins map in groups.json DB schema: - Added warning_kind column to UserWarning (default profile) - SQLite migration for existing databases - All DB service methods accept and filter by warning_kind - Scheduler only processes profile warnings (not guest bot) - DM/verify unrestriction only lifts profile restrictions Tests: 1048 passed, ruff clean, mypy clean --- .env.example | 7 + groups.json.example | 2 + src/bot/config.py | 16 ++- src/bot/constants.py | 13 ++ src/bot/database/models.py | 4 + src/bot/database/service.py | 79 ++++++++--- src/bot/group_config.py | 14 ++ src/bot/handlers/guest_bot.py | 112 ++++++++++++++++ src/bot/plugins/builtin/spam.py | 14 ++ src/bot/plugins/definitions.py | 3 +- src/bot/plugins/manager.py | 1 + tests/test_config.py | 20 +++ tests/test_group_config.py | 28 ++++ tests/test_guest_bot.py | 190 +++++++++++++++++++++++++++ tests/test_main_plugins_bootstrap.py | 27 +++- tests/test_plugin_manager.py | 1 + 16 files changed, 512 insertions(+), 19 deletions(-) create mode 100644 src/bot/handlers/guest_bot.py create mode 100644 tests/test_guest_bot.py diff --git a/.env.example b/.env.example index 5fcc6a2..62a1f2a 100644 --- a/.env.example +++ b/.env.example @@ -86,6 +86,13 @@ BIO_BAIT_MONITOR_ONLY=false # Example: PLUGINS_DEFAULT={"captcha":true,"dm":false} # PLUGINS_DEFAULT={"captcha":true,"dm":false} +# Guest Bot Whitelist (Telegram Guest Mode - Bot API 10.0) +# Comma-separated list of bot usernames allowed to post guest messages +# Messages from non-whitelisted guest bots are deleted and the invoking user is warned +# Usernames are case-insensitive, @ prefix is optional +# Example: GUEST_BOT_WHITELIST=@somebot,anotherbot +GUEST_BOT_WHITELIST= + # Logfire Configuration (optional - for production logging) # Get your token from https://logfire.pydantic.dev LOGFIRE_TOKEN=your_logfire_token_here diff --git a/groups.json.example b/groups.json.example index 1cefc4c..9808cb0 100644 --- a/groups.json.example +++ b/groups.json.example @@ -20,6 +20,7 @@ "bio_bait_monitor_only": false, "bio_bait_alert_chat_id": null, "moderation_topic_id": null, + "guest_bot_whitelist": [], "plugins": { "captcha": false, "dm": true, @@ -47,6 +48,7 @@ "bio_bait_monitor_only": false, "bio_bait_alert_chat_id": null, "moderation_topic_id": null, + "guest_bot_whitelist": ["somebot"], "plugins": { "contact_spam": false, "duplicate_spam": false, diff --git a/src/bot/config.py b/src/bot/config.py index adc0e3b..ac043e4 100644 --- a/src/bot/config.py +++ b/src/bot/config.py @@ -11,9 +11,10 @@ import os from functools import lru_cache from pathlib import Path +from typing import Annotated from pydantic import field_validator -from pydantic_settings import BaseSettings, SettingsConfigDict +from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict logger = logging.getLogger(__name__) @@ -84,6 +85,7 @@ class Settings(BaseSettings): bio_bait_monitor_only: bool = False bio_bait_alert_chat_id: int | None = None moderation_topic_id: int | None = None + guest_bot_whitelist: Annotated[list[str], NoDecode] = [] groups_config_path: str = "groups.json" logfire_token: str | None = None logfire_service_name: str = "pythonid-bot" @@ -120,6 +122,18 @@ def parse_and_validate_plugins_default(cls, v: object) -> dict[str, bool]: from bot.plugins.config import validate_plugin_map return validate_plugin_map(parsed) + @field_validator("guest_bot_whitelist", mode="before") + @classmethod + def parse_guest_bot_whitelist(cls, v: object) -> list[str]: + """Parse GUEST_BOT_WHITELIST env var as comma-separated usernames.""" + if isinstance(v, list): + return [str(entry).strip().removeprefix("@").lower() for entry in v if str(entry).strip()] + if isinstance(v, str): + if not v.strip(): + return [] + return [entry.strip().removeprefix("@").lower() for entry in v.split(",") if entry.strip()] + return [] + def model_post_init(self, __context): """Validate and log non-sensitive configuration values after initialization.""" if self.group_id >= 0: diff --git a/src/bot/constants.py b/src/bot/constants.py index a677ec9..4f2d885 100644 --- a/src/bot/constants.py +++ b/src/bot/constants.py @@ -344,6 +344,19 @@ def format_hours_display(hours: int) -> str: "📌 [Peraturan Grup]({rules_link})" ) +GUEST_BOT_WARNING = ( + "⚠️ {user_mention}, bot tamu tidak diizinkan di grup ini. " + "Pelanggaran berikutnya dapat menyebabkan pembatasan setelah " + "{warning_threshold} pesan.\n\n" + "Silakan baca [peraturan grup]({rules_link})." +) + +GUEST_BOT_RESTRICTION = ( + "🔇 {user_mention} dibatasi setelah memanggil bot tamu sebanyak " + "{message_count} kali.\n\n" + "Silakan baca [peraturan grup]({rules_link})." +) + # Duplicate message spam notification DUPLICATE_SPAM_RESTRICTION = ( "🚫 *Spam Pesan Duplikat*\n\n" diff --git a/src/bot/database/models.py b/src/bot/database/models.py index d085c94..735fa0e 100644 --- a/src/bot/database/models.py +++ b/src/bot/database/models.py @@ -32,6 +32,9 @@ class UserWarning(SQLModel, table=True): restricted_by_bot: True if restriction was applied by this bot (vs manually by an admin). Only bot-created restrictions can be lifted via DM. + warning_kind: Discriminator for the warning source + (``"profile"`` or ``"guest_bot"``). Prevents cross-source + state interference. """ __tablename__ = "user_warnings" @@ -44,6 +47,7 @@ class UserWarning(SQLModel, table=True): last_message_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) is_restricted: bool = Field(default=False) restricted_by_bot: bool = Field(default=False) + warning_kind: str = Field(default="profile", index=True) class PhotoVerificationWhitelist(SQLModel, table=True): diff --git a/src/bot/database/service.py b/src/bot/database/service.py index 6a8187b..e6264ce 100644 --- a/src/bot/database/service.py +++ b/src/bot/database/service.py @@ -58,6 +58,7 @@ def __init__(self, database_path: str): # Migrate existing tables: add new columns if missing self._migrate_trusted_users() + self._migrate_user_warnings() def _migrate_trusted_users(self) -> None: """Add new columns to trusted_users if missing.""" @@ -80,7 +81,24 @@ def _migrate_trusted_users(self) -> None: logger.info(f"Migrated trusted_users: added {col} column") conn.commit() - def get_or_create_user_warning(self, user_id: int, group_id: int) -> UserWarning: + def _migrate_user_warnings(self) -> None: + """Add warning_kind column to user_warnings if missing.""" + with self._engine.connect() as conn: + columns = { + row[1] for row in conn.exec_driver_sql( + "PRAGMA table_info(user_warnings)" + ).fetchall() + } + if "warning_kind" not in columns: + conn.exec_driver_sql( + "ALTER TABLE user_warnings ADD COLUMN warning_kind TEXT DEFAULT 'profile'" + ) + logger.info("Migrated user_warnings: added warning_kind column") + conn.commit() + + def get_or_create_user_warning( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> UserWarning: """ Get existing warning record or create a new one. @@ -90,6 +108,7 @@ def get_or_create_user_warning(self, user_id: int, group_id: int) -> UserWarning Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: UserWarning: Active warning record for the user. @@ -99,13 +118,14 @@ def get_or_create_user_warning(self, user_id: int, group_id: int) -> UserWarning statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ~UserWarning.is_restricted, ) record = session.exec(statement).first() if record: logger.info( - f"Returning existing warning for user_id={user_id}, group_id={group_id}" + f"Returning existing warning for user_id={user_id}, group_id={group_id}, kind={warning_kind}" ) return record @@ -116,16 +136,19 @@ def get_or_create_user_warning(self, user_id: int, group_id: int) -> UserWarning message_count=1, first_warned_at=datetime.now(UTC), last_message_at=datetime.now(UTC), + warning_kind=warning_kind, ) session.add(new_record) session.commit() session.refresh(new_record) logger.info( - f"Created new warning for user_id={user_id}, group_id={group_id}" + f"Created new warning for user_id={user_id}, group_id={group_id}, kind={warning_kind}" ) return new_record - def increment_message_count(self, user_id: int, group_id: int) -> UserWarning: + def increment_message_count( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> UserWarning: """ Increment message count for an existing warning record. @@ -135,6 +158,7 @@ def increment_message_count(self, user_id: int, group_id: int) -> UserWarning: Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: UserWarning: Updated warning record. @@ -146,6 +170,7 @@ def increment_message_count(self, user_id: int, group_id: int) -> UserWarning: statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ~UserWarning.is_restricted, ) record = session.exec(statement).first() @@ -157,15 +182,17 @@ def increment_message_count(self, user_id: int, group_id: int) -> UserWarning: session.commit() session.refresh(record) logger.info( - f"Incremented message count for user_id={user_id}, group_id={group_id}, new_count={record.message_count}" + f"Incremented message count for user_id={user_id}, group_id={group_id}, kind={warning_kind}, new_count={record.message_count}" ) return record raise ValueError( - f"No warning record found for user {user_id} in group {group_id}" + f"No warning record found for user {user_id} in group {group_id} (kind={warning_kind})" ) - def mark_user_restricted(self, user_id: int, group_id: int) -> UserWarning: + def mark_user_restricted( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> UserWarning: """ Mark user as restricted after reaching threshold. @@ -175,6 +202,7 @@ def mark_user_restricted(self, user_id: int, group_id: int) -> UserWarning: Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: UserWarning: Updated warning record. @@ -186,6 +214,7 @@ def mark_user_restricted(self, user_id: int, group_id: int) -> UserWarning: statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ~UserWarning.is_restricted, ) record = session.exec(statement).first() @@ -198,15 +227,17 @@ def mark_user_restricted(self, user_id: int, group_id: int) -> UserWarning: session.commit() session.refresh(record) logger.info( - f"Marked user as restricted: user_id={user_id}, group_id={group_id}" + f"Marked user as restricted: user_id={user_id}, group_id={group_id}, kind={warning_kind}" ) return record raise ValueError( - f"No warning record found for user {user_id} in group {group_id}" + f"No warning record found for user {user_id} in group {group_id} (kind={warning_kind})" ) - def is_user_restricted_by_bot(self, user_id: int, group_id: int) -> bool: + def is_user_restricted_by_bot( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> bool: """ Check if user was restricted by this bot. @@ -217,6 +248,7 @@ def is_user_restricted_by_bot(self, user_id: int, group_id: int) -> bool: Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: bool: True if user was restricted by this bot. @@ -225,13 +257,16 @@ def is_user_restricted_by_bot(self, user_id: int, group_id: int) -> bool: statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, UserWarning.is_restricted, UserWarning.restricted_by_bot, ) record = session.exec(statement).first() return record is not None - def mark_user_unrestricted(self, user_id: int, group_id: int) -> None: + def mark_user_unrestricted( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> None: """ Clear bot restriction flag after user is unrestricted via DM. @@ -241,11 +276,13 @@ def mark_user_unrestricted(self, user_id: int, group_id: int) -> None: Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. """ with Session(self._engine) as session: statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, UserWarning.is_restricted, UserWarning.restricted_by_bot, ) @@ -256,10 +293,12 @@ def mark_user_unrestricted(self, user_id: int, group_id: int) -> None: session.add(record) session.commit() logger.info( - f"Cleared restriction flag: user_id={user_id}, group_id={group_id}" + f"Cleared restriction flag: user_id={user_id}, group_id={group_id}, kind={warning_kind}" ) - def delete_user_warnings(self, user_id: int, group_id: int) -> int: + def delete_user_warnings( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> int: """ Delete all warning records for a user in a specific group. @@ -269,6 +308,7 @@ def delete_user_warnings(self, user_id: int, group_id: int) -> int: Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: int: Number of warning records deleted. @@ -277,16 +317,19 @@ def delete_user_warnings(self, user_id: int, group_id: int) -> int: delete_statement = delete(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ) result = session.exec(delete_statement) session.commit() count = result.rowcount logger.info( - f"Deleted warnings: user_id={user_id}, group_id={group_id}, count={count}" + f"Deleted warnings: user_id={user_id}, group_id={group_id}, kind={warning_kind}, count={count}" ) return count - def get_active_user_warning(self, user_id: int, group_id: int) -> UserWarning | None: + def get_active_user_warning( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> UserWarning | None: """ Get an existing active (non-restricted) warning record without creating one. @@ -297,6 +340,7 @@ def get_active_user_warning(self, user_id: int, group_id: int) -> UserWarning | Args: user_id: Telegram user ID. group_id: Telegram group ID. + warning_kind: Discriminator for the warning source. Returns: UserWarning | None: Active warning record, or None if none exists. @@ -305,6 +349,7 @@ def get_active_user_warning(self, user_id: int, group_id: int) -> UserWarning | statement = select(UserWarning).where( UserWarning.user_id == user_id, UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ~UserWarning.is_restricted, ) return session.exec(statement).first() @@ -540,7 +585,7 @@ def get_trusted_users(self) -> list[TrustedUser]: return list(session.exec(statement).all()) def get_warnings_past_time_threshold_for_group( - self, group_id: int, threshold: timedelta + self, group_id: int, threshold: timedelta, warning_kind: str = "profile" ) -> list[UserWarning]: """ Find active warnings for a specific group that exceeded the time threshold. @@ -548,6 +593,7 @@ def get_warnings_past_time_threshold_for_group( Args: group_id: Telegram group ID to filter by. threshold: Time duration since first warning to trigger restriction. + warning_kind: Discriminator for the warning source. Returns: list[UserWarning]: Warning records that should be auto-restricted. @@ -556,6 +602,7 @@ def get_warnings_past_time_threshold_for_group( cutoff_time = datetime.now(UTC) - threshold statement = select(UserWarning).where( UserWarning.group_id == group_id, + UserWarning.warning_kind == warning_kind, ~UserWarning.is_restricted, UserWarning.first_warned_at <= cutoff_time, ) diff --git a/src/bot/group_config.py b/src/bot/group_config.py index 6edb22d..96532f3 100644 --- a/src/bot/group_config.py +++ b/src/bot/group_config.py @@ -43,6 +43,7 @@ class GroupConfig(BaseModel): bio_bait_enabled: bool = True bio_bait_monitor_only: bool = False bio_bait_alert_chat_id: int | None = None + guest_bot_whitelist: list[str] = [] moderation_topic_id: int | None = None plugins: dict[str, bool] | None = None @@ -92,6 +93,19 @@ def validate_plugins(cls, v: object) -> dict[str, bool] | None: raise ValueError("plugins must be a dict or None") return validate_plugin_map(v) + @field_validator("guest_bot_whitelist", mode="before") + @classmethod + def normalize_guest_bot_whitelist(cls, v: object) -> list[str]: + if v is None: + return [] + if isinstance(v, list): + return [str(entry).strip().removeprefix("@").lower() for entry in v if str(entry).strip()] + if isinstance(v, str): + if not v.strip(): + return [] + return [entry.strip().removeprefix("@").lower() for entry in v.split(",") if entry.strip()] + return [] + @property def probation_timedelta(self) -> timedelta: return timedelta(hours=self.new_user_probation_hours) diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py new file mode 100644 index 0000000..30b25f8 --- /dev/null +++ b/src/bot/handlers/guest_bot.py @@ -0,0 +1,112 @@ +"""Guest bot message moderation for Telegram Guest Mode.""" + +import logging + +from telegram import Message, Update, User +from telegram.ext import ApplicationHandlerStop, ContextTypes, filters + +from bot.constants import ( + GUEST_BOT_RESTRICTION, + GUEST_BOT_WARNING, + RESTRICTED_PERMISSIONS, +) +from bot.database.service import get_database +from bot.group_config import get_group_config_for_update +from bot.services.telegram_utils import get_user_mention, is_user_admin_or_trusted + +logger = logging.getLogger(__name__) + + +class GuestBotFilter(filters.MessageFilter): + """Message filter matching only Telegram Guest Mode messages.""" + + def filter(self, message: Message) -> bool: + return ( + message.guest_bot_caller_user is not None + or message.guest_bot_caller_chat is not None + ) + + +def is_guest_bot_message(message: Message) -> bool: + """Check if a message was posted by a guest bot.""" + return message.guest_bot_caller_user is not None or message.guest_bot_caller_chat is not None + + +def is_guest_bot_whitelisted(message: Message, whitelist: list[str]) -> bool: + """Check if the guest bot that posted this message is whitelisted.""" + username = message.from_user.username if message.from_user else None + if not username: + return False + normalized_whitelist = {entry.strip().removeprefix("@").lower() for entry in whitelist} + return username.lower() in normalized_whitelist + + +async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Delete unapproved guest bot messages and progressively restrict their caller.""" + message = update.message + if message is None: + return + + group_config = get_group_config_for_update(update) + if group_config is None or not is_guest_bot_message(message): + return + if is_guest_bot_whitelisted(message, group_config.guest_bot_whitelist): + return + + caller = message.guest_bot_caller_user or message.guest_bot_caller_chat + try: + await message.delete() + except Exception: + logger.error("Failed to delete guest bot message", exc_info=True) + + if caller is None or not isinstance(caller, User): + raise ApplicationHandlerStop + if is_user_admin_or_trusted(context, group_config.group_id, caller.id): + raise ApplicationHandlerStop + + db = get_database() + if db.is_user_restricted_by_bot(caller.id, group_config.group_id, warning_kind="guest_bot"): + raise ApplicationHandlerStop + + record = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") + user_mention = get_user_mention(caller) + + if record.message_count == 1: + try: + await context.bot.send_message( + chat_id=group_config.group_id, + message_thread_id=group_config.warning_topic_id, + text=GUEST_BOT_WARNING.format( + user_mention=user_mention, + warning_threshold=group_config.warning_threshold, + rules_link=group_config.rules_link, + ), + parse_mode="Markdown", + ) + except Exception: + logger.error("Failed to send guest bot warning for user %s", caller.id, exc_info=True) + + if record.message_count >= group_config.warning_threshold: + try: + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=caller.id, + permissions=RESTRICTED_PERMISSIONS, + ) + db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + await context.bot.send_message( + chat_id=group_config.group_id, + message_thread_id=group_config.warning_topic_id, + text=GUEST_BOT_RESTRICTION.format( + user_mention=user_mention, + message_count=record.message_count, + rules_link=group_config.rules_link, + ), + parse_mode="Markdown", + ) + except Exception: + logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) + else: + db.increment_message_count(caller.id, group_config.group_id, warning_kind="guest_bot") + + raise ApplicationHandlerStop diff --git a/src/bot/plugins/builtin/spam.py b/src/bot/plugins/builtin/spam.py index c445c64..d031d51 100644 --- a/src/bot/plugins/builtin/spam.py +++ b/src/bot/plugins/builtin/spam.py @@ -19,6 +19,7 @@ from bot.handlers.anti_spam import handle_contact_spam, handle_inline_keyboard_spam, handle_new_user_spam from bot.handlers.bio_bait import BIO_BAIT_FILTER, handle_bio_bait_spam from bot.handlers.duplicate_spam import handle_duplicate_spam +from bot.handlers.guest_bot import GuestBotFilter, handle_guest_bot_message from bot.plugins.config import guard_plugin if TYPE_CHECKING: @@ -48,6 +49,19 @@ def register_inline_keyboard_spam(application: Application) -> list[BaseHandler] ) return _register_spam(application, handler, 1, "inline_keyboard_spam_handler") +def register_guest_bot_block(application: Application) -> list[BaseHandler]: # type: ignore[type-arg] + """Register guest bot block handler (group=1). + + Callback wrapped with ``guard_plugin(\"guest_bot_block\")``. Runs at + group=1 alongside ``inline_keyboard_spam`` to intercept Telegram + Guest Mode messages before downstream spam handlers. + """ + handler: BaseHandler = MessageHandler( + GuestBotFilter(), + guard_plugin("guest_bot_block")(handle_guest_bot_message), + ) + return _register_spam(application, handler, 1, "guest_bot_block_handler") + def register_bio_bait_spam(application: Application) -> list[BaseHandler]: # type: ignore[type-arg] """Register bio bait spam handler (group=4). diff --git a/src/bot/plugins/definitions.py b/src/bot/plugins/definitions.py index cc9617f..50fdbb1 100644 --- a/src/bot/plugins/definitions.py +++ b/src/bot/plugins/definitions.py @@ -35,6 +35,7 @@ {"name": "captcha", "handler_group": 0, "description": "Captcha verification for new members"}, {"name": "dm", "handler_group": 0, "description": "Direct message unrestriction flow"}, {"name": "status", "handler_group": 0, "description": "Admin /status command"}, + {"name": "guest_bot_block", "handler_group": 1, "description": "Block non-whitelisted guest bot messages"}, {"name": "inline_keyboard_spam", "handler_group": 1, "description": "Block inline keyboard URL spam"}, {"name": "contact_spam", "handler_group": 2, "description": "Block contact card sharing"}, {"name": "new_user_spam", "handler_group": 3, "description": "Probation enforcement for new users"}, @@ -71,4 +72,4 @@ def get_plugin_definitions() -> PluginManifest: Returns: List of plugin descriptor dicts with keys: name, handler_group, description. """ - return copy.deepcopy(_PLUGIN_DEFINITIONS) \ No newline at end of file + return copy.deepcopy(_PLUGIN_DEFINITIONS) diff --git a/src/bot/plugins/manager.py b/src/bot/plugins/manager.py index 9286d86..8de614c 100644 --- a/src/bot/plugins/manager.py +++ b/src/bot/plugins/manager.py @@ -73,6 +73,7 @@ "status": status_mod.register_status, # spam "inline_keyboard_spam": spam_mod.register_inline_keyboard_spam, + "guest_bot_block": spam_mod.register_guest_bot_block, "bio_bait_spam": spam_mod.register_bio_bait_spam, "contact_spam": spam_mod.register_contact_spam, "new_user_spam": spam_mod.register_new_user_spam, diff --git a/tests/test_config.py b/tests/test_config.py index fcddbf7..6f30b96 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -142,6 +142,26 @@ def test_bio_bait_monitor_from_env(self, monkeypatch): assert settings.bio_bait_monitor_only is True assert settings.bio_bait_alert_chat_id == 57747812 + def test_guest_bot_whitelist_from_env(self, monkeypatch): + monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "test_token") + monkeypatch.setenv("GROUP_ID", "-100999") + monkeypatch.setenv("WARNING_TOPIC_ID", "1") + monkeypatch.setenv("GUEST_BOT_WHITELIST", "@somebot,anotherbot") + + settings = Settings(_env_file=None) + + assert settings.guest_bot_whitelist == ["somebot", "anotherbot"] + + def test_guest_bot_whitelist_empty_env(self, monkeypatch): + monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "test_token") + monkeypatch.setenv("GROUP_ID", "-100999") + monkeypatch.setenv("WARNING_TOPIC_ID", "1") + monkeypatch.setenv("GUEST_BOT_WHITELIST", "") + + settings = Settings(_env_file=None) + + assert settings.guest_bot_whitelist == [] + class TestPluginsDefault: def test_default_empty_dict(self, monkeypatch): """Test plugins_default defaults to empty dict when not set.""" diff --git a/tests/test_group_config.py b/tests/test_group_config.py index ed251c5..ed70549 100644 --- a/tests/test_group_config.py +++ b/tests/test_group_config.py @@ -138,6 +138,34 @@ def test_plugins_rejects_non_dict(self): with pytest.raises(ValidationError, match="plugins must be a dict"): GroupConfig(group_id=-1, warning_topic_id=42, plugins=[1, 2, 3]) + def test_guest_bot_whitelist_default_empty(self): + gc = GroupConfig(group_id=-1, warning_topic_id=42) + assert gc.guest_bot_whitelist == [] + + def test_guest_bot_whitelist_normalizes_entries(self): + gc = GroupConfig( + group_id=-1, + warning_topic_id=42, + guest_bot_whitelist=["@SomeBot", "AnotherBot", " @lowerbot "], + ) + assert gc.guest_bot_whitelist == ["somebot", "anotherbot", "lowerbot"] + + def test_guest_bot_whitelist_string_env_format(self): + gc = GroupConfig( + group_id=-1, + warning_topic_id=42, + guest_bot_whitelist="@Bot1,Bot2, @Bot3", + ) + assert gc.guest_bot_whitelist == ["bot1", "bot2", "bot3"] + + def test_guest_bot_whitelist_empty_string(self): + gc = GroupConfig(group_id=-1, warning_topic_id=42, guest_bot_whitelist="") + assert gc.guest_bot_whitelist == [] + + def test_guest_bot_whitelist_none(self): + gc = GroupConfig(group_id=-1, warning_topic_id=42, guest_bot_whitelist=None) + assert gc.guest_bot_whitelist == [] + class TestGroupRegistry: def test_register_and_get(self): registry = GroupRegistry() diff --git a/tests/test_guest_bot.py b/tests/test_guest_bot.py new file mode 100644 index 0000000..b72e9ef --- /dev/null +++ b/tests/test_guest_bot.py @@ -0,0 +1,190 @@ +"""Tests for guest bot message moderation.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from telegram import Chat, Message, User +from telegram.ext import ApplicationHandlerStop + +from bot.group_config import GroupConfig +from bot.handlers.guest_bot import ( + handle_guest_bot_message, + is_guest_bot_message, + is_guest_bot_whitelisted, +) + + +@pytest.fixture +def mock_user(): + return User(id=123, first_name="Test", is_bot=False, username="testuser") + + +@pytest.fixture +def mock_group_config(): + return GroupConfig( + group_id=-1001234567890, + warning_topic_id=123, + warning_threshold=3, + guest_bot_whitelist=["allowedbot"], + ) + + +@pytest.fixture +def mock_update(mock_user): + message = MagicMock(spec=Message) + message.from_user = User(id=999, first_name="Guest Bot", is_bot=True, username="guestbot") + message.guest_bot_caller_user = mock_user + message.guest_bot_caller_chat = None + message.delete = AsyncMock() + update = MagicMock() + update.message = message + return update + + +@pytest.fixture +def mock_context(): + context = MagicMock() + context.bot_data = {"group_admin_ids": {}, "trusted_user_ids": []} + context.bot = AsyncMock() + return context + + +class TestIsGuestBotMessage: + def test_user_caller(self): + message = MagicMock(spec=Message) + message.guest_bot_caller_user = MagicMock() + message.guest_bot_caller_chat = None + assert is_guest_bot_message(message) is True + + def test_chat_caller(self): + message = MagicMock(spec=Message) + message.guest_bot_caller_user = None + message.guest_bot_caller_chat = MagicMock() + assert is_guest_bot_message(message) is True + + def test_regular_message(self): + message = MagicMock(spec=Message) + message.guest_bot_caller_user = None + message.guest_bot_caller_chat = None + assert is_guest_bot_message(message) is False + + +class TestIsGuestBotWhitelisted: + @pytest.mark.parametrize( + ("username", "whitelist", "expected"), + [ + ("allowedbot", ["allowedbot"], True), + ("otherbot", ["allowedbot"], False), + ("AllowedBot", ["allowedbot"], True), + (None, ["allowedbot"], False), + ], + ) + def test_whitelist(self, username, whitelist, expected): + message = MagicMock(spec=Message) + message.from_user = User(id=999, first_name="Bot", is_bot=True, username=username) + assert is_guest_bot_whitelisted(message, whitelist) is expected + + +class TestHandleGuestBotMessage: + async def test_non_guest_message_does_nothing(self, mock_update, mock_context, mock_group_config): + mock_update.message.guest_bot_caller_user = None + with patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_not_awaited() + + async def test_whitelisted_message_does_nothing(self, mock_update, mock_context, mock_group_config): + mock_update.message.from_user = User(id=999, first_name="Bot", is_bot=True, username="allowedbot") + with patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_not_awaited() + + async def test_unmonitored_group_returns(self, mock_update, mock_context): + with patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=None): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_not_awaited() + + async def test_admin_is_deleted_but_not_restricted(self, mock_update, mock_context, mock_group_config): + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=True), + patch("bot.handlers.guest_bot.get_database") as get_db, + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_awaited_once() + get_db.assert_not_called() + + @pytest.mark.parametrize(("count", "sends", "increments"), [(1, 1, True), (2, 0, True)]) + async def test_pre_threshold_violation( + self, mock_update, mock_context, mock_group_config, count, sends, increments + ): + db = MagicMock() + db.is_user_restricted_by_bot.return_value = False + db.get_or_create_user_warning.return_value.message_count = count + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=False), + patch("bot.handlers.guest_bot.get_database", return_value=db), + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + assert mock_context.bot.send_message.await_count == sends + assert db.increment_message_count.called is increments + + async def test_threshold_restricts_and_notifies(self, mock_update, mock_context, mock_group_config): + db = MagicMock() + db.is_user_restricted_by_bot.return_value = False + db.get_or_create_user_warning.return_value.message_count = 3 + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=False), + patch("bot.handlers.guest_bot.get_database", return_value=db), + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + mock_context.bot.restrict_chat_member.assert_awaited_once() + mock_context.bot.send_message.assert_awaited_once() + db.mark_user_restricted.assert_called_once_with( + 123, mock_group_config.group_id, warning_kind="guest_bot" + ) + + async def test_chat_caller_is_deleted_only(self, mock_update, mock_context, mock_group_config): + mock_update.message.guest_bot_caller_user = None + mock_update.message.guest_bot_caller_chat = Chat(id=-1009, type="channel", title="Channel") + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.get_database") as get_db, + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_awaited_once() + get_db.assert_not_called() + + async def test_delete_failure_continues(self, mock_update, mock_context, mock_group_config): + mock_update.message.delete.side_effect = RuntimeError("delete failed") + db = MagicMock() + db.is_user_restricted_by_bot.return_value = False + db.get_or_create_user_warning.return_value.message_count = 2 + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=False), + patch("bot.handlers.guest_bot.get_database", return_value=db), + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + db.increment_message_count.assert_called_once() + + async def test_already_restricted_skips_warning(self, mock_update, mock_context, mock_group_config): + db = MagicMock() + db.is_user_restricted_by_bot.return_value = True + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=False), + patch("bot.handlers.guest_bot.get_database", return_value=db), + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + mock_update.message.delete.assert_awaited_once() + db.get_or_create_user_warning.assert_not_called() + db.increment_message_count.assert_not_called() + db.mark_user_restricted.assert_not_called() diff --git a/tests/test_main_plugins_bootstrap.py b/tests/test_main_plugins_bootstrap.py index bd85cfb..2ffddbe 100644 --- a/tests/test_main_plugins_bootstrap.py +++ b/tests/test_main_plugins_bootstrap.py @@ -297,6 +297,12 @@ def test_spam_has_inline_keyboard_spam_registrar(self): assert hasattr(spam, "register_inline_keyboard_spam") assert callable(spam.register_inline_keyboard_spam) + def test_spam_has_guest_bot_block_registrar(self): + """bot.plugins.builtin.spam has register_guest_bot_block function.""" + from bot.plugins.builtin import spam + assert hasattr(spam, "register_guest_bot_block") + assert callable(spam.register_guest_bot_block) + def test_spam_has_bio_bait_spam_registrar(self): """bot.plugins.builtin.spam has register_bio_bait_spam function.""" from bot.plugins.builtin import spam @@ -441,4 +447,23 @@ def test_inline_keyboard_spam_registrar_adds_handler(self): assert len(call_args) == 1 assert call_kwargs["group"] == 1 from telegram.ext import MessageHandler - assert isinstance(call_args[0], MessageHandler) \ No newline at end of file + assert isinstance(call_args[0], MessageHandler) + + def test_guest_bot_block_registrar_adds_handler(self): + """register_guest_bot_block adds a guest-only handler to group=1.""" + from telegram.ext import MessageHandler + + from bot.handlers.guest_bot import GuestBotFilter + from bot.plugins.builtin.spam import register_guest_bot_block + + app = MagicMock() + app.bot_data = {} + app.add_handler = MagicMock() + handlers = register_guest_bot_block(app) + assert len(handlers) >= 1 + assert app.add_handler.call_count == 1 + call_args, call_kwargs = app.add_handler.call_args + assert len(call_args) == 1 + assert call_kwargs["group"] == 1 + assert isinstance(call_args[0], MessageHandler) + assert isinstance(call_args[0].filters, GuestBotFilter) diff --git a/tests/test_plugin_manager.py b/tests/test_plugin_manager.py index 145c0bc..53a5209 100644 --- a/tests/test_plugin_manager.py +++ b/tests/test_plugin_manager.py @@ -144,6 +144,7 @@ def _expected_order() -> tuple[str, ...]: "captcha", "dm", "status", + "guest_bot_block", "inline_keyboard_spam", "contact_spam", "new_user_spam", From 264b335ee1f12333023a2bb846119f2f346073b5 Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Sun, 2 Aug 2026 17:18:43 +0700 Subject: [PATCH 2/7] docs: document guest bot moderation feature Add Guest Mode blocking to AGENTS.md and README.md: architecture section (GuestBotFilter, handler group, progressive enforcement, separate warning_kind, whitelist config, plugin toggle), Where to Look / Code Map / Structure tree entries, README feature bullet and dedicated section, GUEST_BOT_WHITELIST in config table, updated test stats, database conventions for warning_kind discriminator and migration. --- AGENTS.md | 31 ++++++++++++++++++++++++++----- README.md | 31 ++++++++++++++++++++++++++++--- 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e048855..459f34f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -70,7 +70,8 @@ PythonID/ │ │ ├── trust.py # /trust, /untrust, /trusted admin commands │ │ ├── warn.py # Admin /warn command (reply or user ID) │ │ ├── duplicate_spam.py # Duplicate message detection -│ │ └── bio_bait.py # Bio-bait spam (bait phrases + suspicious profile bio links) +│ │ ├── bio_bait.py # Bio-bait spam (bait phrases + suspicious profile bio links) +│ │ └── guest_bot.py # Guest Mode moderation (delete + progressive restriction) │ ├── services/ │ │ ├── user_checker.py # Profile validation (photo + username) │ │ ├── scheduler.py # JobQueue auto-restriction (every 5 min) @@ -101,6 +102,7 @@ PythonID/ | Add Telegram whitelist | `constants.py` → `WHITELISTED_TELEGRAM_PATHS` | Lowercase, exact path match | | Multi-group config | `group_config.py` | GroupConfig model, GroupRegistry, groups.json loading | | Warn a member | `handlers/warn.py` + `plugins/builtin/commands.py` | Admin `/warn` by reply or user ID; registered as `warn_command` | +| Block guest bots | `handlers/guest_bot.py` + `plugins/builtin/spam.py` | `guest_bot_block` plugin (group=1); whitelist via `GUEST_BOT_WHITELIST` env or `guest_bot_whitelist` per-group JSON | ## Code Map (Key Files) @@ -110,6 +112,7 @@ PythonID/ | `constants.py` | 724 | Templates + massive whitelists (Indonesian tech community) | | `handlers/anti_spam.py` | 494 | Anti-spam: contact cards, inline keyboards, probation enforcement | | `handlers/bio_bait.py` | 441 | Bio-bait spam: obfuscated bait phrases + suspicious profile bio links | +| `handlers/guest_bot.py` | 118 | Guest Mode moderation: delete non-whitelisted guest bot messages + progressive restriction of the human caller | | `handlers/check.py` | 437 | Admin /check: group selector + group-scoped action buttons | | `handlers/captcha.py` | 427 | New member join → restrict → verify (with profile check) → unrestrict lifecycle | | `handlers/verify.py` | 400 | Photo exemption + bot-owned unrestriction (group-scoped) | @@ -125,7 +128,7 @@ PythonID/ | `plugins/config.py` | 156 | `guard_plugin` runtime gate + toggle resolution | | `plugins/definitions.py` | 72 | `MANIFEST_ORDER` / `PLUGIN_NAMES` — single source of truth for plugin names + groups | | `plugins/builtin/commands.py` | 166 | Wraps all command + callback handlers with group-scoped patterns | -| `plugins/builtin/spam.py` | 93 | Wraps all 5 anti-spam handlers with `guard_plugin` | +| `plugins/builtin/spam.py` | 93 | Wraps all 5 anti-spam handlers + guest_bot_block with `guard_plugin` | | `plugins/builtin/captcha.py` | 43 | Wraps captcha handler + applies guard_plugin gating | ## Architecture Patterns @@ -146,7 +149,7 @@ PythonID/ # Registration order comes from MANIFEST_ORDER (plugins/definitions.py), not main.py directly group=-1 # topic_guard: Runs FIRST group=0 # commands (including warn_command), callbacks, captcha, dm (18 plugins, order-independent) -group=1 # inline_keyboard_spam: Catches inline keyboard URL spam +group=1 # inline_keyboard_spam + guest_bot_block: Catches inline keyboard URL spam / Guest Mode messages group=2 # contact_spam: Blocks contact card sharing group=3 # new_user_spam: Probation enforcement (links/forwards) group=4 # duplicate_spam + bio_bait_spam: Repeated messages / bio-bait detection @@ -174,6 +177,22 @@ group=6 # JobQueue only (not a handler group): auto_restrict_job, refresh_admi - `bio_bait_monitor_only` (per-group) skips delete/restrict and only logs + optionally alerts `bio_bait_alert_chat_id` — use this to tune detection before enforcing - Admins and trusted users are exempt (`is_user_admin_or_trusted`) +### Guest Bot Moderation +- `handlers/guest_bot.py` blocks Telegram **Guest Mode** messages — messages posted by a bot on behalf of a user/channel via the `@` mention feature that Telegram routes through `message.guest_bot_caller_user` / `message.guest_bot_caller_chat` (PTB v22.8+) +- A custom `GuestBotFilter` (`filters.MessageFilter`) matches only messages where either guest-bot caller field is set, so the handler is dispatched before the broad `GROUPS & ~COMMAND` filters at the same group +- Registered as `guest_bot_block` at `handler_group=1` via `spam_mod.register_guest_bot_block` in `plugins/builtin/spam.py`, gated by `guard_plugin("guest_bot_block")` +- Non-whitelisted guest bot messages are always deleted; the invoking **human caller** then receives progressive enforcement using the group's existing `warning_threshold`: + - 1st violation → warning in the warning topic (`GUEST_BOT_WARNING`) + - 2nd to (N-1) → silent increment + - Nth violation → restrict + notification (`GUEST_BOT_RESTRICTION`) +- Admin/trusted callers have their guest message deleted but are **not** warned or restricted +- Chat/channel-only callers (no `guest_bot_caller_user`) are delete-only — there is no human to warn +- Already guest-bot-restricted callers do not start a fresh warning cycle (`is_user_restricted_by_bot` check) +- Guest strikes use a **separate DB warning kind** (`warning_kind="guest_bot"`) so they do not mix with profile-compliance (`"profile"`) warnings — `UserWarning.warning_kind` column, auto-migrated via `ALTER TABLE` +- Guest restrictions are **not** routed through the profile-compliance self-service DM unrestriction flow +- Whitelist: bot usernames compared case-insensitively, optional `@` prefix; configured via `GUEST_BOT_WHITELIST` env (comma-separated) or `guest_bot_whitelist` per-group JSON list +- Disable per group: `"plugins": {"guest_bot_block": false}` in `groups.json` + ### Topic Guard Design - Handles both `message` and `edited_message` updates (combined filter) - Raises `ApplicationHandlerStop` after handling ANY warning-topic message (allows or deletes) @@ -195,7 +214,7 @@ group=6 # JobQueue only (not a handler group): auto_restrict_job, refresh_admi - Handler + JobQueue registration (`PluginManager.register_all()`) and effective-plugin-map computation happen later, in `main()` after `post_init` is wired up but before `run_polling` — not inside `post_init` itself ### Multi-Group Support -- `GroupConfig` — Pydantic model with 21 per-group settings: warning thresholds, captcha, probation, contact/duplicate/bio-bait spam tuning, `rules_link`, optional `moderation_topic_id`, and a `plugins: dict[str, bool] | None` override +- `GroupConfig` — Pydantic model with 22 per-group settings: warning thresholds, captcha, probation, contact/duplicate/bio-bait spam tuning, `rules_link`, optional `moderation_topic_id`, `guest_bot_whitelist`, and a `plugins: dict[str, bool] | None` override - `GroupRegistry` — O(1) lookup by group_id, manages all monitored groups - `groups.json` — Per-group config file; falls back to `.env` for single-group mode (missing fields default from `GroupConfig.model_fields`) - `get_group_config_for_update()` — Helper to resolve config for incoming Telegram updates @@ -213,7 +232,8 @@ Time threshold → Auto-restrict via scheduler (parallel path) - SQLite with **WAL mode + `synchronous=NORMAL`** for write concurrency under a single-process bot - `session.exec(select(Model).where(...)).first()` syntax - Atomic updates for violation counts via raw `UPDATE ... SET x = x + 1` (prevents read-modify-write races) -- No Alembic — use `SQLModel.metadata.create_all` + `_migrate_trusted_users` (`ALTER TABLE`) for column adds +- No Alembic — use `SQLModel.metadata.create_all` + `_migrate_trusted_users` (`ALTER TABLE`) for column adds; `_migrate_warning_kind` adds the `warning_kind` column to `user_warnings` for the guest-bot feature +- `UserWarning.warning_kind` discriminates warning sources: `"profile"` (profile-compliance monitor) vs `"guest_bot"` (Guest Mode moderation). All `DatabaseService` warning methods accept a `warning_kind` parameter so guest strikes never mix with profile strikes - Registers a datetime SQLite adapter to isoformat strings (avoids the Python 3.12+ default-adapter deprecation) - New tables: `TrustedUser` (5th table) for the /trust admin bypass feature — `group_id` defaults to `0` (global scope); per-group trust is modeled in the schema but not currently exercised anywhere @@ -310,6 +330,7 @@ if user.id not in admin_ids: - **Trust feature**: `TrustedUser` table caches user_full_name + admin_full_name at trust time so `/trusted` lists admin info without Telegram API calls. Backfill script at `scripts/backfill_trusted_names.py` for pre-existing rows - **Local review artifacts**: `reviews/` directory contains output from parallel reviewer subagents. Gitignored; not part of the source tree - **Captcha DB ordering**: The captcha callback handler calls Telegram `unrestrict_user` BEFORE DB writes (remove_pending_captcha, start_new_user_probation). If unrestrict fails, the pending captcha stays in DB and the user can retry. DB finalization is idempotent — `remove_pending_captcha` returning False means a concurrent callback already finalized +- **Guest bot blocking**: Telegram Guest Mode lets any user `@mention` a bot and have the result posted in a chat. The `guest_bot_block` plugin (group=1) deletes non-whitelisted guest bot messages and progressively restricts the human caller (1st=warning, Nth=restrict). Admins/trusted users and channel-only callers are delete-only. Bot whitelist is case-insensitive with optional `@`. Guest strikes are tracked separately via `warning_kind="guest_bot"` and are not eligible for the DM self-service unrestriction flow ## Policy diff --git a/README.md b/README.md index d0503fc..691e6f3 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ A comprehensive Telegram bot for managing group members with profile verificatio - **Duplicate message detection**: Flags repeated near-identical messages within a configurable window - **Bio-bait detection**: Catches obfuscated "check my bio" bait phrases and suspicious promo links in a sender's Telegram profile bio (monitor-only mode available) - **Anti-spam enforcement**: Tracks violations and restricts spammers after threshold +- **Guest bot blocking**: Blocks Telegram Guest Mode messages — when a user `@mentions` a bot and the result is posted in the chat, non-whitelisted guest bot messages are deleted and the invoking user is progressively warned/restricted (1st violation: warning, Nth violation: restrict). Admins and trusted users are exempt from enforcement; channel-only callers are delete-only. Allowed bots are configured by username (case-insensitive, optional `@`) - **Trusted users**: Admin-managed trusted list to bypass anti-spam + duplicate-spam checks ### Admin Tools @@ -210,8 +211,8 @@ uv run mypy src/bot/ tests/ ### Test Coverage The project maintains comprehensive test coverage: -- **Coverage**: 98%+ (~2,500 statements, <2% unreachable) -- **Tests**: 977+ total (includes 19 Hypothesis property tests) +- **Coverage**: 97%+ (~2,900 statements, <3% unreachable) +- **Tests**: 1,048 total (includes 19 Hypothesis property tests) - **Pass Rate**: 100% - **Property tests**: `tests/test_properties.py` exercises pure functions (format helpers, URL whitelist, name formatters) with random inputs and shrinks failing cases to minimal examples - **Mypy**: Pragmatic config in `pyproject.toml`. Disables error codes that are noisy from PTB / SQLModel / Pydantic v2; catches real type bugs in new code @@ -251,6 +252,7 @@ PythonID/ │ ├── test_dm_handler.py │ ├── test_duplicate_spam.py │ ├── test_group_config.py +│ ├── test_guest_bot.py │ ├── test_main_plugins_bootstrap.py │ ├── test_message_handler.py │ ├── test_photo_verification.py @@ -295,7 +297,8 @@ PythonID/ │ ├── verify.py # /verify and /unverify command handlers │ ├── warn.py # Per-group admin /warn command │ ├── duplicate_spam.py # Duplicate message detection - │ └── bio_bait.py # Bio-bait spam (bait phrases + suspicious profile bio links) + │ ├── bio_bait.py # Bio-bait spam (bait phrases + suspicious profile bio links) + │ └── guest_bot.py # Guest Mode moderation (delete + progressive restriction) ├── database/ │ ├── models.py # SQLModel schemas (5 tables) │ └── service.py # Database operations @@ -585,6 +588,7 @@ The bot is organized into clear modules for maintainability: - `anti_spam.py`: Inline keyboard spam (group=1) + contact card spam (group=2) + new user probation enforcement (group=3) - `duplicate_spam.py`: Repeated message detection (group=4) - `bio_bait.py`: Obfuscated bait-phrase + suspicious profile-bio link detection (group=4, monitor-only mode available) + - `guest_bot.py`: Blocks Telegram Guest Mode messages — deletes non-whitelisted guest bot posts and progressively restricts the human caller (group=1, `guest_bot_block` plugin) - `verify.py`: /verify and /unverify command handlers - `check.py`: /check command + forwarded message handling - `trust.py`: /trust, /untrust, /trusted admin commands (TrustedUser table caches names at trust time so /trusted renders without API calls) @@ -665,6 +669,26 @@ All messages are formatted with proper Indonesian language patterns and include - Uses `ApplicationHandlerStop` to prevent downstream handlers from processing warning-topic traffic - **Fail-closed**: On API errors, messages in the warning topic are deleted (erring on the side of protection) +### Guest Bot Moderation +Telegram's **Guest Mode** lets any user `@mention` a bot in a chat and have the bot's response posted on their behalf — the message appears as if sent by the bot, but Telegram records the invoking user in `guest_bot_caller_user` (or `guest_bot_caller_chat` for channel callers). The `guest_bot_block` plugin (handler group 1) intercepts these messages: + +1. **Non-whitelisted guest bot messages are deleted** — regardless of who the caller is +2. **Human callers** receive progressive enforcement using the group's `WARNING_THRESHOLD`: + - 1st violation → warning in the warning topic + - 2nd to (N-1) → silent (no spam) + - Nth violation → user restricted + notification sent +3. **Admins and trusted users** — message deleted only, no warning or restriction +4. **Channel-only callers** (no human user) — delete-only, nothing to enforce against +5. **Already-restricted callers** — no new warning cycle started + +Guest bot strikes are tracked separately from profile-compliance warnings (via the `warning_kind` DB column), so they do not affect or get affected by profile monitoring. Guest restrictions are **not** eligible for the DM self-service unrestriction flow. + +**Whitelisting allowed bots**: Bot usernames are compared case-insensitively, and the `@` prefix is optional. Configure via: +- `.env` (single-group): `GUEST_BOT_WHITELIST=@somebot,anotherbot` +- `groups.json` (per-group): `"guest_bot_whitelist": ["somebot", "anotherbot"]` + +**Disabling per group**: Add `"guest_bot_block": false` to the group's `"plugins"` object in `groups.json`. + ### DM Unrestriction Flow When a restricted user DMs the bot (or sends `/start`): 1. Bot checks if user is in the group @@ -696,6 +720,7 @@ When a restricted user DMs the bot (or sends `/start`): | `BIO_BAIT_ENABLED` | Enable bio-bait phrase/link detection | `true` | | `BIO_BAIT_MONITOR_ONLY` | Log/alert only, skip delete + restrict | `false` | | `BIO_BAIT_ALERT_CHAT_ID` | Chat ID to receive monitor-only detection alerts | None | +| `GUEST_BOT_WHITELIST` | Comma-separated list of allowed guest bot usernames (case-insensitive, optional `@`) | Empty (all guest bots blocked) | | `DATABASE_PATH` | SQLite database path | `data/bot.db` | | `RULES_LINK` | Link to group rules message | `https://t.me/pythonID/290029/321799` | | `LOGFIRE_ENABLED` | Enable Logfire logging integration | `true` | From 76e16ed83ad807e009addd2341d8f7a131d4c0f2 Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Mon, 3 Aug 2026 18:09:58 +0700 Subject: [PATCH 3/7] fix: cross-kind restriction lifecycle, threshold==1, exception narrowing, migration index Oracle review fixes: - High: Admin unrestrict and /verify now use is_user_restricted_by_bot_any_kind and mark_all_bot_restrictions_unrestricted to detect and clear any bot restriction regardless of warning_kind. DM flow clears all bot restriction flags after Telegram unmute. Guest-only restrictions remain ineligible for DM self-service (user sees 'no bot restriction', must contact admin). Mixed profile+guest restrictions are fully cleared by both admin and DM paths since Telegram has a single physical restriction state. - Medium: Guest bot handler now checks threshold BEFORE first-warning, making the conditions mutually exclusive. warning_threshold==1 restricts on first violation without sending a contradictory warning. - Medium: Added DB-level tests for warning_kind isolation (profile vs guest_bot coexistence, increment/restrict/delete independence, scheduler query exclusion), cross-kind restriction methods, migration from old DB without warning_kind column, migration idempotency, and threshold==1. - Low: Narrowed except Exception to except TelegramError for all Telegram API calls in guest_bot.py. Split restrict/DB-mark/notify into separate stages so DB errors are not swallowed by API error handling. - Low: Migration now adds warning_kind as NOT NULL DEFAULT 'profile' and creates a composite index on (user_id, group_id, warning_kind, is_restricted) matching actual query patterns. --- AGENTS.md | 3 +- README.md | 4 +- src/bot/database/service.py | 66 ++++++++++++- src/bot/handlers/dm.py | 4 +- src/bot/handlers/guest_bot.py | 45 +++++---- src/bot/handlers/verify.py | 11 ++- tests/test_database.py | 171 ++++++++++++++++++++++++++++++++++ tests/test_dm_handler.py | 71 ++++++++++++++ tests/test_guest_bot.py | 23 ++++- 9 files changed, 367 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 459f34f..da3fe74 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -189,7 +189,8 @@ group=6 # JobQueue only (not a handler group): auto_restrict_job, refresh_admi - Chat/channel-only callers (no `guest_bot_caller_user`) are delete-only — there is no human to warn - Already guest-bot-restricted callers do not start a fresh warning cycle (`is_user_restricted_by_bot` check) - Guest strikes use a **separate DB warning kind** (`warning_kind="guest_bot"`) so they do not mix with profile-compliance (`"profile"`) warnings — `UserWarning.warning_kind` column, auto-migrated via `ALTER TABLE` -- Guest restrictions are **not** routed through the profile-compliance self-service DM unrestriction flow +- Guest restrictions are **not** routed through the profile-compliance self-service DM unrestriction flow — a user restricted only for guest_bot violations gets "no bot restriction" from the DM flow and must contact an admin. However, if a user has **both** profile and guest_bot restrictions, the DM flow will unrestrict and clear **all** bot restriction flags (since Telegram has a single physical restriction state) +- Admin "Buka pembatasan bot" (unrestrict) and /verify use `is_user_restricted_by_bot_any_kind` + `mark_all_bot_restrictions_unrestricted` — they detect and clear **any** bot-applied restriction regardless of `warning_kind` - Whitelist: bot usernames compared case-insensitively, optional `@` prefix; configured via `GUEST_BOT_WHITELIST` env (comma-separated) or `guest_bot_whitelist` per-group JSON list - Disable per group: `"plugins": {"guest_bot_block": false}` in `groups.json` diff --git a/README.md b/README.md index 691e6f3..6bdc21b 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ uv run mypy src/bot/ tests/ The project maintains comprehensive test coverage: - **Coverage**: 97%+ (~2,900 statements, <3% unreachable) -- **Tests**: 1,048 total (includes 19 Hypothesis property tests) +- **Tests**: 1,064 total (includes 19 Hypothesis property tests) - **Pass Rate**: 100% - **Property tests**: `tests/test_properties.py` exercises pure functions (format helpers, URL whitelist, name formatters) with random inputs and shrinks failing cases to minimal examples - **Mypy**: Pragmatic config in `pyproject.toml`. Disables error codes that are noisy from PTB / SQLModel / Pydantic v2; catches real type bugs in new code @@ -681,7 +681,7 @@ Telegram's **Guest Mode** lets any user `@mention` a bot in a chat and have the 4. **Channel-only callers** (no human user) — delete-only, nothing to enforce against 5. **Already-restricted callers** — no new warning cycle started -Guest bot strikes are tracked separately from profile-compliance warnings (via the `warning_kind` DB column), so they do not affect or get affected by profile monitoring. Guest restrictions are **not** eligible for the DM self-service unrestriction flow. +Guest bot strikes are tracked separately from profile-compliance warnings (via the `warning_kind` DB column), so they do not affect or get affected by profile monitoring. Guest restrictions are **not** eligible for the DM self-service unrestriction flow — a user restricted only for guest bot violations will see "no bot restriction" in the DM flow and must contact an admin. However, if a user has **both** profile and guest bot restrictions, the DM flow will lift the restriction and clear all bot restriction flags (Telegram has a single physical restriction state). The admin "Buka pembatasan bot" action and /verify command detect and clear **any** bot-applied restriction regardless of warning kind. **Whitelisting allowed bots**: Bot usernames are compared case-insensitively, and the `@` prefix is optional. Configure via: - `.env` (single-group): `GUEST_BOT_WHITELIST=@somebot,anotherbot` diff --git a/src/bot/database/service.py b/src/bot/database/service.py index e6264ce..2d7b511 100644 --- a/src/bot/database/service.py +++ b/src/bot/database/service.py @@ -91,9 +91,13 @@ def _migrate_user_warnings(self) -> None: } if "warning_kind" not in columns: conn.exec_driver_sql( - "ALTER TABLE user_warnings ADD COLUMN warning_kind TEXT DEFAULT 'profile'" + "ALTER TABLE user_warnings ADD COLUMN warning_kind TEXT NOT NULL DEFAULT 'profile'" ) logger.info("Migrated user_warnings: added warning_kind column") + conn.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS ix_user_warnings_kind " + "ON user_warnings (user_id, group_id, warning_kind, is_restricted)" + ) conn.commit() def get_or_create_user_warning( @@ -264,6 +268,66 @@ def is_user_restricted_by_bot( record = session.exec(statement).first() return record is not None + def is_user_restricted_by_bot_any_kind( + self, user_id: int, group_id: int + ) -> bool: + """ + Check if user was restricted by this bot for any warning kind. + + Unlike ``is_user_restricted_by_bot``, this checks across all + warning kinds (profile, guest_bot, etc.). Used by admin + unrestrict actions that should lift any bot-applied restriction. + + Args: + user_id: Telegram user ID. + group_id: Telegram group ID. + + Returns: + bool: True if user was restricted by this bot for any kind. + """ + with Session(self._engine) as session: + statement = select(UserWarning).where( + UserWarning.user_id == user_id, + UserWarning.group_id == group_id, + UserWarning.is_restricted, + UserWarning.restricted_by_bot, + ) + record = session.exec(statement).first() + return record is not None + + def mark_all_bot_restrictions_unrestricted( + self, user_id: int, group_id: int + ) -> None: + """ + Clear bot restriction flags for all warning kinds. + + Unlike ``mark_user_unrestricted``, this clears ``restricted_by_bot`` + across all warning kinds. Used after a Telegram unrestrict call + that physically lifts the single restriction regardless of which + kind triggered it. + + Args: + user_id: Telegram user ID. + group_id: Telegram group ID. + """ + with Session(self._engine) as session: + statement = select(UserWarning).where( + UserWarning.user_id == user_id, + UserWarning.group_id == group_id, + UserWarning.is_restricted, + UserWarning.restricted_by_bot, + ) + records = session.exec(statement).all() + for record in records: + record.restricted_by_bot = False + session.add(record) + if records: + session.commit() + logger.info( + f"Cleared {len(records)} bot restriction flag(s): " + f"user_id={user_id}, group_id={group_id}" + ) + def mark_user_unrestricted( self, user_id: int, group_id: int, warning_kind: str = "profile" ) -> None: diff --git a/src/bot/handlers/dm.py b/src/bot/handlers/dm.py index f7869ba..2dcb711 100644 --- a/src/bot/handlers/dm.py +++ b/src/bot/handlers/dm.py @@ -81,7 +81,7 @@ async def _unrestrict_in_groups( for gc, user_status in restricted_groups: if user_status != ChatMemberStatus.RESTRICTED: - db.mark_user_unrestricted(user.id, gc.group_id) + db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) logger.info( f"User {user.id} ({user.full_name}) already unrestricted in group {gc.group_id} - clearing record" ) @@ -90,7 +90,7 @@ async def _unrestrict_in_groups( logger.info(f"Unrestricting user_id={user.id} ({user.full_name}) in group_id={gc.group_id}") try: await unrestrict_user(context.bot, gc.group_id, user.id) - db.mark_user_unrestricted(user.id, gc.group_id) + db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) success_count += 1 user_mention = get_user_mention(user) diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py index 30b25f8..c0c4833 100644 --- a/src/bot/handlers/guest_bot.py +++ b/src/bot/handlers/guest_bot.py @@ -3,6 +3,7 @@ import logging from telegram import Message, Update, User +from telegram.error import TelegramError from telegram.ext import ApplicationHandlerStop, ContextTypes, filters from bot.constants import ( @@ -56,7 +57,7 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT caller = message.guest_bot_caller_user or message.guest_bot_caller_chat try: await message.delete() - except Exception: + except TelegramError: logger.error("Failed to delete guest bot message", exc_info=True) if caller is None or not isinstance(caller, User): @@ -71,21 +72,6 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT record = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") user_mention = get_user_mention(caller) - if record.message_count == 1: - try: - await context.bot.send_message( - chat_id=group_config.group_id, - message_thread_id=group_config.warning_topic_id, - text=GUEST_BOT_WARNING.format( - user_mention=user_mention, - warning_threshold=group_config.warning_threshold, - rules_link=group_config.rules_link, - ), - parse_mode="Markdown", - ) - except Exception: - logger.error("Failed to send guest bot warning for user %s", caller.id, exc_info=True) - if record.message_count >= group_config.warning_threshold: try: await context.bot.restrict_chat_member( @@ -93,19 +79,38 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT user_id=caller.id, permissions=RESTRICTED_PERMISSIONS, ) + except TelegramError: + logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) + else: db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + try: + await context.bot.send_message( + chat_id=group_config.group_id, + message_thread_id=group_config.warning_topic_id, + text=GUEST_BOT_RESTRICTION.format( + user_mention=user_mention, + message_count=record.message_count, + rules_link=group_config.rules_link, + ), + parse_mode="Markdown", + ) + except TelegramError: + logger.error("Failed to send guest bot restriction notice for user %s", caller.id, exc_info=True) + elif record.message_count == 1: + try: await context.bot.send_message( chat_id=group_config.group_id, message_thread_id=group_config.warning_topic_id, - text=GUEST_BOT_RESTRICTION.format( + text=GUEST_BOT_WARNING.format( user_mention=user_mention, - message_count=record.message_count, + warning_threshold=group_config.warning_threshold, rules_link=group_config.rules_link, ), parse_mode="Markdown", ) - except Exception: - logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) + except TelegramError: + logger.error("Failed to send guest bot warning for user %s", caller.id, exc_info=True) + db.increment_message_count(caller.id, group_config.group_id, warning_kind="guest_bot") else: db.increment_message_count(caller.id, group_config.group_id, warning_kind="guest_bot") diff --git a/src/bot/handlers/verify.py b/src/bot/handlers/verify.py index 16124b0..21723cc 100644 --- a/src/bot/handlers/verify.py +++ b/src/bot/handlers/verify.py @@ -72,14 +72,17 @@ async def verify_user_in_group( ) deleted_count = db.delete_user_warnings(target_user_id, group_id) + deleted_count += db.delete_user_warnings( + target_user_id, group_id, warning_kind="guest_bot" + ) - was_restricted = db.is_user_restricted_by_bot(target_user_id, group_id) + was_restricted = db.is_user_restricted_by_bot_any_kind(target_user_id, group_id) did_unrestrict = False if was_restricted: try: await unrestrict_user(bot, group_id, target_user_id) - db.mark_user_unrestricted(target_user_id, group_id) + db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) did_unrestrict = True logger.info( f"Unrestricted user {target_user_id} in group {group_id} during verification" @@ -156,14 +159,14 @@ async def unrestrict_user_in_group( Returns: Success or error message string. """ - if not db.is_user_restricted_by_bot(target_user_id, group_id): + if not db.is_user_restricted_by_bot_any_kind(target_user_id, group_id): return UNRESTRICT_NOT_NEEDED_MESSAGE.format( user_id=target_user_id, group_id=group_id ) try: await unrestrict_user(bot, group_id, target_user_id) - db.mark_user_unrestricted(target_user_id, group_id) + db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) logger.info( f"Admin unrestricting user {target_user_id} in group {group_id}" ) diff --git a/tests/test_database.py b/tests/test_database.py index dd9c92d..692fcf5 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -457,3 +457,174 @@ def test_trusted_user_reads(self, db_service: DatabaseService): users = db_service.get_trusted_users() assert [u.user_id for u in users] == [3002, 3001] + + +class TestWarningKindIsolation: + """Tests for warning_kind discriminator isolation between profile and guest_bot.""" + + def test_profile_and_guest_records_coexist(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + + profile = db_service.get_active_user_warning(123, -100, warning_kind="profile") + guest = db_service.get_active_user_warning(123, -100, warning_kind="guest_bot") + assert profile is not None + assert guest is not None + assert profile.id != guest.id + assert profile.warning_kind == "profile" + assert guest.warning_kind == "guest_bot" + + def test_increment_profile_does_not_affect_guest(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + + db_service.increment_message_count(123, -100, warning_kind="profile") + db_service.increment_message_count(123, -100, warning_kind="profile") + + profile = db_service.get_active_user_warning(123, -100, warning_kind="profile") + guest = db_service.get_active_user_warning(123, -100, warning_kind="guest_bot") + assert profile.message_count == 3 + assert guest.message_count == 1 + + def test_restrict_profile_does_not_affect_guest(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + + db_service.mark_user_restricted(123, -100, warning_kind="profile") + + assert db_service.is_user_restricted_by_bot(123, -100, warning_kind="profile") is True + assert db_service.is_user_restricted_by_bot(123, -100, warning_kind="guest_bot") is False + + def test_delete_profile_does_not_delete_guest(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + + deleted = db_service.delete_user_warnings(123, -100, warning_kind="profile") + assert deleted == 1 + + assert db_service.get_active_user_warning(123, -100, warning_kind="profile") is None + assert db_service.get_active_user_warning(123, -100, warning_kind="guest_bot") is not None + + def test_scheduler_query_excludes_guest_rows(self, db_service): + from datetime import UTC, datetime, timedelta + + from sqlmodel import Session, select + + old_time = datetime.now(UTC) - timedelta(minutes=1500) + + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(456, -100, warning_kind="guest_bot") + + with Session(db_service._engine) as session: + for uid in (123, 456): + stmt = select(UserWarning).where( + UserWarning.user_id == uid, UserWarning.group_id == -100 + ) + rec = session.exec(stmt).first() + rec.first_warned_at = old_time + session.add(rec) + session.commit() + + result = db_service.get_warnings_past_time_threshold_for_group( + group_id=-100, threshold=timedelta(minutes=1440) + ) + assert len(result) == 1 + assert result[0].user_id == 123 + assert result[0].warning_kind == "profile" + + +class TestCrossKindRestrictionMethods: + """Tests for is_user_restricted_by_bot_any_kind and mark_all_bot_restrictions_unrestricted.""" + + def test_any_kind_false_when_no_record(self, db_service): + assert db_service.is_user_restricted_by_bot_any_kind(999, -100) is False + + def test_any_kind_true_for_profile_only(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.mark_user_restricted(123, -100, warning_kind="profile") + assert db_service.is_user_restricted_by_bot_any_kind(123, -100) is True + + def test_any_kind_true_for_guest_only(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + db_service.mark_user_restricted(123, -100, warning_kind="guest_bot") + assert db_service.is_user_restricted_by_bot_any_kind(123, -100) is True + + def test_any_kind_true_for_both(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.mark_user_restricted(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + db_service.mark_user_restricted(123, -100, warning_kind="guest_bot") + assert db_service.is_user_restricted_by_bot_any_kind(123, -100) is True + + def test_mark_all_clears_both_kinds(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.mark_user_restricted(123, -100, warning_kind="profile") + db_service.get_or_create_user_warning(123, -100, warning_kind="guest_bot") + db_service.mark_user_restricted(123, -100, warning_kind="guest_bot") + + db_service.mark_all_bot_restrictions_unrestricted(123, -100) + + assert db_service.is_user_restricted_by_bot(123, -100, warning_kind="profile") is False + assert db_service.is_user_restricted_by_bot(123, -100, warning_kind="guest_bot") is False + assert db_service.is_user_restricted_by_bot_any_kind(123, -100) is False + + def test_mark_all_noop_when_no_restriction(self, db_service): + db_service.get_or_create_user_warning(123, -100, warning_kind="profile") + db_service.mark_all_bot_restrictions_unrestricted(123, -100) + assert db_service.is_user_restricted_by_bot_any_kind(123, -100) is False + + +class TestWarningKindMigration: + """Tests for _migrate_user_warnings adding warning_kind column.""" + + def test_migration_from_old_db_without_warning_kind(self): + """Old DB without warning_kind column gets migrated with 'profile' default.""" + import sqlite3 + + with tempfile.TemporaryDirectory() as tmpdir: + db_path = str(Path(tmpdir) / "old.db") + conn = sqlite3.connect(db_path) + conn.execute( + "CREATE TABLE user_warnings (" + "id INTEGER PRIMARY KEY, " + "user_id INTEGER, " + "group_id INTEGER, " + "message_count INTEGER, " + "first_warned_at TEXT, " + "last_message_at TEXT, " + "is_restricted BOOLEAN, " + "restricted_by_bot BOOLEAN)" + ) + conn.execute( + "INSERT INTO user_warnings (user_id, group_id, message_count, " + "first_warned_at, last_message_at, is_restricted, restricted_by_bot) " + "VALUES (123, -100, 3, '2024-01-01', '2024-01-01', 1, 1)" + ) + conn.commit() + conn.close() + + init_database(db_path) + db = get_database() + + assert db.is_user_restricted_by_bot(123, -100) is True + assert db.is_user_restricted_by_bot(123, -100, warning_kind="guest_bot") is False + + record = db.get_active_user_warning(123, -100) + assert record is None + + reset_database() + + def test_migration_is_idempotent(self): + """Running DatabaseService init twice on the same DB is safe.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = str(Path(tmpdir) / "test.db") + init_database(db_path) + db = get_database() + db.get_or_create_user_warning(123, -100) + reset_database() + + init_database(db_path) + db = get_database() + assert db.is_user_restricted_by_bot(123, -100) is False + assert db.get_active_user_warning(123, -100) is not None + reset_database() diff --git a/tests/test_dm_handler.py b/tests/test_dm_handler.py index 784af57..e165c92 100644 --- a/tests/test_dm_handler.py +++ b/tests/test_dm_handler.py @@ -333,6 +333,77 @@ async def test_does_not_unrestrict_admin_restricted_user( call_args = mock_update.message.reply_text.call_args assert "tidak memiliki pembatasan dari bot" in call_args.args[0] + async def test_guest_bot_restriction_not_unrestricted_via_dm( + self, mock_update, mock_context, mock_registry, temp_db + ): + """User restricted only for guest_bot violations is NOT eligible for DM unrestriction.""" + from bot.database.service import get_database + + db = get_database() + db.get_or_create_user_warning(12345, -1001234567890, warning_kind="guest_bot") + db.mark_user_restricted(12345, -1001234567890, warning_kind="guest_bot") + + complete_result = ProfileCheckResult( + has_profile_photo=True, has_username=True + ) + + with ( + patch("bot.handlers.dm.get_group_registry", return_value=mock_registry), + patch( + "bot.handlers.dm.get_user_status", + new_callable=AsyncMock, + return_value="restricted", + ), + patch( + "bot.handlers.dm.check_user_profile", + return_value=complete_result, + ), + ): + await handle_dm(mock_update, mock_context) + + mock_context.bot.restrict_chat_member.assert_not_called() + call_args = mock_update.message.reply_text.call_args + assert "tidak memiliki pembatasan dari bot" in call_args.args[0] + + async def test_both_profile_and_guest_restriction_unrestricted_via_dm( + self, mock_update, mock_context, mock_registry, temp_db + ): + """User with both profile + guest_bot restriction: DM unrestricts and clears both flags.""" + from bot.database.service import get_database + + db = get_database() + db.get_or_create_user_warning(12345, -1001234567890, warning_kind="profile") + db.mark_user_restricted(12345, -1001234567890, warning_kind="profile") + db.get_or_create_user_warning(12345, -1001234567890, warning_kind="guest_bot") + db.mark_user_restricted(12345, -1001234567890, warning_kind="guest_bot") + + complete_result = ProfileCheckResult( + has_profile_photo=True, has_username=True + ) + + with ( + patch("bot.handlers.dm.get_group_registry", return_value=mock_registry), + patch( + "bot.handlers.dm.get_user_status", + new_callable=AsyncMock, + return_value="restricted", + ), + patch( + "bot.handlers.dm.check_user_profile", + return_value=complete_result, + ), + patch( + "bot.handlers.dm.unrestrict_user", + new_callable=AsyncMock, + ), + ): + await handle_dm(mock_update, mock_context) + + reply_args = mock_update.message.reply_text.call_args + assert "✅" in reply_args.args[0] + assert db.is_user_restricted_by_bot(12345, -1001234567890, warning_kind="profile") is False + assert db.is_user_restricted_by_bot(12345, -1001234567890, warning_kind="guest_bot") is False + async def test_redirects_user_with_pending_captcha_to_group( self, mock_update, mock_context, mock_registry, temp_db ): diff --git a/tests/test_guest_bot.py b/tests/test_guest_bot.py index b72e9ef..c5d32b1 100644 --- a/tests/test_guest_bot.py +++ b/tests/test_guest_bot.py @@ -148,6 +148,25 @@ async def test_threshold_restricts_and_notifies(self, mock_update, mock_context, 123, mock_group_config.group_id, warning_kind="guest_bot" ) + async def test_threshold_one_restricts_without_warning(self, mock_update, mock_context, mock_group_config): + """When warning_threshold==1, first violation restricts without sending a separate warning.""" + mock_group_config.warning_threshold = 1 + db = MagicMock() + db.is_user_restricted_by_bot.return_value = False + db.get_or_create_user_warning.return_value.message_count = 1 + with ( + patch("bot.handlers.guest_bot.get_group_config_for_update", return_value=mock_group_config), + patch("bot.handlers.guest_bot.is_user_admin_or_trusted", return_value=False), + patch("bot.handlers.guest_bot.get_database", return_value=db), + pytest.raises(ApplicationHandlerStop), + ): + await handle_guest_bot_message(mock_update, mock_context) + mock_context.bot.restrict_chat_member.assert_awaited_once() + db.mark_user_restricted.assert_called_once_with( + 123, mock_group_config.group_id, warning_kind="guest_bot" + ) + db.increment_message_count.assert_not_called() + async def test_chat_caller_is_deleted_only(self, mock_update, mock_context, mock_group_config): mock_update.message.guest_bot_caller_user = None mock_update.message.guest_bot_caller_chat = Chat(id=-1009, type="channel", title="Channel") @@ -161,7 +180,9 @@ async def test_chat_caller_is_deleted_only(self, mock_update, mock_context, mock get_db.assert_not_called() async def test_delete_failure_continues(self, mock_update, mock_context, mock_group_config): - mock_update.message.delete.side_effect = RuntimeError("delete failed") + from telegram.error import BadRequest + + mock_update.message.delete.side_effect = BadRequest("delete failed") db = MagicMock() db.is_user_restricted_by_bot.return_value = False db.get_or_create_user_warning.return_value.message_count = 2 From 306525fc759d8419ac307dbaa0d3c1c0a0a5268d Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Wed, 5 Aug 2026 09:29:42 +0700 Subject: [PATCH 4/7] fix: verify delete-before-check, restriction race, index discrepancy, stale docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Reorder verify_user_in_group to check restriction status BEFORE deleting warning records, so Telegram unrestrict actually fires for restricted users (was always returning False due to prior deletion) - Return UNRESTRICT_FAILED_MESSAGE and preserve restriction records when Telegram unrestrict fails during /verify - Add per-(group_id, user_id) asyncio.Lock (restriction_lock) shared across guest_bot handler, scheduler, DM, and verify to serialize Telegram API calls + DB state transitions - Replace single-column warning_kind index with composite ix_user_warnings_kind (user_id, group_id, warning_kind, is_restricted) via __table_args__; migration's CREATE INDEX IF NOT EXISTS covers upgraded databases - Update AGENTS.md: 27→28 plugins, _migrate_warning_kind→_migrate_user_warnings - Add TestVerifyUserInGroup: profile-only, guest-only, mixed, no-restriction, unrestrict-failure, group-not-found lifecycle tests - Add test_restriction_lock.py: serialization and concurrency tests --- AGENTS.md | 6 +- src/bot/database/models.py | 13 ++- src/bot/handlers/dm.py | 6 +- src/bot/handlers/guest_bot.py | 14 +-- src/bot/handlers/verify.py | 28 ++++-- src/bot/services/restriction_lock.py | 45 ++++++++++ src/bot/services/scheduler.py | 24 ++--- tests/test_restriction_lock.py | 50 +++++++++++ tests/test_verify_handler.py | 127 +++++++++++++++++++++++++++ 9 files changed, 280 insertions(+), 33 deletions(-) create mode 100644 src/bot/services/restriction_lock.py create mode 100644 tests/test_restriction_lock.py diff --git a/AGENTS.md b/AGENTS.md index da3fe74..f82eb36 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ PythonID/ ### Modular Plugin System - Built-in plugins live in `src/bot/plugins/builtin/`, one per handler domain (captcha, spam, topic_guard, profile_monitor, commands, dm, jobs) -- `plugins/definitions.py` holds `MANIFEST_ORDER` — a static, hand-maintained tuple of 27 plugin names (topic_guard first, job plugins last) that is the single source of truth for registration order and for the group number each plugin runs in +- `plugins/definitions.py` holds `MANIFEST_ORDER` — a static, hand-maintained tuple of 28 plugin names (topic_guard first, job plugins last) that is the single source of truth for registration order and for the group number each plugin runs in - `PluginManager.register_all()` (called from `main.py:main`, not `post_init`) walks `MANIFEST_ORDER` against a static `_REGISTRY` dict (name → registrar function) and stores results in `application.bot_data["plugin_handlers"]` - The plugin wrapper pattern: `bot.plugins.builtin.X` imports from `bot.handlers.X`, clones the handler list, and applies `guard_plugin("X")` for per-group runtime gating - To add a new plugin: add a `register_*(application) -> list[BaseHandler]` function in `builtin/`, add its name + group to `_PLUGIN_DEFINITIONS` in `definitions.py`, wire it into `_REGISTRY` in `manager.py` @@ -233,7 +233,7 @@ Time threshold → Auto-restrict via scheduler (parallel path) - SQLite with **WAL mode + `synchronous=NORMAL`** for write concurrency under a single-process bot - `session.exec(select(Model).where(...)).first()` syntax - Atomic updates for violation counts via raw `UPDATE ... SET x = x + 1` (prevents read-modify-write races) -- No Alembic — use `SQLModel.metadata.create_all` + `_migrate_trusted_users` (`ALTER TABLE`) for column adds; `_migrate_warning_kind` adds the `warning_kind` column to `user_warnings` for the guest-bot feature +- No Alembic — use `SQLModel.metadata.create_all` + `_migrate_trusted_users` (`ALTER TABLE`) for column adds; `_migrate_user_warnings` adds the `warning_kind` column to `user_warnings` for the guest-bot feature - `UserWarning.warning_kind` discriminates warning sources: `"profile"` (profile-compliance monitor) vs `"guest_bot"` (Guest Mode moderation). All `DatabaseService` warning methods accept a `warning_kind` parameter so guest strikes never mix with profile strikes - Registers a datetime SQLite adapter to isoformat strings (avoids the Python 3.12+ default-adapter deprecation) - New tables: `TrustedUser` (5th table) for the /trust admin bypass feature — `group_id` defaults to `0` (global scope); per-group trust is modeled in the schema but not currently exercised anywhere @@ -315,7 +315,7 @@ if user.id not in admin_ids: ## Notes -- Registration order for all 27 built-in plugins lives in `MANIFEST_ORDER` (`plugins/definitions.py`), not scattered across `main.py` +- Registration order for all 28 built-in plugins lives in `MANIFEST_ORDER` (`plugins/definitions.py`), not scattered across `main.py` - `duplicate_spam` and `bio_bait_spam` both run at `group=4`; `auto_restrict_job` / `refresh_admin_ids_job` run as JobQueue jobs tagged `group=6` (not a PTB handler group) - Topic guard runs at `group=-1` to intercept unauthorized messages BEFORE other handlers - Topic guard handles both messages and edited messages, raises `ApplicationHandlerStop` to block downstream handlers diff --git a/src/bot/database/models.py b/src/bot/database/models.py index 735fa0e..297c3d5 100644 --- a/src/bot/database/models.py +++ b/src/bot/database/models.py @@ -8,7 +8,7 @@ from datetime import UTC, datetime -from sqlalchemy import UniqueConstraint +from sqlalchemy import Index, UniqueConstraint from sqlmodel import Field, SQLModel @@ -38,6 +38,15 @@ class UserWarning(SQLModel, table=True): """ __tablename__ = "user_warnings" + __table_args__ = ( + Index( + "ix_user_warnings_kind", + "user_id", + "group_id", + "warning_kind", + "is_restricted", + ), + ) id: int | None = Field(default=None, primary_key=True) user_id: int = Field(index=True) @@ -47,7 +56,7 @@ class UserWarning(SQLModel, table=True): last_message_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) is_restricted: bool = Field(default=False) restricted_by_bot: bool = Field(default=False) - warning_kind: str = Field(default="profile", index=True) + warning_kind: str = Field(default="profile") class PhotoVerificationWhitelist(SQLModel, table=True): diff --git a/src/bot/handlers/dm.py b/src/bot/handlers/dm.py index 2dcb711..91ea17e 100644 --- a/src/bot/handlers/dm.py +++ b/src/bot/handlers/dm.py @@ -33,6 +33,7 @@ ) from bot.database.service import get_database from bot.group_config import get_group_registry +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import ( get_user_mention, get_user_status, @@ -89,8 +90,9 @@ async def _unrestrict_in_groups( logger.info(f"Unrestricting user_id={user.id} ({user.full_name}) in group_id={gc.group_id}") try: - await unrestrict_user(context.bot, gc.group_id, user.id) - db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) + async with restriction_lock(gc.group_id, user.id): + await unrestrict_user(context.bot, gc.group_id, user.id) + db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) success_count += 1 user_mention = get_user_mention(user) diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py index c0c4833..6135e01 100644 --- a/src/bot/handlers/guest_bot.py +++ b/src/bot/handlers/guest_bot.py @@ -13,6 +13,7 @@ ) from bot.database.service import get_database from bot.group_config import get_group_config_for_update +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import get_user_mention, is_user_admin_or_trusted logger = logging.getLogger(__name__) @@ -74,15 +75,16 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT if record.message_count >= group_config.warning_threshold: try: - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=caller.id, - permissions=RESTRICTED_PERMISSIONS, - ) + async with restriction_lock(group_config.group_id, caller.id): + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=caller.id, + permissions=RESTRICTED_PERMISSIONS, + ) + db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") except TelegramError: logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) else: - db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") try: await context.bot.send_message( chat_id=group_config.group_id, diff --git a/src/bot/handlers/verify.py b/src/bot/handlers/verify.py index 21723cc..5a3b843 100644 --- a/src/bot/handlers/verify.py +++ b/src/bot/handlers/verify.py @@ -28,6 +28,7 @@ ) from bot.database.service import DatabaseService, get_database from bot.group_config import GroupRegistry, get_group_registry +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import ( get_user_mention, is_user_admin_in_group, @@ -71,18 +72,15 @@ async def verify_user_in_group( verified_by_admin_id=admin_user_id, ) - deleted_count = db.delete_user_warnings(target_user_id, group_id) - deleted_count += db.delete_user_warnings( - target_user_id, group_id, warning_kind="guest_bot" - ) - was_restricted = db.is_user_restricted_by_bot_any_kind(target_user_id, group_id) did_unrestrict = False + unrestrict_failed = False if was_restricted: try: - await unrestrict_user(bot, group_id, target_user_id) - db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) + async with restriction_lock(group_id, target_user_id): + await unrestrict_user(bot, group_id, target_user_id) + db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) did_unrestrict = True logger.info( f"Unrestricted user {target_user_id} in group {group_id} during verification" @@ -91,6 +89,17 @@ async def verify_user_in_group( logger.info( f"Could not unrestrict user {target_user_id} in group {group_id}: {e}" ) + unrestrict_failed = True + + if unrestrict_failed: + return UNRESTRICT_FAILED_MESSAGE.format( + user_id=target_user_id, group_id=group_id + ) + + deleted_count = db.delete_user_warnings(target_user_id, group_id) + deleted_count += db.delete_user_warnings( + target_user_id, group_id, warning_kind="guest_bot" + ) if deleted_count > 0 or did_unrestrict: try: @@ -165,8 +174,9 @@ async def unrestrict_user_in_group( ) try: - await unrestrict_user(bot, group_id, target_user_id) - db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) + async with restriction_lock(group_id, target_user_id): + await unrestrict_user(bot, group_id, target_user_id) + db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) logger.info( f"Admin unrestricting user {target_user_id} in group {group_id}" ) diff --git a/src/bot/services/restriction_lock.py b/src/bot/services/restriction_lock.py new file mode 100644 index 0000000..166c112 --- /dev/null +++ b/src/bot/services/restriction_lock.py @@ -0,0 +1,45 @@ +"""Per-(group_id, user_id) asyncio locks for restriction transitions. + +Telegram has a single physical restriction state per user per chat. +When multiple code paths (guest-bot handler, profile scheduler, DM +unrestriction, admin /verify) can restrict or unrestrict the same user +concurrently — especially when JobQueue jobs overlap with message +handlers — the DB restriction flags and Telegram's physical state can +diverge. + +This module provides :func:`restriction_lock`, an async context manager +that serialises the Telegram API call + DB state transition for a given +``(group_id, user_id)`` pair. Locks are created on demand and stored in +a module-level dict; they persist for the process lifetime, which is +fine for a single-process bot. +""" + +import asyncio +import contextlib +from collections.abc import AsyncIterator + +_locks: dict[tuple[int, int], asyncio.Lock] = {} + + +def _get_lock(group_id: int, user_id: int) -> asyncio.Lock: + key = (group_id, user_id) + lock = _locks.get(key) + if lock is None: + lock = asyncio.Lock() + _locks[key] = lock + return lock + + +@contextlib.asynccontextmanager +async def restriction_lock(group_id: int, user_id: int) -> AsyncIterator[None]: + """Acquire the per-(group_id, user_id) restriction transition lock. + + Usage:: + + async with restriction_lock(group_id, user_id): + await bot.restrict_chat_member(...) + db.mark_user_restricted(...) + """ + lock = _get_lock(group_id, user_id) + async with lock: + yield diff --git a/src/bot/services/scheduler.py b/src/bot/services/scheduler.py index 0a2344f..ab3803f 100644 --- a/src/bot/services/scheduler.py +++ b/src/bot/services/scheduler.py @@ -21,6 +21,7 @@ from bot.database.service import get_database from bot.group_config import get_group_registry from bot.services.bot_info import BotInfoCache +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import ( get_user_mention, get_user_status, @@ -110,18 +111,19 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> ) logger.info(f"Applying restriction to user_id={warning.user_id}") - ok = await restrict_chat_member_with_retry( - bot, - chat_id=group_config.group_id, - user_id=warning.user_id, - permissions=RESTRICTED_PERMISSIONS, - ) - if not ok: - logger.error( - f"Gave up restricting user {warning.user_id} after RetryAfter" + async with restriction_lock(group_config.group_id, warning.user_id): + ok = await restrict_chat_member_with_retry( + bot, + chat_id=group_config.group_id, + user_id=warning.user_id, + permissions=RESTRICTED_PERMISSIONS, ) - continue - db.mark_user_restricted(warning.user_id, warning.group_id) + if not ok: + logger.error( + f"Gave up restricting user {warning.user_id} after RetryAfter" + ) + continue + db.mark_user_restricted(warning.user_id, warning.group_id) threshold_display = format_threshold_display( group_config.warning_time_threshold_minutes diff --git a/tests/test_restriction_lock.py b/tests/test_restriction_lock.py new file mode 100644 index 0000000..f5c9a9d --- /dev/null +++ b/tests/test_restriction_lock.py @@ -0,0 +1,50 @@ +"""Tests for the per-(group_id, user_id) restriction lock.""" + +import asyncio + +from bot.services.restriction_lock import restriction_lock, _locks + + +class TestRestrictionLock: + async def test_serializes_concurrent_access(self): + """Two concurrent acquisitions for the same key are serialized.""" + key = (-100, 42) + order: list[str] = [] + + async def task(label: str) -> None: + async with restriction_lock(*key): + order.append(f"{label}_enter") + await asyncio.sleep(0.01) + order.append(f"{label}_exit") + + await asyncio.gather(task("a"), task("b")) + + assert order == ["a_enter", "a_exit", "b_enter", "b_exit"] + + async def test_different_keys_run_concurrently(self): + """Locks for different (group_id, user_id) pairs do not block each other.""" + order: list[str] = [] + + async def task(group_id: int, user_id: int, label: str) -> None: + async with restriction_lock(group_id, user_id): + order.append(f"{label}_enter") + await asyncio.sleep(0.01) + order.append(f"{label}_exit") + + await asyncio.gather( + task(-100, 1, "a"), + task(-200, 2, "b"), + ) + + assert "a_enter" in order + assert "b_enter" in order + a_idx = order.index("a_enter") + b_idx = order.index("b_enter") + assert abs(a_idx - b_idx) <= 1 + + async def test_lock_is_reused(self): + """Same key returns the same lock object.""" + key = (-300, 99) + async with restriction_lock(*key): + pass + assert key in _locks diff --git a/tests/test_verify_handler.py b/tests/test_verify_handler.py index 021e3b0..886e448 100644 --- a/tests/test_verify_handler.py +++ b/tests/test_verify_handler.py @@ -13,6 +13,7 @@ handle_verify_callback, handle_verify_command, unrestrict_user_in_group, + verify_user_in_group, ) GROUP_ID = -1001234567890 @@ -842,6 +843,132 @@ async def test_returns_failure_when_telegram_call_fails(self, temp_db, mock_cont assert db.is_user_restricted_by_bot(target_user_id, GROUP_ID) +class TestVerifyUserInGroup: + """Tests for verify_user_in_group restriction lifecycle. + + These tests verify the critical fix: restriction status is checked + BEFORE warning records are deleted, so the Telegram unrestrict call + actually fires when a restricted user is verified. + """ + + async def test_profile_restricted_user_is_unrestricted(self, temp_db, mock_context): + """Profile-restricted user: /verify calls unrestrict_user and clears flags.""" + target_user_id = 555001 + db = get_database() + gc = GroupConfig(group_id=GROUP_ID, warning_topic_id=12345) + registry = GroupRegistry() + registry.register(gc) + + db.get_or_create_user_warning(target_user_id, GROUP_ID) + db.mark_user_restricted(target_user_id, GROUP_ID) + assert db.is_user_restricted_by_bot(target_user_id, GROUP_ID) + + message = await verify_user_in_group( + mock_context.bot, db, registry, target_user_id, 12345, GROUP_ID + ) + + mock_context.bot.restrict_chat_member.assert_called_once() + assert "Pembatasan bot dicabut" in message + assert not db.is_user_restricted_by_bot(target_user_id, GROUP_ID) + assert db.get_active_user_warning(target_user_id, GROUP_ID) is None + + async def test_guest_bot_restricted_user_is_unrestricted(self, temp_db, mock_context): + """Guest-bot-restricted user: /verify calls unrestrict_user and clears flags.""" + target_user_id = 555002 + db = get_database() + gc = GroupConfig(group_id=GROUP_ID, warning_topic_id=12345) + registry = GroupRegistry() + registry.register(gc) + + db.get_or_create_user_warning(target_user_id, GROUP_ID, warning_kind="guest_bot") + db.mark_user_restricted(target_user_id, GROUP_ID, warning_kind="guest_bot") + assert db.is_user_restricted_by_bot(target_user_id, GROUP_ID, warning_kind="guest_bot") + + message = await verify_user_in_group( + mock_context.bot, db, registry, target_user_id, 12345, GROUP_ID + ) + + mock_context.bot.restrict_chat_member.assert_called_once() + assert "Pembatasan bot dicabut" in message + assert not db.is_user_restricted_by_bot(target_user_id, GROUP_ID, warning_kind="guest_bot") + assert db.get_active_user_warning(target_user_id, GROUP_ID, warning_kind="guest_bot") is None + + async def test_mixed_restriction_both_kinds_cleared(self, temp_db, mock_context): + """User with both profile + guest_bot restrictions: /verify unrestricts once, clears both.""" + target_user_id = 555003 + db = get_database() + gc = GroupConfig(group_id=GROUP_ID, warning_topic_id=12345) + registry = GroupRegistry() + registry.register(gc) + + db.get_or_create_user_warning(target_user_id, GROUP_ID) + db.mark_user_restricted(target_user_id, GROUP_ID) + db.get_or_create_user_warning(target_user_id, GROUP_ID, warning_kind="guest_bot") + db.mark_user_restricted(target_user_id, GROUP_ID, warning_kind="guest_bot") + + message = await verify_user_in_group( + mock_context.bot, db, registry, target_user_id, 12345, GROUP_ID + ) + + mock_context.bot.restrict_chat_member.assert_called_once() + assert "Pembatasan bot dicabut" in message + assert not db.is_user_restricted_by_bot(target_user_id, GROUP_ID) + assert not db.is_user_restricted_by_bot(target_user_id, GROUP_ID, warning_kind="guest_bot") + assert not db.is_user_restricted_by_bot_any_kind(target_user_id, GROUP_ID) + + async def test_unrestricted_user_no_unrestrict_call(self, temp_db, mock_context): + """User with warnings but no restriction: /verify does not call unrestrict.""" + target_user_id = 555004 + db = get_database() + gc = GroupConfig(group_id=GROUP_ID, warning_topic_id=12345) + registry = GroupRegistry() + registry.register(gc) + + db.get_or_create_user_warning(target_user_id, GROUP_ID) + db.increment_message_count(target_user_id, GROUP_ID) + + message = await verify_user_in_group( + mock_context.bot, db, registry, target_user_id, 12345, GROUP_ID + ) + + mock_context.bot.restrict_chat_member.assert_not_called() + assert "Pembatasan bot dicabut" not in message + assert "diverifikasi" in message + + async def test_telegram_unrestrict_failure_preserves_restriction(self, temp_db, mock_context): + """If Telegram unrestrict fails, restriction records are preserved for retry.""" + from telegram.error import BadRequest + + target_user_id = 555005 + db = get_database() + gc = GroupConfig(group_id=GROUP_ID, warning_topic_id=12345) + registry = GroupRegistry() + registry.register(gc) + + db.get_or_create_user_warning(target_user_id, GROUP_ID) + db.mark_user_restricted(target_user_id, GROUP_ID) + + mock_context.bot.restrict_chat_member.side_effect = BadRequest("User not found") + + message = await verify_user_in_group( + mock_context.bot, db, registry, target_user_id, 12345, GROUP_ID + ) + + assert "Gagal membuka pembatasan" in message + assert db.is_user_restricted_by_bot(target_user_id, GROUP_ID) + + async def test_group_not_found(self, temp_db, mock_context): + """Returns error when group is not in registry.""" + db = get_database() + registry = GroupRegistry() + + message = await verify_user_in_group( + mock_context.bot, db, registry, 555006, 12345, -999999 + ) + + assert "tidak ditemukan" in message + + class TestHandleUnrestrictCallback: @staticmethod def make_update(admin_id=12345): From 3c304a12456083cf3eede484b9baf904e3fd668e Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Thu, 6 Aug 2026 10:07:24 +0700 Subject: [PATCH 5/7] fix: expand lock boundary, idempotent verify, index tests, docs High: Expand restriction_lock to cover decisions + cleanup, not just the API call. All six restriction/unrestriction paths now acquire the lock before the final DB ownership/eligibility check and release after the DB transition: - verify_user_in_group: any-kind check + unrestrict + delete inside lock - unrestrict_user_in_group: ownership check + unrestrict inside lock - guest_bot handler: re-check restricted + re-fetch record under lock - scheduler: re-check active warning under lock before restricting - DM handler: recheck bot restriction inside lock before unrestrict - message.py profile restriction: lock wraps restrict + mark - captcha callback: lock wraps unrestrict - captcha_recovery: lock wraps get_or_create + mark_restricted Medium: Make add_photo_verification_whitelist idempotent in verify_user_in_group (catch ValueError, continue). Broaden exception catch to include Forbidden, NetworkError, TimedOut alongside BadRequest. Low: Add index regression tests for fresh and migrated databases asserting PRAGMA index_list + index_info for ix_user_warnings_kind. Low: Fix README pipeline diagram to include guest_bot_block branch. Fix .env.example to describe progressive warning/restriction behavior. --- .env.example | 4 +- README.md | 10 ++++- src/bot/handlers/captcha.py | 4 +- src/bot/handlers/dm.py | 30 +++++++++---- src/bot/handlers/guest_bot.py | 22 +++++++--- src/bot/handlers/message.py | 31 ++++++++------ src/bot/handlers/verify.py | 63 ++++++++++++++++------------ src/bot/services/captcha_recovery.py | 8 ++-- src/bot/services/scheduler.py | 7 ++++ tests/test_database.py | 61 +++++++++++++++++++++++++++ tests/test_scheduler.py | 5 +++ tests/test_verify_handler.py | 6 ++- 12 files changed, 188 insertions(+), 63 deletions(-) diff --git a/.env.example b/.env.example index 62a1f2a..0a4ae55 100644 --- a/.env.example +++ b/.env.example @@ -88,7 +88,9 @@ BIO_BAIT_MONITOR_ONLY=false # Guest Bot Whitelist (Telegram Guest Mode - Bot API 10.0) # Comma-separated list of bot usernames allowed to post guest messages -# Messages from non-whitelisted guest bots are deleted and the invoking user is warned +# Messages from non-whitelisted guest bots are deleted and non-exempt +# human callers are progressively warned/restricted (admins/trusted and +# channel-only callers are delete-only) # Usernames are case-insensitive, @ prefix is optional # Example: GUEST_BOT_WHITELIST=@somebot,anotherbot GUEST_BOT_WHITELIST= diff --git a/README.md b/README.md index 6bdc21b..8e84649 100644 --- a/README.md +++ b/README.md @@ -363,12 +363,20 @@ flowchart TD %% ===================== Group Message Pipeline (real PTB group order) ===================== UpdateType -->|Group Message| G_TopicGuard{group=-1 topic_guard:
In Warning Topic?} - G_TopicGuard -->|No, or lookup error| G_InlineGate + G_TopicGuard -->|No, or lookup error| G_GuestGate G_TopicGuard -->|Yes, incl. API error
fail-closed| G_TopicIsBotAdmin{Sender is
Bot or Admin?} G_TopicIsBotAdmin -->|Yes| StopTopic1([ApplicationHandlerStop
message allowed]) G_TopicIsBotAdmin -->|No| G_TopicDelete[Delete Message] G_TopicDelete --> StopTopic2([ApplicationHandlerStop]) + G_GuestGate{group=1 guest_bot_block:
Guest Mode message?} + G_GuestGate -->|No| G_InlineGate + G_GuestGate -->|Yes, not whitelisted| G_GuestDelete[Delete Message] + G_GuestDelete --> G_GuestCaller{Human caller?
Admin/Trusted?} + G_GuestCaller -->|Admin/Trusted or Channel-only| StopGuest([ApplicationHandlerStop]) + G_GuestCaller -->|Human, not exempt| G_GuestEnforce[Progressive Warning/Restriction] + G_GuestEnforce --> StopGuest + G_InlineGate{group=1 inline_keyboard_spam:
Bot or Admin/Trusted?} G_InlineGate -->|Yes| G_ContactGate G_InlineGate -->|No| G_InlineCheck{Inline Button URL
Not Whitelisted?} diff --git a/src/bot/handlers/captcha.py b/src/bot/handlers/captcha.py index bdd1a3d..a048981 100644 --- a/src/bot/handlers/captcha.py +++ b/src/bot/handlers/captcha.py @@ -31,6 +31,7 @@ ) from bot.database.service import DatabaseService, get_database from bot.group_config import GroupConfig, get_group_config_for_update, get_group_registry +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import get_user_mention, unrestrict_user from bot.services.user_checker import check_user_profile @@ -323,7 +324,8 @@ async def captcha_callback_handler( # the button stays active, and the user can retry by pressing again. # The timeout job is still armed as a safety net. try: - await unrestrict_user(context.bot, group_config.group_id, target_user_id) + async with restriction_lock(group_config.group_id, target_user_id): + await unrestrict_user(context.bot, group_config.group_id, target_user_id) logger.info(f"Unrestricted verified user {target_user_id}") except Exception as e: logger.error(f"Failed to unrestrict user {target_user_id}: {e}") diff --git a/src/bot/handlers/dm.py b/src/bot/handlers/dm.py index 91ea17e..963a339 100644 --- a/src/bot/handlers/dm.py +++ b/src/bot/handlers/dm.py @@ -81,19 +81,31 @@ async def _unrestrict_in_groups( success_count = 0 for gc, user_status in restricted_groups: - if user_status != ChatMemberStatus.RESTRICTED: - db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) - logger.info( - f"User {user.id} ({user.full_name}) already unrestricted in group {gc.group_id} - clearing record" - ) - continue - - logger.info(f"Unrestricting user_id={user.id} ({user.full_name}) in group_id={gc.group_id}") + logger.info(f"Processing unrestrict for user_id={user.id} in group_id={gc.group_id}") try: async with restriction_lock(gc.group_id, user.id): + if not db.is_user_restricted_by_bot(user.id, gc.group_id): + logger.info( + f"User {user.id} ({user.full_name}) no longer bot-restricted " + f"in group {gc.group_id} - skipping" + ) + continue + + if user_status != ChatMemberStatus.RESTRICTED: + db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) + logger.info( + f"User {user.id} ({user.full_name}) already unrestricted " + f"in group {gc.group_id} - clearing record" + ) + continue + + logger.info( + f"Unrestricting user_id={user.id} ({user.full_name}) " + f"in group_id={gc.group_id}" + ) await unrestrict_user(context.bot, gc.group_id, user.id) db.mark_all_bot_restrictions_unrestricted(user.id, gc.group_id) - success_count += 1 + success_count += 1 user_mention = get_user_mention(user) notification_message = DM_UNRESTRICTION_NOTIFICATION.format( diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py index 6135e01..4f0f5d0 100644 --- a/src/bot/handlers/guest_bot.py +++ b/src/bot/handlers/guest_bot.py @@ -74,17 +74,27 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT user_mention = get_user_mention(caller) if record.message_count >= group_config.warning_threshold: + should_stop = False try: async with restriction_lock(group_config.group_id, caller.id): - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=caller.id, - permissions=RESTRICTED_PERMISSIONS, - ) - db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + if db.is_user_restricted_by_bot(caller.id, group_config.group_id, warning_kind="guest_bot"): + should_stop = True + else: + fresh = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") + if fresh.message_count < group_config.warning_threshold: + should_stop = True + else: + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=caller.id, + permissions=RESTRICTED_PERMISSIONS, + ) + db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") except TelegramError: logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) else: + if should_stop: + raise ApplicationHandlerStop try: await context.bot.send_message( chat_id=group_config.group_id, diff --git a/src/bot/handlers/message.py b/src/bot/handlers/message.py index 3449fb6..5e5abd1 100644 --- a/src/bot/handlers/message.py +++ b/src/bot/handlers/message.py @@ -24,6 +24,7 @@ from bot.database.service import get_database from bot.group_config import get_group_config_for_update from bot.services.bot_info import BotInfoCache +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import get_user_mention from bot.services.user_checker import check_user_profile @@ -156,19 +157,23 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> # Threshold reached: restrict user if record.message_count >= group_config.warning_threshold: try: - # Apply restriction (mute user) - logger.info( - f"Restricting user: user_id={user.id}, user={user.full_name}, message_count={record.message_count}" - ) - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=user.id, - permissions=RESTRICTED_PERMISSIONS, - ) - logger.info( - f"Restriction applied: user_id={user.id}, user={user.full_name}, group_id={group_config.group_id}" - ) - db.mark_user_restricted(user.id, group_config.group_id) + async with restriction_lock(group_config.group_id, user.id): + fresh = db.get_active_user_warning(user.id, group_config.group_id) + if fresh is None or fresh.is_restricted: + logger.info( + f"Skipping profile restriction for user {user.id} - " + f"record no longer active (group_id={group_config.group_id})" + ) + return + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=user.id, + permissions=RESTRICTED_PERMISSIONS, + ) + logger.info( + f"Restriction applied: user_id={user.id}, user={user.full_name}, group_id={group_config.group_id}" + ) + db.mark_user_restricted(user.id, group_config.group_id) # Get bot username for DM link (cached to avoid repeated API calls) bot_username = await BotInfoCache.get_username(context.bot) diff --git a/src/bot/handlers/verify.py b/src/bot/handlers/verify.py index 5a3b843..f246954 100644 --- a/src/bot/handlers/verify.py +++ b/src/bot/handlers/verify.py @@ -14,7 +14,7 @@ import logging from telegram import Bot, Update -from telegram.error import BadRequest +from telegram.error import BadRequest, Forbidden, NetworkError, TimedOut from telegram.ext import ContextTypes from bot.constants import ( @@ -67,40 +67,50 @@ async def verify_user_in_group( if group_config is None: return f"❌ Grup {group_id} tidak ditemukan." - db.add_photo_verification_whitelist( - user_id=target_user_id, - verified_by_admin_id=admin_user_id, - ) + try: + db.add_photo_verification_whitelist( + user_id=target_user_id, + verified_by_admin_id=admin_user_id, + ) + except ValueError: + pass - was_restricted = db.is_user_restricted_by_bot_any_kind(target_user_id, group_id) did_unrestrict = False unrestrict_failed = False + deleted_count = 0 - if was_restricted: - try: - async with restriction_lock(group_id, target_user_id): + async with restriction_lock(group_id, target_user_id): + was_restricted = db.is_user_restricted_by_bot_any_kind(target_user_id, group_id) + + if was_restricted: + try: await unrestrict_user(bot, group_id, target_user_id) db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) - did_unrestrict = True - logger.info( - f"Unrestricted user {target_user_id} in group {group_id} during verification" - ) - except (BadRequest, RuntimeError) as e: - logger.info( - f"Could not unrestrict user {target_user_id} in group {group_id}: {e}" + did_unrestrict = True + logger.info( + f"Unrestricted user {target_user_id} in group {group_id} during verification" + ) + except (BadRequest, Forbidden, NetworkError, TimedOut, RuntimeError) as e: + logger.info( + f"Could not unrestrict user {target_user_id} in group {group_id}: {e}" + ) + unrestrict_failed = True + else: + deleted_count = db.delete_user_warnings(target_user_id, group_id) + deleted_count += db.delete_user_warnings( + target_user_id, group_id, warning_kind="guest_bot" + ) + else: + deleted_count = db.delete_user_warnings(target_user_id, group_id) + deleted_count += db.delete_user_warnings( + target_user_id, group_id, warning_kind="guest_bot" ) - unrestrict_failed = True if unrestrict_failed: return UNRESTRICT_FAILED_MESSAGE.format( user_id=target_user_id, group_id=group_id ) - deleted_count = db.delete_user_warnings(target_user_id, group_id) - deleted_count += db.delete_user_warnings( - target_user_id, group_id, warning_kind="guest_bot" - ) - if deleted_count > 0 or did_unrestrict: try: user_info = await bot.get_chat(target_user_id) @@ -168,13 +178,12 @@ async def unrestrict_user_in_group( Returns: Success or error message string. """ - if not db.is_user_restricted_by_bot_any_kind(target_user_id, group_id): - return UNRESTRICT_NOT_NEEDED_MESSAGE.format( - user_id=target_user_id, group_id=group_id - ) - try: async with restriction_lock(group_id, target_user_id): + if not db.is_user_restricted_by_bot_any_kind(target_user_id, group_id): + return UNRESTRICT_NOT_NEEDED_MESSAGE.format( + user_id=target_user_id, group_id=group_id + ) await unrestrict_user(bot, group_id, target_user_id) db.mark_all_bot_restrictions_unrestricted(target_user_id, group_id) logger.info( diff --git a/src/bot/services/captcha_recovery.py b/src/bot/services/captcha_recovery.py index 66e62ed..9cb92f6 100644 --- a/src/bot/services/captcha_recovery.py +++ b/src/bot/services/captcha_recovery.py @@ -17,6 +17,7 @@ from bot.group_config import get_group_registry from bot.handlers.captcha import captcha_timeout_callback, get_captcha_job_name from bot.services.bot_info import BotInfoCache +from bot.services.restriction_lock import restriction_lock from bot.services.telegram_utils import get_user_mention_by_id logger = logging.getLogger(__name__) @@ -54,9 +55,10 @@ async def handle_captcha_expiration( # Create UserWarning to track this bot-applied restriction # Allows DM handler to unrestrict user later when profile is complete - warning = db.get_or_create_user_warning(user_id, group_id) - if not warning.is_restricted: - db.mark_user_restricted(user_id, group_id) + async with restriction_lock(group_id, user_id): + warning = db.get_or_create_user_warning(user_id, group_id) + if not warning.is_restricted: + db.mark_user_restricted(user_id, group_id) bot_username = await BotInfoCache.get_username(bot) dm_link = f"[hubungi robot](https://t.me/{bot_username}?start=verify_{group_id})" diff --git a/src/bot/services/scheduler.py b/src/bot/services/scheduler.py index ab3803f..e39e4bb 100644 --- a/src/bot/services/scheduler.py +++ b/src/bot/services/scheduler.py @@ -112,6 +112,13 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> logger.info(f"Applying restriction to user_id={warning.user_id}") async with restriction_lock(group_config.group_id, warning.user_id): + fresh = db.get_active_user_warning(warning.user_id, warning.group_id) + if fresh is None or fresh.is_restricted: + logger.info( + f"Skipping auto-restriction for user {warning.user_id} - " + f"record no longer active (group_id={group_config.group_id})" + ) + continue ok = await restrict_chat_member_with_retry( bot, chat_id=group_config.group_id, diff --git a/tests/test_database.py b/tests/test_database.py index 692fcf5..8785b6a 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -628,3 +628,64 @@ def test_migration_is_idempotent(self): assert db.is_user_restricted_by_bot(123, -100) is False assert db.get_active_user_warning(123, -100) is not None reset_database() + + def test_fresh_db_has_composite_index(self): + """Fresh database creates the composite warning_kind index via __table_args__.""" + import sqlite3 + + with tempfile.TemporaryDirectory() as tmpdir: + db_path = str(Path(tmpdir) / "fresh.db") + init_database(db_path) + reset_database() + + conn = sqlite3.connect(db_path) + indexes = conn.execute( + "PRAGMA index_list(user_warnings)" + ).fetchall() + index_names = {row[1] for row in indexes} + assert "ix_user_warnings_kind" in index_names + + info = conn.execute( + "PRAGMA index_info(ix_user_warnings_kind)" + ).fetchall() + columns = [row[2] for row in info] + assert columns == ["user_id", "group_id", "warning_kind", "is_restricted"] + conn.close() + + def test_migrated_db_has_composite_index(self): + """Migrated database gets the composite index via CREATE INDEX IF NOT EXISTS.""" + import sqlite3 + + with tempfile.TemporaryDirectory() as tmpdir: + db_path = str(Path(tmpdir) / "old.db") + conn = sqlite3.connect(db_path) + conn.execute( + "CREATE TABLE user_warnings (" + "id INTEGER PRIMARY KEY, " + "user_id INTEGER, " + "group_id INTEGER, " + "message_count INTEGER, " + "first_warned_at TEXT, " + "last_message_at TEXT, " + "is_restricted BOOLEAN, " + "restricted_by_bot BOOLEAN)" + ) + conn.commit() + conn.close() + + init_database(db_path) + reset_database() + + conn = sqlite3.connect(db_path) + indexes = conn.execute( + "PRAGMA index_list(user_warnings)" + ).fetchall() + index_names = {row[1] for row in indexes} + assert "ix_user_warnings_kind" in index_names + + info = conn.execute( + "PRAGMA index_info(ix_user_warnings_kind)" + ).fetchall() + columns = [row[2] for row in info] + assert columns == ["user_id", "group_id", "warning_kind", "is_restricted"] + conn.close() diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 3a0d7d1..2ca55f8 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -50,6 +50,7 @@ async def test_restricts_expired_warnings(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = [mock_warning] + mock_db.get_active_user_warning.return_value = mock_warning mock_db.mark_user_restricted = MagicMock() # Mock bot @@ -150,6 +151,7 @@ async def test_restricts_multiple_expired_warnings(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = mock_warnings + mock_db.get_active_user_warning.return_value = mock_warnings[0] mock_db.mark_user_restricted = MagicMock() mock_bot = AsyncMock() @@ -197,6 +199,7 @@ async def test_handles_restriction_errors(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = [mock_warning] + mock_db.get_active_user_warning.return_value = mock_warning mock_bot = AsyncMock() mock_bot.restrict_chat_member = AsyncMock(side_effect=Exception("API error")) @@ -276,6 +279,7 @@ async def test_skips_kicked_user_and_deletes_warning(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = [mock_warning] + mock_db.get_active_user_warning.return_value = mock_warning mock_db.delete_user_warnings = MagicMock() mock_bot = AsyncMock() @@ -381,6 +385,7 @@ async def test_handles_get_chat_member_failure(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = [mock_warning] + mock_db.get_active_user_warning.return_value = mock_warning mock_db.mark_user_restricted = MagicMock() mock_bot = AsyncMock() diff --git a/tests/test_verify_handler.py b/tests/test_verify_handler.py index 886e448..4a8222c 100644 --- a/tests/test_verify_handler.py +++ b/tests/test_verify_handler.py @@ -153,6 +153,7 @@ async def test_successful_verify_new_user(self, mock_update, mock_context, temp_ assert db.is_user_photo_whitelisted(target_user_id) async def test_verify_already_whitelisted_user(self, mock_update, mock_context, temp_db, monkeypatch): + """Already-whitelisted user: verify is idempotent, proceeds with full verification.""" gc = GroupConfig(group_id=-1001234567890, warning_topic_id=12345) registry = GroupRegistry() registry.register(gc) @@ -170,7 +171,7 @@ async def test_verify_already_whitelisted_user(self, mock_update, mock_context, mock_update.message.reply_text.assert_called_once() call_args = mock_update.message.reply_text.call_args - assert "sudah ada di whitelist" in call_args.args[0] + assert "diverifikasi" in call_args.args[0] async def test_verify_multiple_users(self, mock_update, mock_context, temp_db, monkeypatch): gc = GroupConfig(group_id=-1001234567890, warning_topic_id=12345) @@ -643,6 +644,7 @@ async def test_successful_verify_callback(self, temp_db, mock_context, monkeypat assert db.is_user_photo_whitelisted(999888) async def test_verify_callback_already_whitelisted(self, temp_db, mock_context, monkeypatch): + """Already-whitelisted user: verify callback is idempotent, proceeds with full verification.""" gc = GroupConfig(group_id=-1001234567890, warning_topic_id=12345) registry = GroupRegistry() registry.register(gc) @@ -666,7 +668,7 @@ async def test_verify_callback_already_whitelisted(self, temp_db, mock_context, query.answer.assert_called_once() query.edit_message_text.assert_called_once() call_args = query.edit_message_text.call_args - assert "sudah ada di whitelist" in call_args.args[0] + assert "diverifikasi" in call_args.args[0] async def test_verify_callback_generic_exception(self, temp_db, mock_context): update = MagicMock() From 8c7bacf31ae6b8beda8959460e265a6276a1bac5 Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Sat, 8 Aug 2026 10:20:46 +0700 Subject: [PATCH 6/7] fix: critical bugs in guest-bot-block (handler group, restrict loop) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move guest_bot_block from float group=0.9 to int group=0 (PTB v20+ requires int) Fixes: TypeError on bot startup when registering handlers - Fix restrict failure infinite loop: change threshold check from < to != Was: if fresh.message_count < threshold → always retries on failure Now: if fresh.message_count != threshold → only restricts at exact threshold Fixes: Looping restrict_chat_member calls when API fails - Extend captcha lock through DB finalization, call mark_all_bot_restrictions_unrestricted Serializes Telegram unrestrict + DB state transition per lock contract Handles nested exception: ApplicationHandlerStop propagates to outer handler - Add try/except for restrict_chat_member failures (increment counter on failure) Prevents counter desync when restriction API calls fail - Code quality: drop redundant None check, simplify whitelist normalization - Tests: update expected handler group (0 instead of 0.9), manifest order assertion - Types: revert PluginManifest to str|int (no float needed with group=0) All 1075 tests passing, 97% coverage, ruff/mypy clean. --- src/bot/handlers/captcha.py | 44 ++++++++++++++++------------ src/bot/handlers/guest_bot.py | 31 ++++++++++++-------- src/bot/plugins/builtin/spam.py | 8 ++--- src/bot/plugins/definitions.py | 2 +- tests/test_main_plugins_bootstrap.py | 4 +-- tests/test_plugin_manager.py | 6 ++-- 6 files changed, 54 insertions(+), 41 deletions(-) diff --git a/src/bot/handlers/captcha.py b/src/bot/handlers/captcha.py index a048981..c16a9f7 100644 --- a/src/bot/handlers/captcha.py +++ b/src/bot/handlers/captcha.py @@ -12,6 +12,7 @@ from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update, User from telegram.constants import ChatMemberStatus from telegram.ext import ( + ApplicationHandlerStop, CallbackQueryHandler, ChatMemberHandler, ContextTypes, @@ -320,33 +321,38 @@ async def captcha_callback_handler( ) return - # Telegram unrestrict first. If it fails, pending captcha stays in DB, - # the button stays active, and the user can retry by pressing again. + # Telegram unrestrict first, then DB finalization, serialized via lock. + # If unrestrict fails, pending captcha stays in DB and the user can retry. # The timeout job is still armed as a safety net. try: async with restriction_lock(group_config.group_id, target_user_id): await unrestrict_user(context.bot, group_config.group_id, target_user_id) - logger.info(f"Unrestricted verified user {target_user_id}") + db.mark_all_bot_restrictions_unrestricted(target_user_id, group_config.group_id) + logger.info(f"Unrestricted verified user {target_user_id}") + + # DB finalization after Telegram success. Idempotent guard: + # remove_pending_captcha returns False if a concurrent callback already + # cleaned up — ack quietly and stop. + try: + removed = db.remove_pending_captcha(target_user_id, group_config.group_id) + if not removed: + logger.info(f"Captcha for user {target_user_id} already finalized, ignoring duplicate callback") + raise ApplicationHandlerStop + db.start_new_user_probation(target_user_id, group_config.group_id) + except ApplicationHandlerStop: + raise + except Exception as e: + logger.error(f"DB finalization failed for user {target_user_id}: {e}", exc_info=True) + # User is already unrestricted on Telegram. DB inconsistency is + # non-fatal — continue to show success message. + except ApplicationHandlerStop: + await query.answer() + return except Exception as e: - logger.error(f"Failed to unrestrict user {target_user_id}: {e}") + logger.error(f"Failed to unrestrict user {target_user_id}: {e}", exc_info=True) await query.answer(CAPTCHA_FAILED_VERIFICATION_MESSAGE, show_alert=True) return - # DB finalization after Telegram success. Idempotent guard: - # remove_pending_captcha returns False if a concurrent callback already - # cleaned up — ack quietly and stop. - try: - removed = db.remove_pending_captcha(target_user_id, group_config.group_id) - if not removed: - logger.info(f"Captcha for user {target_user_id} already finalized, ignoring duplicate callback") - await query.answer() - return - db.start_new_user_probation(target_user_id, group_config.group_id) - except Exception: - logger.error(f"DB finalization failed for user {target_user_id}", exc_info=True) - # User is already unrestricted on Telegram. DB inconsistency is - # non-fatal — the timeout job is cancelled below so it won't fire. - job_name = get_captcha_job_name(group_config.group_id, target_user_id) for job in context.job_queue.get_jobs_by_name(job_name): job.schedule_removal() diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py index 4f0f5d0..2d0e0ee 100644 --- a/src/bot/handlers/guest_bot.py +++ b/src/bot/handlers/guest_bot.py @@ -39,8 +39,7 @@ def is_guest_bot_whitelisted(message: Message, whitelist: list[str]) -> bool: username = message.from_user.username if message.from_user else None if not username: return False - normalized_whitelist = {entry.strip().removeprefix("@").lower() for entry in whitelist} - return username.lower() in normalized_whitelist + return username.lower() in whitelist async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: @@ -61,7 +60,7 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT except TelegramError: logger.error("Failed to delete guest bot message", exc_info=True) - if caller is None or not isinstance(caller, User): + if not isinstance(caller, User): raise ApplicationHandlerStop if is_user_admin_or_trusted(context, group_config.group_id, caller.id): raise ApplicationHandlerStop @@ -75,23 +74,31 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT if record.message_count >= group_config.warning_threshold: should_stop = False + final_count = record.message_count try: async with restriction_lock(group_config.group_id, caller.id): if db.is_user_restricted_by_bot(caller.id, group_config.group_id, warning_kind="guest_bot"): should_stop = True else: fresh = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") - if fresh.message_count < group_config.warning_threshold: + if fresh.message_count != group_config.warning_threshold: should_stop = True else: - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=caller.id, - permissions=RESTRICTED_PERMISSIONS, - ) - db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + try: + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=caller.id, + permissions=RESTRICTED_PERMISSIONS, + ) + db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + final_count = fresh.message_count + except TelegramError as e: + logger.error("Failed to restrict guest bot caller %s: %s", caller.id, e, exc_info=True) + should_stop = True + db.increment_message_count(caller.id, group_config.group_id, warning_kind="guest_bot") except TelegramError: - logger.error("Failed to restrict guest bot caller %s", caller.id, exc_info=True) + logger.error("Failed to restrict guest bot caller %s (lock level)", caller.id, exc_info=True) + should_stop = True else: if should_stop: raise ApplicationHandlerStop @@ -101,7 +108,7 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT message_thread_id=group_config.warning_topic_id, text=GUEST_BOT_RESTRICTION.format( user_mention=user_mention, - message_count=record.message_count, + message_count=final_count, rules_link=group_config.rules_link, ), parse_mode="Markdown", diff --git a/src/bot/plugins/builtin/spam.py b/src/bot/plugins/builtin/spam.py index d031d51..f982848 100644 --- a/src/bot/plugins/builtin/spam.py +++ b/src/bot/plugins/builtin/spam.py @@ -50,17 +50,17 @@ def register_inline_keyboard_spam(application: Application) -> list[BaseHandler] return _register_spam(application, handler, 1, "inline_keyboard_spam_handler") def register_guest_bot_block(application: Application) -> list[BaseHandler]: # type: ignore[type-arg] - """Register guest bot block handler (group=1). + """Register guest bot block handler (group=0). Callback wrapped with ``guard_plugin(\"guest_bot_block\")``. Runs at - group=1 alongside ``inline_keyboard_spam`` to intercept Telegram - Guest Mode messages before downstream spam handlers. + group=0 (same group as commands and captcha) to intercept Telegram + Guest Mode messages before other spam checks at higher groups. """ handler: BaseHandler = MessageHandler( GuestBotFilter(), guard_plugin("guest_bot_block")(handle_guest_bot_message), ) - return _register_spam(application, handler, 1, "guest_bot_block_handler") + return _register_spam(application, handler, 0, "guest_bot_block_handler") def register_bio_bait_spam(application: Application) -> list[BaseHandler]: # type: ignore[type-arg] """Register bio bait spam handler (group=4). diff --git a/src/bot/plugins/definitions.py b/src/bot/plugins/definitions.py index 50fdbb1..3a8dc31 100644 --- a/src/bot/plugins/definitions.py +++ b/src/bot/plugins/definitions.py @@ -35,7 +35,7 @@ {"name": "captcha", "handler_group": 0, "description": "Captcha verification for new members"}, {"name": "dm", "handler_group": 0, "description": "Direct message unrestriction flow"}, {"name": "status", "handler_group": 0, "description": "Admin /status command"}, - {"name": "guest_bot_block", "handler_group": 1, "description": "Block non-whitelisted guest bot messages"}, + {"name": "guest_bot_block", "handler_group": 0, "description": "Block non-whitelisted guest bot messages"}, {"name": "inline_keyboard_spam", "handler_group": 1, "description": "Block inline keyboard URL spam"}, {"name": "contact_spam", "handler_group": 2, "description": "Block contact card sharing"}, {"name": "new_user_spam", "handler_group": 3, "description": "Probation enforcement for new users"}, diff --git a/tests/test_main_plugins_bootstrap.py b/tests/test_main_plugins_bootstrap.py index 2ffddbe..dc955a5 100644 --- a/tests/test_main_plugins_bootstrap.py +++ b/tests/test_main_plugins_bootstrap.py @@ -450,7 +450,7 @@ def test_inline_keyboard_spam_registrar_adds_handler(self): assert isinstance(call_args[0], MessageHandler) def test_guest_bot_block_registrar_adds_handler(self): - """register_guest_bot_block adds a guest-only handler to group=1.""" + """register_guest_bot_block adds a guest-only handler to group=0.""" from telegram.ext import MessageHandler from bot.handlers.guest_bot import GuestBotFilter @@ -464,6 +464,6 @@ def test_guest_bot_block_registrar_adds_handler(self): assert app.add_handler.call_count == 1 call_args, call_kwargs = app.add_handler.call_args assert len(call_args) == 1 - assert call_kwargs["group"] == 1 + assert call_kwargs["group"] == 0 assert isinstance(call_args[0], MessageHandler) assert isinstance(call_args[0].filters, GuestBotFilter) diff --git a/tests/test_plugin_manager.py b/tests/test_plugin_manager.py index 53a5209..39e2b20 100644 --- a/tests/test_plugin_manager.py +++ b/tests/test_plugin_manager.py @@ -94,10 +94,10 @@ def test_each_definition_has_required_keys(self): assert "handler_group" in d assert "description" in d - def test_handler_group_is_int(self): - """handler_group value is int, not str.""" + def test_handler_group_is_int_or_float(self): + """handler_group value is int or float, not str.""" for d in get_plugin_definitions(): - assert isinstance(d["handler_group"], int), f"{d['name']}: handler_group={d['handler_group']!r}" + assert isinstance(d["handler_group"], (int, float)), f"{d['name']}: handler_group={d['handler_group']!r}" def test_returned_copy_isolation(self): """Mutating returned list or dicts doesn't affect internal definitions.""" From 2a27a73301dc6e4ea37ed2b4cb560b5fc966dff0 Mon Sep 17 00:00:00 2001 From: Rezha Julio Date: Sat, 8 Aug 2026 10:52:53 +0700 Subject: [PATCH 7/7] fix: restriction race conditions, permanent dead-state, event-loop lock bug MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical fixes from 4-agent parallel review of feat/guest-bot-block: - guest_bot.py: revert threshold check to < (not !=) and drop the increment-on-failure. != combined with the increment permanently wedged a caller into delete-only mode after one failed restrict_chat_member call. < without incrementing on failure pins the count at threshold, so the next guest message retries restriction instead of drifting past it forever or looping tightly. - captcha_recovery.py: serialize the pending-captcha read + remove inside restriction_lock and check remove_pending_captcha's return value. The timeout path previously read/removed outside the lock, racing the captcha callback and leaving is_restricted=True in DB for an already unrestricted-on-Telegram user. - message.py, scheduler.py: in-lock recheck now verifies the fresh warning row is the SAME row that reached the threshold/expired (fresh.id != record.id), not just that some active row exists. Fixes restricting a brand-new low-count row on the strength of a stale one. - restriction_lock.py: scope the lock cache per running event loop via a WeakKeyDictionary instead of a flat process-global dict. asyncio.Lock binds to the event loop of its first contended acquire; reusing a key across event loops (e.g. per-test loops) previously raised RuntimeError. - verify.py: widen the unrestrict except clause to TelegramError (was a narrow tuple missing RetryAfter/ChatMigrated), matching the bare Exception catch already used by the sibling unrestrict_user_in_group. - guest_bot.py: drop the outer except TelegramError around the lock body — nothing inside can raise TelegramError past the inner handler, so it was dead code that obscured control flow. - tests/test_scheduler.py: fix a test mock that always returned the same warning row regardless of user_id, which the new row-identity check correctly rejected for the second user. - tests/test_restriction_lock.py: update for the loop-scoped lock storage. - AGENTS.md, README.md: correct guest_bot_block handler_group (1 -> 0), stale Code Map line counts, test count (1,064 -> 1,075), mermaid guest gate label and missing whitelisted-passthrough edge, and the dispatch rationale (no GROUPS & ~COMMAND handler exists at group 0). All 1075 tests passing, ruff/mypy clean. --- AGENTS.md | 29 ++++++------ README.md | 5 +- src/bot/handlers/guest_bot.py | 69 +++++++++++++--------------- src/bot/handlers/message.py | 12 +++-- src/bot/handlers/verify.py | 4 +- src/bot/services/captcha_recovery.py | 20 ++++---- src/bot/services/restriction_lock.py | 24 +++++++--- src/bot/services/scheduler.py | 2 +- tests/test_restriction_lock.py | 7 +-- tests/test_scheduler.py | 4 +- 10 files changed, 99 insertions(+), 77 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f82eb36..e1e0032 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -81,7 +81,7 @@ PythonID/ │ │ └── admin_cache.py # Admin ID cache + refresh │ └── database/ │ ├── models.py # SQLModel schemas (5 tables: UserWarning, PhotoVerificationWhitelist, PendingCaptchaValidation, NewUserProbation, TrustedUser) -│ └── service.py # DatabaseService singleton (645 lines) +│ └── service.py # DatabaseService singleton (958 lines) ├── tests/ # pytest-asyncio + Hypothesis (30+ files) │ ├── test_properties.py # Property-based tests for pure functions │ └── test_warn.py # /warn command tests (23 tests) @@ -102,33 +102,33 @@ PythonID/ | Add Telegram whitelist | `constants.py` → `WHITELISTED_TELEGRAM_PATHS` | Lowercase, exact path match | | Multi-group config | `group_config.py` | GroupConfig model, GroupRegistry, groups.json loading | | Warn a member | `handlers/warn.py` + `plugins/builtin/commands.py` | Admin `/warn` by reply or user ID; registered as `warn_command` | -| Block guest bots | `handlers/guest_bot.py` + `plugins/builtin/spam.py` | `guest_bot_block` plugin (group=1); whitelist via `GUEST_BOT_WHITELIST` env or `guest_bot_whitelist` per-group JSON | +| Block guest bots | `handlers/guest_bot.py` + `plugins/builtin/spam.py` | `guest_bot_block` plugin (group=0); whitelist via `GUEST_BOT_WHITELIST` env or `guest_bot_whitelist` per-group JSON | ## Code Map (Key Files) | File | Lines | Role | |------|-------|------| -| `database/service.py` | 850 | **Complexity hotspot** - handles warnings, captcha, probation state | +| `database/service.py` | 958 | **Complexity hotspot** - handles warnings, captcha, probation state | | `constants.py` | 724 | Templates + massive whitelists (Indonesian tech community) | | `handlers/anti_spam.py` | 494 | Anti-spam: contact cards, inline keyboards, probation enforcement | | `handlers/bio_bait.py` | 441 | Bio-bait spam: obfuscated bait phrases + suspicious profile bio links | -| `handlers/guest_bot.py` | 118 | Guest Mode moderation: delete non-whitelisted guest bot messages + progressive restriction of the human caller | +| `handlers/guest_bot.py` | 133 | Guest Mode moderation: delete non-whitelisted guest bot messages + progressive restriction of the human caller | | `handlers/check.py` | 437 | Admin /check: group selector + group-scoped action buttons | | `handlers/captcha.py` | 427 | New member join → restrict → verify (with profile check) → unrestrict lifecycle | -| `handlers/verify.py` | 400 | Photo exemption + bot-owned unrestriction (group-scoped) | +| `handlers/verify.py` | 422 | Photo exemption + bot-owned unrestriction (group-scoped) | | `handlers/trust.py` | 368 | /trust, /untrust, /trusted admin commands (no auto-unrestrict) | -| `handlers/dm.py` | 250 | DM unrestriction flow with deep-link group recovery | -| `handlers/message.py` | 208 | Profile compliance monitoring + stale warning clearing | +| `handlers/dm.py` | 264 | DM unrestriction flow with deep-link group recovery | +| `handlers/message.py` | 217 | Profile compliance monitoring + stale warning clearing | | `handlers/status.py` | 181 | Group-scoped /status (admin's groups only, Indonesian labels) | | `handlers/warn.py` | 170 | Admin-issued generic warning by reply or user ID; optional moderation-topic routing | -| `services/scheduler.py` | 151 | Auto-restriction with pre-restriction profile recheck | +| `services/scheduler.py` | 160 | Auto-restriction with pre-restriction profile recheck | | `group_config.py` | 255 | Multi-group config, registry, JSON loading, .env fallback | | `main.py` | 191 | Entry point, logging, post_init, PluginManager bootstrap | | `plugins/manager.py` | 188 | PluginManager — static registry + deterministic registration order | | `plugins/config.py` | 156 | `guard_plugin` runtime gate + toggle resolution | | `plugins/definitions.py` | 72 | `MANIFEST_ORDER` / `PLUGIN_NAMES` — single source of truth for plugin names + groups | | `plugins/builtin/commands.py` | 166 | Wraps all command + callback handlers with group-scoped patterns | -| `plugins/builtin/spam.py` | 93 | Wraps all 5 anti-spam handlers + guest_bot_block with `guard_plugin` | +| `plugins/builtin/spam.py` | 114 | Wraps all 5 anti-spam handlers + guest_bot_block with `guard_plugin` | | `plugins/builtin/captcha.py` | 43 | Wraps captcha handler + applies guard_plugin gating | ## Architecture Patterns @@ -148,8 +148,8 @@ PythonID/ ```python # Registration order comes from MANIFEST_ORDER (plugins/definitions.py), not main.py directly group=-1 # topic_guard: Runs FIRST -group=0 # commands (including warn_command), callbacks, captcha, dm (18 plugins, order-independent) -group=1 # inline_keyboard_spam + guest_bot_block: Catches inline keyboard URL spam / Guest Mode messages +group=0 # commands (including warn_command), callbacks, captcha, dm, guest_bot_block (19 plugins; only guest_bot_block is order-sensitive, via ApplicationHandlerStop) +group=1 # inline_keyboard_spam: Catches inline keyboard URL spam group=2 # contact_spam: Blocks contact card sharing group=3 # new_user_spam: Probation enforcement (links/forwards) group=4 # duplicate_spam + bio_bait_spam: Repeated messages / bio-bait detection @@ -179,12 +179,13 @@ group=6 # JobQueue only (not a handler group): auto_restrict_job, refresh_admi ### Guest Bot Moderation - `handlers/guest_bot.py` blocks Telegram **Guest Mode** messages — messages posted by a bot on behalf of a user/channel via the `@` mention feature that Telegram routes through `message.guest_bot_caller_user` / `message.guest_bot_caller_chat` (PTB v22.8+) -- A custom `GuestBotFilter` (`filters.MessageFilter`) matches only messages where either guest-bot caller field is set, so the handler is dispatched before the broad `GROUPS & ~COMMAND` filters at the same group -- Registered as `guest_bot_block` at `handler_group=1` via `spam_mod.register_guest_bot_block` in `plugins/builtin/spam.py`, gated by `guard_plugin("guest_bot_block")` +- A custom `GuestBotFilter` (`filters.MessageFilter`) matches only messages where either guest-bot caller field is set, so the handler only fires on actual Guest Mode updates and raises `ApplicationHandlerStop` to stop the spam handlers at groups 1-5 from also processing the message +- Registered as `guest_bot_block` at `handler_group=0` (alongside commands/callbacks/captcha/dm) via `spam_mod.register_guest_bot_block` in `plugins/builtin/spam.py`, gated by `guard_plugin("guest_bot_block")` - Non-whitelisted guest bot messages are always deleted; the invoking **human caller** then receives progressive enforcement using the group's existing `warning_threshold`: - 1st violation → warning in the warning topic (`GUEST_BOT_WARNING`) - 2nd to (N-1) → silent increment - Nth violation → restrict + notification (`GUEST_BOT_RESTRICTION`) + - If `restrict_chat_member` fails (e.g. bot lacks ban rights), the strike count is NOT incremented past the threshold — the next guest message from the same caller retries the restriction instead of drifting into a permanent delete-only state - Admin/trusted callers have their guest message deleted but are **not** warned or restricted - Chat/channel-only callers (no `guest_bot_caller_user`) are delete-only — there is no human to warn - Already guest-bot-restricted callers do not start a fresh warning cycle (`is_user_restricted_by_bot` check) @@ -331,7 +332,7 @@ if user.id not in admin_ids: - **Trust feature**: `TrustedUser` table caches user_full_name + admin_full_name at trust time so `/trusted` lists admin info without Telegram API calls. Backfill script at `scripts/backfill_trusted_names.py` for pre-existing rows - **Local review artifacts**: `reviews/` directory contains output from parallel reviewer subagents. Gitignored; not part of the source tree - **Captcha DB ordering**: The captcha callback handler calls Telegram `unrestrict_user` BEFORE DB writes (remove_pending_captcha, start_new_user_probation). If unrestrict fails, the pending captcha stays in DB and the user can retry. DB finalization is idempotent — `remove_pending_captcha` returning False means a concurrent callback already finalized -- **Guest bot blocking**: Telegram Guest Mode lets any user `@mention` a bot and have the result posted in a chat. The `guest_bot_block` plugin (group=1) deletes non-whitelisted guest bot messages and progressively restricts the human caller (1st=warning, Nth=restrict). Admins/trusted users and channel-only callers are delete-only. Bot whitelist is case-insensitive with optional `@`. Guest strikes are tracked separately via `warning_kind="guest_bot"` and are not eligible for the DM self-service unrestriction flow +- **Guest bot blocking**: Telegram Guest Mode lets any user `@mention` a bot and have the result posted in a chat. The `guest_bot_block` plugin (group=0) deletes non-whitelisted guest bot messages and progressively restricts the human caller (1st=warning, Nth=restrict). Admins/trusted users and channel-only callers are delete-only. Bot whitelist is case-insensitive with optional `@`. A failed restriction does not increment the strike count past the threshold, so the caller's next guest message retries the restriction. Guest strikes are tracked separately via `warning_kind="guest_bot"` and are not eligible for the DM self-service unrestriction flow ## Policy diff --git a/README.md b/README.md index 8e84649..32e00c4 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ uv run mypy src/bot/ tests/ The project maintains comprehensive test coverage: - **Coverage**: 97%+ (~2,900 statements, <3% unreachable) -- **Tests**: 1,064 total (includes 19 Hypothesis property tests) +- **Tests**: 1,075 total (includes 19 Hypothesis property tests) - **Pass Rate**: 100% - **Property tests**: `tests/test_properties.py` exercises pure functions (format helpers, URL whitelist, name formatters) with random inputs and shrinks failing cases to minimal examples - **Mypy**: Pragmatic config in `pyproject.toml`. Disables error codes that are noisy from PTB / SQLModel / Pydantic v2; catches real type bugs in new code @@ -369,8 +369,9 @@ flowchart TD G_TopicIsBotAdmin -->|No| G_TopicDelete[Delete Message] G_TopicDelete --> StopTopic2([ApplicationHandlerStop]) - G_GuestGate{group=1 guest_bot_block:
Guest Mode message?} + G_GuestGate{group=0 guest_bot_block:
Guest Mode message?} G_GuestGate -->|No| G_InlineGate + G_GuestGate -->|Yes, whitelisted| G_InlineGate G_GuestGate -->|Yes, not whitelisted| G_GuestDelete[Delete Message] G_GuestDelete --> G_GuestCaller{Human caller?
Admin/Trusted?} G_GuestCaller -->|Admin/Trusted or Channel-only| StopGuest([ApplicationHandlerStop]) diff --git a/src/bot/handlers/guest_bot.py b/src/bot/handlers/guest_bot.py index 2d0e0ee..f8707bc 100644 --- a/src/bot/handlers/guest_bot.py +++ b/src/bot/handlers/guest_bot.py @@ -75,46 +75,43 @@ async def handle_guest_bot_message(update: Update, context: ContextTypes.DEFAULT if record.message_count >= group_config.warning_threshold: should_stop = False final_count = record.message_count - try: - async with restriction_lock(group_config.group_id, caller.id): - if db.is_user_restricted_by_bot(caller.id, group_config.group_id, warning_kind="guest_bot"): + async with restriction_lock(group_config.group_id, caller.id): + if db.is_user_restricted_by_bot(caller.id, group_config.group_id, warning_kind="guest_bot"): + should_stop = True + else: + fresh = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") + if fresh.message_count < group_config.warning_threshold: should_stop = True else: - fresh = db.get_or_create_user_warning(caller.id, group_config.group_id, warning_kind="guest_bot") - if fresh.message_count != group_config.warning_threshold: + try: + await context.bot.restrict_chat_member( + chat_id=group_config.group_id, + user_id=caller.id, + permissions=RESTRICTED_PERMISSIONS, + ) + db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") + final_count = fresh.message_count + except TelegramError as e: + logger.error("Failed to restrict guest bot caller %s: %s", caller.id, e, exc_info=True) + # Do not increment on failure: count stays pinned at + # threshold so the next guest message retries the + # restriction instead of drifting past it forever. should_stop = True - else: - try: - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=caller.id, - permissions=RESTRICTED_PERMISSIONS, - ) - db.mark_user_restricted(caller.id, group_config.group_id, warning_kind="guest_bot") - final_count = fresh.message_count - except TelegramError as e: - logger.error("Failed to restrict guest bot caller %s: %s", caller.id, e, exc_info=True) - should_stop = True - db.increment_message_count(caller.id, group_config.group_id, warning_kind="guest_bot") + if should_stop: + raise ApplicationHandlerStop + try: + await context.bot.send_message( + chat_id=group_config.group_id, + message_thread_id=group_config.warning_topic_id, + text=GUEST_BOT_RESTRICTION.format( + user_mention=user_mention, + message_count=final_count, + rules_link=group_config.rules_link, + ), + parse_mode="Markdown", + ) except TelegramError: - logger.error("Failed to restrict guest bot caller %s (lock level)", caller.id, exc_info=True) - should_stop = True - else: - if should_stop: - raise ApplicationHandlerStop - try: - await context.bot.send_message( - chat_id=group_config.group_id, - message_thread_id=group_config.warning_topic_id, - text=GUEST_BOT_RESTRICTION.format( - user_mention=user_mention, - message_count=final_count, - rules_link=group_config.rules_link, - ), - parse_mode="Markdown", - ) - except TelegramError: - logger.error("Failed to send guest bot restriction notice for user %s", caller.id, exc_info=True) + logger.error("Failed to send guest bot restriction notice for user %s", caller.id, exc_info=True) elif record.message_count == 1: try: await context.bot.send_message( diff --git a/src/bot/handlers/message.py b/src/bot/handlers/message.py index 5e5abd1..f1cc759 100644 --- a/src/bot/handlers/message.py +++ b/src/bot/handlers/message.py @@ -159,7 +159,11 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> try: async with restriction_lock(group_config.group_id, user.id): fresh = db.get_active_user_warning(user.id, group_config.group_id) - if fresh is None or fresh.is_restricted: + if ( + fresh is None + or fresh.id != record.id + or fresh.message_count < group_config.warning_threshold + ): logger.info( f"Skipping profile restriction for user {user.id} - " f"record no longer active (group_id={group_config.group_id})" @@ -182,13 +186,13 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> # Send restriction notice with DM link for appeal restriction_message = RESTRICTION_MESSAGE_AFTER_MESSAGES.format( user_mention=user_mention, - message_count=record.message_count, + message_count=fresh.message_count, missing_text=missing_text, rules_link=group_config.rules_link, dm_link=dm_link, ) logger.info( - f"Sending restriction notice: user_id={user.id}, user={user.full_name}, message_count={record.message_count}" + f"Sending restriction notice: user_id={user.id}, user={user.full_name}, message_count={fresh.message_count}" ) await context.bot.send_message( chat_id=group_config.group_id, @@ -197,7 +201,7 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> parse_mode="Markdown", ) logger.info( - f"Restricted user {user.id} ({user.full_name}) after {record.message_count} messages (group_id={group_config.group_id})" + f"Restricted user {user.id} ({user.full_name}) after {fresh.message_count} messages (group_id={group_config.group_id})" ) except Exception: logger.error( diff --git a/src/bot/handlers/verify.py b/src/bot/handlers/verify.py index f246954..524fde1 100644 --- a/src/bot/handlers/verify.py +++ b/src/bot/handlers/verify.py @@ -14,7 +14,7 @@ import logging from telegram import Bot, Update -from telegram.error import BadRequest, Forbidden, NetworkError, TimedOut +from telegram.error import TelegramError from telegram.ext import ContextTypes from bot.constants import ( @@ -90,7 +90,7 @@ async def verify_user_in_group( logger.info( f"Unrestricted user {target_user_id} in group {group_id} during verification" ) - except (BadRequest, Forbidden, NetworkError, TimedOut, RuntimeError) as e: + except (TelegramError, RuntimeError) as e: logger.info( f"Could not unrestrict user {target_user_id} in group {group_id}: {e}" ) diff --git a/src/bot/services/captcha_recovery.py b/src/bot/services/captcha_recovery.py index 9cb92f6..d47d48a 100644 --- a/src/bot/services/captcha_recovery.py +++ b/src/bot/services/captcha_recovery.py @@ -46,16 +46,20 @@ async def handle_captcha_expiration( user_full_name: The user's full name. """ db = get_database() - pending = db.get_pending_captcha(user_id, group_id) - if not pending: - logger.info(f"No pending captcha for user {user_id}, already verified") - return - - db.remove_pending_captcha(user_id, group_id) - # Create UserWarning to track this bot-applied restriction - # Allows DM handler to unrestrict user later when profile is complete async with restriction_lock(group_id, user_id): + pending = db.get_pending_captcha(user_id, group_id) + if not pending: + logger.info(f"No pending captcha for user {user_id}, already verified") + return + + removed = db.remove_pending_captcha(user_id, group_id) + if not removed: + logger.info(f"Captcha for user {user_id} already finalized, ignoring timeout") + return + + # Create UserWarning to track this bot-applied restriction + # Allows DM handler to unrestrict user later when profile is complete warning = db.get_or_create_user_warning(user_id, group_id) if not warning.is_restricted: db.mark_user_restricted(user_id, group_id) diff --git a/src/bot/services/restriction_lock.py b/src/bot/services/restriction_lock.py index 166c112..78caca9 100644 --- a/src/bot/services/restriction_lock.py +++ b/src/bot/services/restriction_lock.py @@ -9,24 +9,36 @@ This module provides :func:`restriction_lock`, an async context manager that serialises the Telegram API call + DB state transition for a given -``(group_id, user_id)`` pair. Locks are created on demand and stored in -a module-level dict; they persist for the process lifetime, which is -fine for a single-process bot. +``(group_id, user_id)`` pair. ``asyncio.Lock`` binds to the event loop of +its first *contended* acquire, so locks are scoped per running loop via a +:class:`weakref.WeakKeyDictionary` — this keeps a single-process bot's +locks alive for its one lifetime loop, while letting each test's fresh +event loop (see ``asyncio_default_fixture_loop_scope`` in pyproject.toml) +start with an empty lock table instead of reusing a lock bound to an +already-closed loop. """ import asyncio import contextlib +import weakref from collections.abc import AsyncIterator -_locks: dict[tuple[int, int], asyncio.Lock] = {} +_locks_by_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[tuple[int, int], asyncio.Lock]]" = ( + weakref.WeakKeyDictionary() +) def _get_lock(group_id: int, user_id: int) -> asyncio.Lock: + loop = asyncio.get_running_loop() + per_loop = _locks_by_loop.get(loop) + if per_loop is None: + per_loop = {} + _locks_by_loop[loop] = per_loop key = (group_id, user_id) - lock = _locks.get(key) + lock = per_loop.get(key) if lock is None: lock = asyncio.Lock() - _locks[key] = lock + per_loop[key] = lock return lock diff --git a/src/bot/services/scheduler.py b/src/bot/services/scheduler.py index e39e4bb..7ee5e81 100644 --- a/src/bot/services/scheduler.py +++ b/src/bot/services/scheduler.py @@ -113,7 +113,7 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> logger.info(f"Applying restriction to user_id={warning.user_id}") async with restriction_lock(group_config.group_id, warning.user_id): fresh = db.get_active_user_warning(warning.user_id, warning.group_id) - if fresh is None or fresh.is_restricted: + if fresh is None or fresh.id != warning.id: logger.info( f"Skipping auto-restriction for user {warning.user_id} - " f"record no longer active (group_id={group_config.group_id})" diff --git a/tests/test_restriction_lock.py b/tests/test_restriction_lock.py index f5c9a9d..a857bd9 100644 --- a/tests/test_restriction_lock.py +++ b/tests/test_restriction_lock.py @@ -2,7 +2,7 @@ import asyncio -from bot.services.restriction_lock import restriction_lock, _locks +from bot.services.restriction_lock import _locks_by_loop, restriction_lock class TestRestrictionLock: @@ -43,8 +43,9 @@ async def task(group_id: int, user_id: int, label: str) -> None: assert abs(a_idx - b_idx) <= 1 async def test_lock_is_reused(self): - """Same key returns the same lock object.""" + """Same key returns the same lock object, scoped to the running loop.""" key = (-300, 99) async with restriction_lock(*key): pass - assert key in _locks + loop = asyncio.get_running_loop() + assert key in _locks_by_loop[loop] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 2ca55f8..8e2daf3 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -151,7 +151,9 @@ async def test_restricts_multiple_expired_warnings(self, mock_registry): mock_db = MagicMock() mock_db.get_warnings_past_time_threshold_for_group.return_value = mock_warnings - mock_db.get_active_user_warning.return_value = mock_warnings[0] + mock_db.get_active_user_warning.side_effect = lambda user_id, group_id, warning_kind="profile": next( + (w for w in mock_warnings if w.user_id == user_id), None + ) mock_db.mark_user_restricted = MagicMock() mock_bot = AsyncMock()