diff --git a/.gitignore b/.gitignore index b4aa9a5..6bcfdfc 100644 --- a/.gitignore +++ b/.gitignore @@ -14,9 +14,11 @@ data/ # Agent/planning docs docs/ +plans/ .pi/ # Local review artifacts (parallel-reviewer outputs) reviews/ +.pi-subagents/ .tokensave/ diff --git a/src/bot/database/service.py b/src/bot/database/service.py index 32b2518..e2588c2 100644 --- a/src/bot/database/service.py +++ b/src/bot/database/service.py @@ -634,6 +634,17 @@ def get_all_pending_captchas(self) -> list[PendingCaptchaValidation]: statement = select(PendingCaptchaValidation) return list(session.exec(statement).all()) + def get_all_new_user_probations(self) -> list[NewUserProbation]: + """ + Get all new-user probation records. + + Returns: + list[NewUserProbation]: All probation records. + """ + with Session(self._engine) as session: + statement = select(NewUserProbation) + return list(session.exec(statement).all()) + def start_new_user_probation(self, user_id: int, group_id: int) -> NewUserProbation: """ Start or refresh probation for a new user. diff --git a/src/bot/handlers/anti_spam.py b/src/bot/handlers/anti_spam.py index 01de11a..ffb6044 100644 --- a/src/bot/handlers/anti_spam.py +++ b/src/bot/handlers/anti_spam.py @@ -348,6 +348,18 @@ async def handle_inline_keyboard_spam( ) +def _should_skip_new_user_spam_check(update, context, group_config) -> bool: + """Check if new user spam handler should skip this message.""" + if group_config is None: + return True + user = update.message.from_user + if user.is_bot: + return True + if is_user_admin_or_trusted(context, group_config.group_id, user.id): + return True + return False + + async def handle_new_user_spam( update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -371,15 +383,7 @@ async def handle_new_user_spam( group_config = get_group_config_for_update(update) user = update.message.from_user - # Only process messages from monitored groups - if group_config is None: - return - - # Ignore bots - if user.is_bot: - return - - if is_user_admin_or_trusted(context, group_config.group_id, user.id): + if _should_skip_new_user_spam_check(update, context, group_config): return db = get_database() diff --git a/src/bot/handlers/bio_bait.py b/src/bot/handlers/bio_bait.py index 8aee4bc..848d583 100644 --- a/src/bot/handlers/bio_bait.py +++ b/src/bot/handlers/bio_bait.py @@ -37,7 +37,13 @@ WHITELISTED_TELEGRAM_PATHS, ) from bot.group_config import get_group_config_for_update -from bot.services.telegram_utils import get_user_mention, is_user_admin_or_trusted, is_url_whitelisted +from bot.services.telegram_utils import ( + get_user_mention, + is_user_admin_or_trusted, + is_url_whitelisted, + restrict_chat_member_with_retry, + send_message_with_retry, +) # Filter for bio-bait handler registration in main.py. # Must NOT restrict to TEXT|CAPTION so non-text messages (e.g. photos @@ -269,7 +275,14 @@ async def send_monitor_alert_to_owner( try: for chunk in _chunk_telegram_text(alert_text): - await context.bot.send_message(chat_id=alert_chat_id, text=chunk) + ok = await send_message_with_retry( + context.bot, chat_id=alert_chat_id, text=chunk + ) + if not ok: + logger.error( + f"Failed to send bio bait monitor alert chunk: user_id={user_id}, group_id={group_id}" + ) + return False return True except Exception: logger.error(f"Failed to send bio bait monitor alert: user_id={user_id}, group_id={group_id}") @@ -312,6 +325,64 @@ async def get_cached_user_bio( cache[user_id] = (now + USER_BIO_FAILURE_CACHE_TTL_SECONDS, _BIO_CACHE_FAILURE) return None + +async def _enforce_bio_bait_restriction( + update: Update, + context: ContextTypes.DEFAULT_TYPE, + group_config, + user, + detection_reason: str, +) -> None: + """Delete, restrict, notify for confirmed bio bait spam. Caller raises ApplicationHandlerStop.""" + user_mention = get_user_mention(user) + + try: + await update.message.delete() + logger.info(f"Deleted bio bait spam from user_id={user.id}") + except Exception: + logger.error(f"Failed to delete bio bait spam: user_id={user.id}", exc_info=True) + + restricted = False + try: + await restrict_chat_member_with_retry( + context.bot, + chat_id=group_config.group_id, + user_id=user.id, + permissions=RESTRICTED_PERMISSIONS, + ) + restricted = True + clear_cached_user_bio(context, user.id) + logger.info(f"Restricted user_id={user.id} for bio bait spam") + except Exception: + logger.error(f"Failed to restrict user for bio bait spam: user_id={user.id}", exc_info=True) + + try: + if detection_reason == "bio_links": + template = ( + BIO_LINK_SPAM_NOTIFICATION if restricted + else BIO_LINK_SPAM_NOTIFICATION_NO_RESTRICT + ) + else: + template = ( + BIO_BAIT_SPAM_NOTIFICATION if restricted + else BIO_BAIT_SPAM_NOTIFICATION_NO_RESTRICT + ) + notification_text = template.format( + user_mention=user_mention, + rules_link=group_config.rules_link, + ) + await send_message_with_retry( + context.bot, + chat_id=group_config.group_id, + message_thread_id=group_config.warning_topic_id, + text=notification_text, + parse_mode="Markdown", + ) + logger.info(f"Sent bio bait spam notification for user_id={user.id}") + except Exception: + logger.error(f"Failed to send bio bait spam notification: user_id={user.id}", exc_info=True) + + async def handle_bio_bait_spam( update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -392,50 +463,5 @@ async def handle_bio_bait_spam( ) return - user_mention = get_user_mention(user) - - try: - await update.message.delete() - logger.info(f"Deleted bio bait spam from user_id={user.id}") - except Exception: - logger.error(f"Failed to delete bio bait spam: user_id={user.id}", exc_info=True) - - restricted = False - try: - await context.bot.restrict_chat_member( - chat_id=group_config.group_id, - user_id=user.id, - permissions=RESTRICTED_PERMISSIONS, - ) - restricted = True - clear_cached_user_bio(context, user.id) - logger.info(f"Restricted user_id={user.id} for bio bait spam") - except Exception: - logger.error(f"Failed to restrict user for bio bait spam: user_id={user.id}", exc_info=True) - - try: - if detection_reason == "bio_links": - template = ( - BIO_LINK_SPAM_NOTIFICATION if restricted - else BIO_LINK_SPAM_NOTIFICATION_NO_RESTRICT - ) - else: - template = ( - BIO_BAIT_SPAM_NOTIFICATION if restricted - else BIO_BAIT_SPAM_NOTIFICATION_NO_RESTRICT - ) - notification_text = template.format( - user_mention=user_mention, - rules_link=group_config.rules_link, - ) - await context.bot.send_message( - chat_id=group_config.group_id, - message_thread_id=group_config.warning_topic_id, - text=notification_text, - parse_mode="Markdown", - ) - logger.info(f"Sent bio bait spam notification for user_id={user.id}") - except Exception: - logger.error(f"Failed to send bio bait spam notification: user_id={user.id}", exc_info=True) - + await _enforce_bio_bait_restriction(update, context, group_config, user, detection_reason) raise ApplicationHandlerStop diff --git a/src/bot/handlers/check.py b/src/bot/handlers/check.py index 901289e..0fa3dda 100644 --- a/src/bot/handlers/check.py +++ b/src/bot/handlers/check.py @@ -30,6 +30,7 @@ get_user_mention, get_user_mention_by_id, require_admin_dm_target, + send_message_with_retry, ) from bot.services.user_checker import check_user_profile @@ -205,6 +206,17 @@ async def handle_check_forwarded_message( logger.error(f"Error checking forwarded user {user_id}: {e}", exc_info=True) +def _parse_warn_callback_data(data: str) -> tuple[int, str] | None: + """Parse warn callback data (warn::). Returns (user_id, missing_code) or None.""" + try: + parts = data.split(":") + user_id = int(parts[1]) + missing_code = parts[2] if len(parts) > 2 else "" + return (user_id, missing_code) + except (IndexError, ValueError): + return None + + async def handle_warn_callback( update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -231,14 +243,12 @@ async def handle_warn_callback( return # Parse callback data: warn:: - try: - parts = query.data.split(":") - target_user_id = int(parts[1]) - missing_code = parts[2] if len(parts) > 2 else "" - except (IndexError, ValueError): + parsed = _parse_warn_callback_data(query.data) + if parsed is None: await query.edit_message_text("❌ Data callback tidak valid.") logger.error(f"Invalid callback_data format: {query.data}") return + target_user_id, missing_code = parsed # Build missing items text missing_items = [] @@ -264,16 +274,18 @@ async def handle_warn_callback( rules_link=group_config.rules_link, ) try: - await context.bot.send_message( + ok = await send_message_with_retry( + context.bot, chat_id=group_config.group_id, message_thread_id=group_config.warning_topic_id, text=warn_message, parse_mode="Markdown", ) - sent_to_any = True - logger.info( - f"Admin {admin_user_id} sent warning to user {target_user_id} in group {group_config.group_id}" - ) + if ok: + sent_to_any = True + logger.info( + f"Admin {admin_user_id} sent warning to user {target_user_id} in group {group_config.group_id}" + ) except Exception as e: logger.error(f"Failed to send warning to group {group_config.group_id}: {e}") diff --git a/src/bot/handlers/dm.py b/src/bot/handlers/dm.py index c7df0fd..38e4584 100644 --- a/src/bot/handlers/dm.py +++ b/src/bot/handlers/dm.py @@ -29,12 +29,64 @@ ) from bot.database.service import get_database from bot.group_config import get_group_registry -from bot.services.telegram_utils import get_user_mention, get_user_status, unrestrict_user +from bot.services.telegram_utils import ( + get_user_mention, + get_user_status, + send_message_with_retry, + unrestrict_user, +) from bot.services.user_checker import check_user_profile logger = logging.getLogger(__name__) + +async def _unrestrict_in_groups( + context: ContextTypes.DEFAULT_TYPE, + user, + restricted_groups: list, +) -> int: + """Unrestrict user in all groups where restricted. Returns success count.""" + db = get_database() + success_count = 0 + + for gc, user_status in restricted_groups: + if user_status != ChatMemberStatus.RESTRICTED: + db.mark_user_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}") + try: + await unrestrict_user(context.bot, gc.group_id, user.id) + db.mark_user_unrestricted(user.id, gc.group_id) + success_count += 1 + + user_mention = get_user_mention(user) + notification_message = DM_UNRESTRICTION_NOTIFICATION.format( + user_mention=user_mention + ) + await send_message_with_retry( + context.bot, + chat_id=gc.group_id, + message_thread_id=gc.warning_topic_id, + text=notification_message, + parse_mode="Markdown", + ) + logger.info( + f"Unrestricted user {user.id} ({user.full_name}) via DM (group_id={gc.group_id})" + ) + except Exception: + logger.error( + f"Failed to unrestrict user {user.id} ({user.full_name}) via DM (group_id={gc.group_id})", + exc_info=True, + ) + + return success_count + + async def handle_dm(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """ Handle direct messages to the bot for unrestriction flow. @@ -133,51 +185,12 @@ async def handle_dm(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: return # Unrestrict user from all groups where restricted by bot - unrestricted_any = False - all_already_unrestricted = True - - for gc, user_status in restricted_groups: - # User was restricted by bot but is no longer restricted on Telegram - # (e.g., admin already unrestricted them) - just clear our record - if user_status != ChatMemberStatus.RESTRICTED: - db.mark_user_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 - - all_already_unrestricted = False - - # Remove restriction - 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) - unrestricted_any = True - - # Send notification to warning topic - user_mention = get_user_mention(user) - notification_message = DM_UNRESTRICTION_NOTIFICATION.format( - user_mention=user_mention - ) - await context.bot.send_message( - chat_id=gc.group_id, - message_thread_id=gc.warning_topic_id, - text=notification_message, - parse_mode="Markdown", - ) - logger.info( - f"Unrestricted user {user.id} ({user.full_name}) via DM (group_id={gc.group_id})" - ) - except Exception: - logger.error( - f"Failed to unrestrict user {user.id} ({user.full_name}) via DM (group_id={gc.group_id})", - exc_info=True, - ) + had_any_restricted = any(status == ChatMemberStatus.RESTRICTED for _, status in restricted_groups) + success_count = await _unrestrict_in_groups(context, user, restricted_groups) - if unrestricted_any: + if success_count > 0: await update.message.reply_text(DM_UNRESTRICTION_SUCCESS_MESSAGE) - elif all_already_unrestricted: + elif not had_any_restricted: await update.message.reply_text(DM_ALREADY_UNRESTRICTED_MESSAGE) else: # All unrestriction attempts failed diff --git a/src/bot/handlers/duplicate_spam.py b/src/bot/handlers/duplicate_spam.py index bfbd124..2ce66c5 100644 --- a/src/bot/handlers/duplicate_spam.py +++ b/src/bot/handlers/duplicate_spam.py @@ -33,6 +33,8 @@ logger = logging.getLogger(__name__) RECENT_MESSAGES_KEY = "duplicate_spam_recent" +RECENT_MESSAGES_MAX_SIZE = 2000 +_RECENT_LAST_TOUCH_KEY = "duplicate_spam_recent_last_touch" @dataclass @@ -134,6 +136,18 @@ async def handle_duplicate_spam( ) dq.append(current_message) + # Track last-touch timestamp for cache eviction + last_touch = context.bot_data.setdefault(_RECENT_LAST_TOUCH_KEY, {}) + last_touch[(group_config.group_id, user.id)] = now + + # Evict oldest entries when cache exceeds the cap + recent = context.bot_data[RECENT_MESSAGES_KEY] + if len(recent) >= RECENT_MESSAGES_MAX_SIZE: + sorted_keys = sorted(last_touch, key=lambda k: last_touch[k]) + for k in sorted_keys[:RECENT_MESSAGES_MAX_SIZE // 2]: + del recent[k] + del last_touch[k] + if len(similar_messages) < group_config.duplicate_spam_threshold - 1: return diff --git a/src/bot/handlers/message.py b/src/bot/handlers/message.py index 5200cd4..9cb9a2f 100644 --- a/src/bot/handlers/message.py +++ b/src/bot/handlers/message.py @@ -63,7 +63,6 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> f"Handler called: user_id={user.id}, user={user.full_name}, chat_id={update.effective_chat.id}" ) - # Ignore messages from bots if user.is_bot: logger.info(f"Skipping message: user is bot (user_id={user.id})") return diff --git a/src/bot/handlers/status.py b/src/bot/handlers/status.py new file mode 100644 index 0000000..196bf8c --- /dev/null +++ b/src/bot/handlers/status.py @@ -0,0 +1,165 @@ +""" +Status command handler for the PythonID bot. + +Provides a DM-only, admin-only ``/status`` command that shows bot +operational state: uptime, per-group config summary, probation and +captcha queue lengths, database file size, and last job timestamps. +""" + +from __future__ import annotations + +import logging +import os +import time + +from telegram import Update +from telegram.ext import CommandHandler, ContextTypes +from telegram.helpers import escape_markdown + +from bot.config import get_settings +from bot.database.service import get_database +from bot.group_config import get_group_registry + +logger = logging.getLogger(__name__) + + +def _format_uptime(seconds: float) -> str: + """Format monotonic seconds into Xd Yh Zm.""" + days, rem = divmod(int(seconds), 86400) + hours, rem = divmod(rem, 3600) + minutes = rem // 60 + parts = [] + if days: + parts.append(f"{days}d") + if hours: + parts.append(f"{hours}h") + parts.append(f"{minutes}m") + return " ".join(parts) + + +def _format_filesize(path: str) -> str: + """Return file size as KB or MB.""" + try: + size = os.path.getsize(path) + except FileNotFoundError: + return "N/A" + if size >= 1024 * 1024: + return f"{size / (1024 * 1024):.1f} MB" + return f"{size / 1024:.1f} KB" + + +async def _check_status_prereqs( + update: Update, context: ContextTypes.DEFAULT_TYPE +) -> bool: + """Validate /status prerequisites: message exists, private chat, admin. + + Returns True if all checks pass. Sends error reply and returns False + on any failure. + """ + if not update.message or not update.message.from_user: + logger.warning("handle_status called without message or sender") + return False + + if update.effective_chat and update.effective_chat.type != "private": + await update.message.reply_text( + "❌ Perintah ini hanya bisa digunakan di chat pribadi dengan bot." + ) + return False + + admin_user_id = update.message.from_user.id + admin_ids = context.bot_data.get("admin_ids", []) + if admin_user_id not in admin_ids: + await update.message.reply_text( + "❌ Kamu tidak memiliki izin untuk menggunakan perintah ini." + ) + logger.warning( + f"Non-admin user {admin_user_id} ({update.message.from_user.full_name}) " + "attempted to use /status command" + ) + return False + + return True + + +async def handle_status( + update: Update, context: ContextTypes.DEFAULT_TYPE +) -> None: + """Handle /status command in bot DM — show operational state.""" + if not await _check_status_prereqs(update, context): + return + + lines: list[str] = [] + + # --- Uptime --- + start = context.bot_data.get("start_time") + if start is not None: + uptime = _format_uptime(time.monotonic() - start) + lines.append(f"*Uptime:* {uptime}") + else: + lines.append("*Uptime:* N/A") + + # --- Per-group summary --- + lines.append("") + lines.append("*Groups:*") + registry = get_group_registry() + effective_map = context.bot_data.get("plugin_effective_map", {}) + for gc in registry.all_groups(): + gid = gc.group_id + captcha = "CAPTCHA" if gc.captcha_enabled else "" + group_line = f" • `{escape_markdown(str(gid), version=1)}`" + if captcha: + group_line += f" _{captcha}_" + toggles = effective_map.get(gid, {}) + disabled = [k for k, v in toggles.items() if not v] + if disabled: + disabled_str = ", ".join(sorted(disabled)) + group_line += ( + f"\n plugins off: {escape_markdown(disabled_str, version=1)}" + ) + lines.append(group_line) + + # --- Probation count --- + lines.append("") + db = get_database() + probation_records = db.get_all_new_user_probations() + lines.append(f"*Probation:* {len(probation_records)} user(s)") + + # --- Pending captcha count --- + pending = db.get_all_pending_captchas() + lines.append(f"*Captcha:* {len(pending)} pending") + + # --- Database size --- + db_path = get_settings().database_path + size_str = _format_filesize(db_path) + lines.append(f"*Database:* {size_str}") + + # --- Last jobs --- + lines.append("") + lines.append("*Last jobs:*") + refresh_ts = context.bot_data.get("last_admin_refresh") + if refresh_ts is not None: + lines.append( + " • admin refresh: " + f"{time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime(refresh_ts))}" + ) + else: + lines.append(" • admin refresh: never") + + restrict_ts = context.bot_data.get("last_auto_restrict") + if restrict_ts is not None: + lines.append( + " • auto restrict: " + f"{time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime(restrict_ts))}" + ) + else: + lines.append(" • auto restrict: never") + + await update.message.reply_text( + "\n".join(lines), + parse_mode="Markdown", + ) + + +def get_handlers() -> list[CommandHandler]: + """Return list of handlers for the status command.""" + return [CommandHandler("status", handle_status)] diff --git a/src/bot/handlers/verify.py b/src/bot/handlers/verify.py index 29e6da7..1465195 100644 --- a/src/bot/handlers/verify.py +++ b/src/bot/handlers/verify.py @@ -15,7 +15,12 @@ from bot.constants import VERIFICATION_CLEARANCE_MESSAGE from bot.database.service import DatabaseService, get_database from bot.group_config import GroupRegistry, get_group_registry -from bot.services.telegram_utils import get_user_mention, require_admin_dm_target, unrestrict_user +from bot.services.telegram_utils import ( + get_user_mention, + require_admin_dm_target, + send_message_with_retry, + unrestrict_user, +) logger = logging.getLogger(__name__) @@ -56,8 +61,9 @@ async def verify_user( try: await unrestrict_user(bot, group_config.group_id, target_user_id) logger.info(f"Unrestricted user {target_user_id} in group {group_config.group_id} during verification") - except BadRequest as e: - # User might not be restricted or not in group - that's okay + except (BadRequest, RuntimeError) as e: + # BadRequest: user might not be restricted or not in group - okay + # RuntimeError: flood control retries exhausted logger.info(f"Could not unrestrict user {target_user_id} in group {group_config.group_id}: {e}") # Delete all warning records for this user in this group @@ -74,11 +80,12 @@ async def verify_user( clearance_message = VERIFICATION_CLEARANCE_MESSAGE.format( user_mention=user_mention ) - await bot.send_message( + await send_message_with_retry( + bot, chat_id=group_config.group_id, message_thread_id=group_config.warning_topic_id, text=clearance_message, - parse_mode="Markdown" + parse_mode="Markdown", ) logger.info(f"Sent clearance notification to warning topic for user {target_user_id} in group {group_config.group_id}") except Exception: diff --git a/src/bot/main.py b/src/bot/main.py index cf5ff9c..e58d4f0 100644 --- a/src/bot/main.py +++ b/src/bot/main.py @@ -6,6 +6,7 @@ """ import logging +import time from typing import Literal import logfire @@ -122,6 +123,9 @@ async def post_init(application: Application) -> None: # type: ignore[type-arg] # Use preload_admin_ids which preserves cache on failures await preload_admin_ids(application) + # Record start time for /status uptime + application.bot_data["start_time"] = time.monotonic() + # Preload trusted users cache db = get_database() trusted_ids = db.get_trusted_user_ids() diff --git a/src/bot/plugins/builtin/status.py b/src/bot/plugins/builtin/status.py new file mode 100644 index 0000000..ea158ca --- /dev/null +++ b/src/bot/plugins/builtin/status.py @@ -0,0 +1,36 @@ +"""Built-in plugin: status. + +Wraps ``bot.handlers.status`` for the ``/status`` DM-admin command. +No ``guard_plugin`` wrap — /status is admin-only by handler checks, not +per-group gated. +""" + +from __future__ import annotations + +import copy +import logging +from typing import TYPE_CHECKING + +from bot.handlers import status + +if TYPE_CHECKING: + from telegram.ext import Application, BaseHandler + +logger = logging.getLogger(__name__) + + +def register_status(application: Application) -> list[BaseHandler]: # type: ignore[type-arg] + """Register /status command handler. + + Each handler is CLONED before adding to avoid mutating the original + handler objects returned by ``get_handlers()``. + """ + handlers = status.get_handlers() + registered = [] + for h in handlers: + cloned = copy.copy(h) + # No guard_plugin wrap — /status is admin-gated by the handler itself + application.add_handler(cloned) + registered.append(cloned) + logger.info("Registered handler: status (group=0)") + return registered diff --git a/src/bot/plugins/definitions.py b/src/bot/plugins/definitions.py index aa24825..055af0b 100644 --- a/src/bot/plugins/definitions.py +++ b/src/bot/plugins/definitions.py @@ -31,6 +31,7 @@ {"name": "untrust_callback", "handler_group": 0, "description": "Admin untrust button callback"}, {"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": "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/src/bot/plugins/manager.py b/src/bot/plugins/manager.py index fc7ebc9..9c35e7a 100644 --- a/src/bot/plugins/manager.py +++ b/src/bot/plugins/manager.py @@ -30,6 +30,7 @@ from bot.plugins.builtin import jobs as jobs_mod from bot.plugins.builtin import profile_monitor as pm_mod from bot.plugins.builtin import spam as spam_mod +from bot.plugins.builtin import status as status_mod from bot.plugins.builtin import topic_guard as tg_mod from bot.plugins.config import resolve_plugin_toggles from bot.plugins.definitions import MANIFEST_ORDER, get_plugin_definitions @@ -65,6 +66,8 @@ "captcha": captcha_mod.register_captcha, # dm "dm": dm_mod.register_dm, + # status + "status": status_mod.register_status, # spam "inline_keyboard_spam": spam_mod.register_inline_keyboard_spam, "bio_bait_spam": spam_mod.register_bio_bait_spam, diff --git a/src/bot/services/admin_cache.py b/src/bot/services/admin_cache.py index 287cb42..7fc7bee 100644 --- a/src/bot/services/admin_cache.py +++ b/src/bot/services/admin_cache.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +import time from typing import TYPE_CHECKING from bot.group_config import get_group_registry @@ -60,6 +61,7 @@ async def refresh_admin_ids(context: ContextTypes.DEFAULT_TYPE) -> None: group_admin_ids = context.bot_data.get("group_admin_ids", {}) all_admin_ids = context.bot_data.get("admin_ids", []) logger.info(f"Refreshed admin IDs: {len(all_admin_ids)} unique admin(s) across {len(group_admin_ids)} group(s)") + context.bot_data["last_admin_refresh"] = time.time() async def preload_admin_ids(context: ContextTypes.DEFAULT_TYPE) -> None: diff --git a/src/bot/services/scheduler.py b/src/bot/services/scheduler.py index 9f754ab..a240e25 100644 --- a/src/bot/services/scheduler.py +++ b/src/bot/services/scheduler.py @@ -7,6 +7,7 @@ """ import logging +import time from telegram.constants import ChatMemberStatus from telegram.ext import ContextTypes @@ -20,7 +21,12 @@ 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.telegram_utils import get_user_mention, get_user_status +from bot.services.telegram_utils import ( + get_user_mention, + get_user_status, + restrict_chat_member_with_retry, + send_message_with_retry, +) logger = logging.getLogger(__name__) @@ -80,11 +86,17 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> logger.info(f"Applying restriction to user_id={warning.user_id}") # Apply restriction (even if user left, they'll be restricted when they rejoin) - await bot.restrict_chat_member( + 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" + ) + continue db.mark_user_restricted(warning.user_id, group_config.group_id) # Get user info for proper mention @@ -109,7 +121,8 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> rules_link=group_config.rules_link, dm_link=dm_link, ) - await bot.send_message( + await send_message_with_retry( + bot, chat_id=group_config.group_id, message_thread_id=group_config.warning_topic_id, text=restriction_message, @@ -123,3 +136,4 @@ async def auto_restrict_expired_warnings(context: ContextTypes.DEFAULT_TYPE) -> logger.error( f"Error auto-restricting user {warning.user_id} in group {group_config.group_id}: {e}", exc_info=True ) + context.bot_data["last_auto_restrict"] = time.time() diff --git a/src/bot/services/telegram_utils.py b/src/bot/services/telegram_utils.py index 99e9387..a59dde6 100644 --- a/src/bot/services/telegram_utils.py +++ b/src/bot/services/telegram_utils.py @@ -5,12 +5,14 @@ Telegram's API across different handlers and services. """ +import asyncio import logging +from datetime import timedelta from urllib.parse import urlparse from telegram import Bot, Chat, Message, Update, User from telegram.constants import ChatMemberStatus -from telegram.error import BadRequest, Forbidden +from telegram.error import BadRequest, Forbidden, RetryAfter from telegram.ext import ContextTypes from telegram.helpers import escape_markdown, mention_markdown @@ -99,15 +101,15 @@ async def unrestrict_user( ) -> None: """ Remove restrictions from a user by applying group's default permissions. - + This restores the user to normal member status in the group. Does NOT update the database - caller must handle that separately. - + Args: bot: Telegram bot instance. group_id: Telegram group ID. user_id: Telegram user ID to unrestrict. - + Raises: BadRequest: If user not found or bot lacks permissions. """ @@ -116,13 +118,16 @@ async def unrestrict_user( # Get group's default permissions chat = await bot.get_chat(group_id) default_permissions = chat.permissions - + # Apply default permissions to remove restrictions - await bot.restrict_chat_member( + ok = await restrict_chat_member_with_retry( + bot, chat_id=group_id, user_id=user_id, permissions=default_permissions, ) + if not ok: + raise RuntimeError("Final RetryAfter exceeded on restrict_chat_member") except Exception as e: logger.error( f"Failed to unrestrict user_id={user_id} in group_id={group_id}: {e}", @@ -252,6 +257,100 @@ def is_user_admin_or_trusted(context: object, group_id: int, user_id: int) -> bo trusted_ids = _get_trusted_ids(bot_data) return user_id in trusted_ids +def _retry_after_seconds(e: RetryAfter) -> float: + """Extract RetryAfter.retry_after as seconds (handles int and timedelta).""" + return e.retry_after.total_seconds() if isinstance(e.retry_after, timedelta) else e.retry_after + + +_MAX_RETRY_SLEEP_SECONDS = 30.0 + + +def _clamped_retry_seconds(e: RetryAfter) -> float: + """RetryAfter sleep capped at ``_MAX_RETRY_SLEEP_SECONDS`` so one bad + flood-control response can't stall a per-group loop for a minute-plus.""" + return min(_retry_after_seconds(e) + 1, _MAX_RETRY_SLEEP_SECONDS) + + +async def send_message_with_retry(bot: Bot, *, chat_id: int, **kwargs: object) -> bool: + """ + Send a message with one retry on RetryAfter (HTTP 429 / flood control). + + Catches ``telegram.error.RetryAfter``, sleeps ``e.retry_after + 1`` seconds, + and retries exactly once. On a second RetryAfter it returns ``False`` + (the error is logged). All other exceptions re-raise so the caller's + existing ``except Exception`` still catches them. + + Args: + bot: Telegram Bot instance. + chat_id: Target chat / group ID. + **kwargs: Extra keyword arguments forwarded to ``bot.send_message``. + + Returns: + bool: ``True`` if the message was sent successfully, ``False`` if a + second consecutive RetryAfter was encountered. + """ + try: + await bot.send_message(chat_id=chat_id, **kwargs) + return True + except RetryAfter as e: + wait_seconds = int(_retry_after_seconds(e)) + logger.warning( + f"RetryAfter on send_message to chat {chat_id}, sleeping {wait_seconds}s before retry" + ) + await asyncio.sleep(_clamped_retry_seconds(e)) + try: + await bot.send_message(chat_id=chat_id, **kwargs) + return True + except RetryAfter: + logger.error( + f"RetryAfter again on send_message to chat {chat_id}, giving up" + ) + return False + + +async def restrict_chat_member_with_retry( + bot: Bot, *, chat_id: int, user_id: int, permissions: object, **kwargs: object +) -> bool: + """ + Restrict a chat member with one retry on RetryAfter. + + Same retry strategy as :func:`send_message_with_retry` but wraps + ``bot.restrict_chat_member``. Returns ``True`` on success, ``False`` after a + second consecutive RetryAfter. Other exceptions re-raise. + + Args: + bot: Telegram Bot instance. + chat_id: Group ID. + user_id: User ID to restrict. + permissions: ``ChatPermissions`` to apply. + **kwargs: Extra keyword arguments forwarded to ``bot.restrict_chat_member``. + + Returns: + bool: ``True`` if restriction applied, ``False`` after second RetryAfter. + """ + try: + await bot.restrict_chat_member( + chat_id=chat_id, user_id=user_id, permissions=permissions, **kwargs + ) + return True + except RetryAfter as e: + wait_seconds = int(_retry_after_seconds(e)) + logger.warning( + f"RetryAfter on restrict_chat_member to chat {chat_id} (user {user_id}), sleeping {wait_seconds}s" + ) + await asyncio.sleep(_clamped_retry_seconds(e)) + try: + await bot.restrict_chat_member( + chat_id=chat_id, user_id=user_id, permissions=permissions, **kwargs + ) + return True + except RetryAfter: + logger.error( + f"RetryAfter again on restrict_chat_member to chat {chat_id} (user {user_id}), giving up" + ) + return False + + async def fetch_group_admin_ids(bot: Bot, group_id: int) -> list[int]: """ Fetch all administrator user IDs from a group. diff --git a/tests/test_bio_bait.py b/tests/test_bio_bait.py index 065eb87..e58f25c 100644 --- a/tests/test_bio_bait.py +++ b/tests/test_bio_bait.py @@ -18,6 +18,7 @@ has_suspicious_bio_links, is_bio_bait_spam, normalize_bio_bait_text, + send_monitor_alert_to_owner, ) class TestNormalizeBioBaitText: @@ -706,3 +707,53 @@ async def test_failure_cache_expires(self): result = await get_cached_user_bio(mock_context, user_id=123) assert result == "new bio" mock_context.bot.get_chat.assert_called_once() + + +class TestSendMonitorAlertToOwner: + """Tests for send_monitor_alert_to_owner error handling.""" + + @pytest.fixture + def mock_context(self): + context = MagicMock() + context.bot = MagicMock() + context.bot.send_message = AsyncMock() + return context + + async def test_propagates_send_error_to_except(self, mock_context): + """Non-RetryAfter error re-raises; caught by outer except Exception, returns False.""" + from telegram.error import BadRequest + mock_context.bot.send_message = AsyncMock(side_effect=BadRequest("Test error")) + + result = await send_monitor_alert_to_owner( + context=mock_context, + alert_chat_id=57747812, + group_id=-100, + user_id=42, + user_name="Test", + username="test", + detection_reason="message_bait", + message_text="cek bio", + profile_bio=None, + ) + assert result is False + + async def test_retry_after_failure_returns_false(self, mock_context): + """Second RetryAfter in send_message_with_retry returns False; logs and returns False.""" + from telegram.error import RetryAfter + mock_context.bot.send_message = AsyncMock( + side_effect=RetryAfter(retry_after=1) + ) + + with patch("bot.services.telegram_utils.asyncio.sleep"): + result = await send_monitor_alert_to_owner( + context=mock_context, + alert_chat_id=57747812, + group_id=-100, + user_id=42, + user_name="Test", + username="test", + detection_reason="message_bait", + message_text="cek bio", + profile_bio=None, + ) + assert result is False diff --git a/tests/test_duplicate_spam.py b/tests/test_duplicate_spam.py index 9686e3e..84ff645 100644 --- a/tests/test_duplicate_spam.py +++ b/tests/test_duplicate_spam.py @@ -408,3 +408,121 @@ async def test_does_not_retry_already_attempted_deletes( mock_context.bot.delete_message.assert_called_once_with( chat_id=-100, message_id=100 ) + + +class TestRecentMessagesCacheEviction: + """Tests for the in-memory cache eviction logic.""" + + @pytest.fixture + def group_config(self): + return GroupConfig( + group_id=-100, + warning_topic_id=0, + duplicate_spam_enabled=True, + duplicate_spam_min_length=20, + duplicate_spam_threshold=2, + duplicate_spam_window_seconds=120, + duplicate_spam_similarity=0.95, + ) + + @pytest.fixture + def mock_update(self): + update = MagicMock() + update.message.from_user.id = 42 + update.message.from_user.is_bot = False + update.message.text = "Barangkali di sini ada yang sedang mencari kerja" + update.message.effective_chat.id = -100 + update.message.message_id = 100 + update.effective_chat.id = -100 + update.effective_user.id = 42 + update.effective_message = update.message + return update + + @pytest.fixture + def mock_context(self): + context = MagicMock() + context.bot_data = {"group_admin_ids": {-100: [1, 2]}} + context.bot.send_message = AsyncMock() + return context + + async def test_eviction_bounds_outer_dict( + self, mock_update, mock_context, group_config + ): + """Outer dict shrinks and stays bounded when it exceeds the cap.""" + now = datetime.now(UTC) + max_size = 10 + overage = 5 + + recent_dict: dict[tuple[int, int], deque[RecentMessage]] = {} + last_touch_dict: dict[tuple[int, int], datetime] = {} + for i in range(max_size + overage): + gid = -(1000 + i) + uid = i + recent_dict[(gid, uid)] = deque() + last_touch_dict[(gid, uid)] = now - timedelta(hours=i * 24) + + mock_context.bot_data[RECENT_MESSAGES_KEY] = recent_dict + mock_context.bot_data["duplicate_spam_recent_last_touch"] = last_touch_dict + + mock_update.message.from_user.id = 9999 + + with ( + patch("bot.handlers.duplicate_spam.get_group_config_for_update", return_value=group_config), + patch("bot.handlers.duplicate_spam.RECENT_MESSAGES_MAX_SIZE", max_size), + ): + await handle_duplicate_spam(mock_update, mock_context) + + outer = mock_context.bot_data[RECENT_MESSAGES_KEY] + assert len(outer) < max_size + overage, ( + f"Dict size {len(outer)} did not shrink below {max_size + overage}" + ) + assert len(outer) <= max_size + 1, ( + f"Dict size {len(outer)} exceeds bounded limit {max_size + 1}" + ) + + async def test_eviction_keeps_recent_entries( + self, mock_update, mock_context, group_config + ): + """Recently-active entries survive eviction; old ones removed.""" + now = datetime.now(UTC) + max_size = 10 + half = max_size // 2 + + recent_dict: dict[tuple[int, int], deque[RecentMessage]] = {} + last_touch_dict: dict[tuple[int, int], datetime] = {} + + # Old entries (oldest timestamps — should be evicted) + for i in range(half): + gid = -(1000 + i) + uid = i + recent_dict[(gid, uid)] = deque() + last_touch_dict[(gid, uid)] = now - timedelta(days=30) + + # Recent entries (should survive) + for i in range(half, max_size): + gid = -(1000 + i) + uid = i + recent_dict[(gid, uid)] = deque() + last_touch_dict[(gid, uid)] = now + + mock_context.bot_data[RECENT_MESSAGES_KEY] = recent_dict + mock_context.bot_data["duplicate_spam_recent_last_touch"] = last_touch_dict + + mock_update.message.from_user.id = 9999 + + with ( + patch("bot.handlers.duplicate_spam.get_group_config_for_update", return_value=group_config), + patch("bot.handlers.duplicate_spam.RECENT_MESSAGES_MAX_SIZE", max_size), + ): + await handle_duplicate_spam(mock_update, mock_context) + + outer = mock_context.bot_data[RECENT_MESSAGES_KEY] + + # Recent entries survive + for i in range(half, max_size): + assert (-(1000 + i), i) in outer, f"Recent entry ({i}) was evicted" + + # Old entries gone + for i in range(half): + assert (-(1000 + i), i) not in outer, f"Old entry ({i}) survived eviction" + diff --git a/tests/test_plugin_manager.py b/tests/test_plugin_manager.py index 728d0d9..c83a227 100644 --- a/tests/test_plugin_manager.py +++ b/tests/test_plugin_manager.py @@ -140,6 +140,7 @@ def _expected_order() -> tuple[str, ...]: "untrust_callback", "captcha", "dm", + "status", "inline_keyboard_spam", "contact_spam", "new_user_spam", diff --git a/tests/test_status_command.py b/tests/test_status_command.py new file mode 100644 index 0000000..7508363 --- /dev/null +++ b/tests/test_status_command.py @@ -0,0 +1,173 @@ +"""Tests for the /status command handler.""" + +import time +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from bot.database.service import get_database, init_database, reset_database +from bot.group_config import GroupConfig, GroupRegistry +from bot.handlers.status import handle_status + + +@pytest.fixture(autouse=True) +def temp_db(): + with TemporaryDirectory() as tmpdir: + db_path = Path(tmpdir) / "test.db" + reset_database() + init_database(str(db_path)) + yield db_path + reset_database() + + +@pytest.fixture +def mock_registry(): + registry = GroupRegistry() + registry.register(GroupConfig( + group_id=-1001, warning_topic_id=11, + captcha_enabled=True, + )) + registry.register(GroupConfig( + group_id=-1002, warning_topic_id=12, + captcha_enabled=False, + )) + return registry + + +@pytest.fixture +def mock_settings(): + settings = MagicMock() + settings.database_path = "/tmp/test.db" + return settings + + +@pytest.fixture +def mock_update(): + update = MagicMock() + update.message = MagicMock() + update.message.from_user = MagicMock() + update.message.from_user.id = 12345 + update.message.from_user.full_name = "Admin User" + update.message.reply_text = AsyncMock() + update.effective_chat = MagicMock() + update.effective_chat.type = "private" + return update + + +@pytest.fixture +def mock_context(): + context = MagicMock() + context.bot = MagicMock() + context.bot_data = { + "admin_ids": [12345], + "start_time": time.monotonic(), + "plugin_effective_map": { + -1001: {"captcha": True, "spam": True}, + -1002: {"captcha": False, "spam": True, "profile_monitor": False}, + }, + } + return context + + +class TestHandleStatus: + + async def test_handle_status_non_private_chat_rejected(self, mock_context): + """Group chat → handler replies with DM-only error.""" + update = MagicMock() + update.message = MagicMock() + update.message.from_user = MagicMock() + update.message.from_user.id = 12345 + update.message.reply_text = AsyncMock() + update.effective_chat = MagicMock() + update.effective_chat.type = "group" + + await handle_status(update, mock_context) + + update.message.reply_text.assert_called_once() + args, _ = update.message.reply_text.call_args + assert "chat pribadi" in args[0] + + async def test_handle_status_non_admin_rejected(self, mock_context): + """Private chat but caller not admin → handler replies with no-permission.""" + mock_context.bot_data["admin_ids"] = [99999] + non_admin_update = MagicMock() + non_admin_update.message = MagicMock() + non_admin_update.message.from_user = MagicMock() + non_admin_update.message.from_user.id = 111 + non_admin_update.message.from_user.full_name = "Bad Actor" + non_admin_update.message.reply_text = AsyncMock() + non_admin_update.effective_chat = MagicMock() + non_admin_update.effective_chat.type = "private" + + await handle_status(non_admin_update, mock_context) + + non_admin_update.message.reply_text.assert_called_once() + args, _ = non_admin_update.message.reply_text.call_args + assert "tidak memiliki izin" in args[0] + + async def test_handle_status_admin_success( + self, mock_update, mock_context, mock_registry, mock_settings, + ): + """Admin in private chat gets full status with all sections.""" + with ( + patch("bot.handlers.status.get_group_registry", return_value=mock_registry), + patch("bot.handlers.status.get_settings", return_value=mock_settings), + ): + await handle_status(mock_update, mock_context) + + mock_update.message.reply_text.assert_called_once() + args, kwargs = mock_update.message.reply_text.call_args + text = args[0] + assert "*Uptime:*" in text + assert "*Groups:*" in text + assert "*Probation:*" in text + assert "*Captcha:*" in text + assert "*Database:*" in text + assert "*Last jobs:*" in text + + async def test_handle_status_shows_pending_captcha_count( + self, mock_update, mock_context, mock_registry, mock_settings, + ): + """Pending captchas appear in status reply.""" + db = get_database() + db.add_pending_captcha( + user_id=111, group_id=-1001, + chat_id=-1001, message_id=1, + user_full_name="User1", + ) + db.add_pending_captcha( + user_id=222, group_id=-1002, + chat_id=-1002, message_id=2, + user_full_name="User2", + ) + + with ( + patch("bot.handlers.status.get_group_registry", return_value=mock_registry), + patch("bot.handlers.status.get_settings", return_value=mock_settings), + ): + await handle_status(mock_update, mock_context) + + args, _ = mock_update.message.reply_text.call_args + text = args[0] + assert "*Captcha:* 2 pending" in text + + async def test_handle_status_shows_last_job_timestamps( + self, mock_update, mock_context, mock_registry, mock_settings, + ): + """Timestamps for last jobs appear in status reply.""" + mock_context.bot_data["last_admin_refresh"] = time.time() - 60 + mock_context.bot_data["last_auto_restrict"] = time.time() - 300 + + with ( + patch("bot.handlers.status.get_group_registry", return_value=mock_registry), + patch("bot.handlers.status.get_settings", return_value=mock_settings), + ): + await handle_status(mock_update, mock_context) + + args, _ = mock_update.message.reply_text.call_args + text = args[0] + assert "admin refresh:" in text + assert "auto restrict:" in text + assert "never" not in text.split("admin refresh:")[1].split("\n")[0] diff --git a/tests/test_telegram_utils.py b/tests/test_telegram_utils.py index 2ac7a42..eb4dea3 100644 --- a/tests/test_telegram_utils.py +++ b/tests/test_telegram_utils.py @@ -2,7 +2,7 @@ import pytest from telegram import Chat, User -from telegram.error import BadRequest, Forbidden +from telegram.error import BadRequest, Forbidden, RetryAfter from bot.services.telegram_utils import ( fetch_group_admin_ids, @@ -10,6 +10,8 @@ get_user_mention_by_id, get_user_status, is_user_admin_or_trusted, + restrict_chat_member_with_retry, + send_message_with_retry, unrestrict_user, ) @@ -643,3 +645,131 @@ async def test_fetch_admins_different_exceptions(self, mock_bot): with pytest.raises(Exception): await fetch_group_admin_ids(mock_bot, group_id=456) + + +class TestSendMessageWithRetry: + """send_message_with_retry handles RetryAfter correctly.""" + + async def test_success_no_retry(self): + """Normal success: no retry, no sleep, returns True.""" + bot = MagicMock() + bot.send_message = AsyncMock(return_value=MagicMock()) + + result = await send_message_with_retry(bot, chat_id=-100, text="hello") + + assert result is True + bot.send_message.assert_awaited_once_with(chat_id=-100, text="hello") + + async def test_retries_on_retry_after(self): + """RetryAfter once then success: sleeps, retries, returns True.""" + bot = MagicMock() + bot.send_message = AsyncMock( + side_effect=[ + RetryAfter(retry_after=2), + MagicMock(), + ] + ) + + with patch("bot.services.telegram_utils.asyncio.sleep") as mock_sleep: + mock_sleep.return_value = None + result = await send_message_with_retry(bot, chat_id=-100, text="hello") + + assert result is True + assert bot.send_message.await_count == 2 + mock_sleep.assert_called_once_with(3) + + async def test_gives_up_after_second_retry_after(self): + """Second RetryAfter: returns False.""" + bot = MagicMock() + bot.send_message = AsyncMock( + side_effect=RetryAfter(retry_after=1), + ) + + with patch("bot.services.telegram_utils.asyncio.sleep") as mock_sleep: + mock_sleep.return_value = None + result = await send_message_with_retry(bot, chat_id=-100, text="hello") + + assert result is False + assert bot.send_message.await_count == 2 + + async def test_propagates_other_telegram_error(self): + """Non-RetryAfter TelegramError re-raises (caller's except catches it).""" + bot = MagicMock() + bot.send_message = AsyncMock(side_effect=BadRequest("User not found")) + + with pytest.raises(BadRequest): + await send_message_with_retry(bot, chat_id=-100, text="hello") + + bot.send_message.assert_awaited_once() + + +class TestRestrictChatMemberWithRetry: + """restrict_chat_member_with_retry handles RetryAfter correctly.""" + + async def test_success_no_retry(self): + """Normal success: returns True.""" + bot = MagicMock() + bot.restrict_chat_member = AsyncMock(return_value=MagicMock()) + permissions = MagicMock() + + result = await restrict_chat_member_with_retry( + bot, chat_id=-100, user_id=123, permissions=permissions, + ) + + assert result is True + bot.restrict_chat_member.assert_awaited_once_with( + chat_id=-100, user_id=123, permissions=permissions, + ) + + async def test_retries_on_retry_after(self): + """RetryAfter once then success: sleeps, retries, returns True.""" + bot = MagicMock() + permissions = MagicMock() + bot.restrict_chat_member = AsyncMock( + side_effect=[ + RetryAfter(retry_after=2), + MagicMock(), + ] + ) + + with patch("bot.services.telegram_utils.asyncio.sleep") as mock_sleep: + mock_sleep.return_value = None + result = await restrict_chat_member_with_retry( + bot, chat_id=-100, user_id=123, permissions=permissions, + ) + + assert result is True + assert bot.restrict_chat_member.await_count == 2 + mock_sleep.assert_called_once_with(3) + + async def test_gives_up_after_second_retry_after(self): + """Second RetryAfter: returns False.""" + bot = MagicMock() + permissions = MagicMock() + bot.restrict_chat_member = AsyncMock( + side_effect=RetryAfter(retry_after=1), + ) + + with patch("bot.services.telegram_utils.asyncio.sleep") as mock_sleep: + mock_sleep.return_value = None + result = await restrict_chat_member_with_retry( + bot, chat_id=-100, user_id=123, permissions=permissions, + ) + + assert result is False + assert bot.restrict_chat_member.await_count == 2 + + async def test_propagates_other_telegram_error(self): + """Non-RetryAfter TelegramError re-raises.""" + bot = MagicMock() + permissions = MagicMock() + bot.restrict_chat_member = AsyncMock( + side_effect=BadRequest("User not found"), + ) + + with pytest.raises(BadRequest): + await restrict_chat_member_with_retry( + bot, chat_id=-100, user_id=123, permissions=permissions, + ) + + bot.restrict_chat_member.assert_awaited_once()