diff --git a/.env.example b/.env.example index 5fcc6a2..0a4ae55 100644 --- a/.env.example +++ b/.env.example @@ -86,6 +86,15 @@ 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 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= + # Logfire Configuration (optional - for production logging) # Get your token from https://logfire.pydantic.dev LOGFIRE_TOKEN=your_logfire_token_here diff --git a/AGENTS.md b/AGENTS.md index e048855..e1e0032 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) @@ -80,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) @@ -101,38 +102,40 @@ 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=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` | 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 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 ### 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` @@ -145,7 +148,7 @@ 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=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) @@ -174,6 +177,24 @@ 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 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) +- 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 — 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` + ### 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 +216,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 +234,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_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 @@ -294,7 +316,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 @@ -310,6 +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=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 d0503fc..32e00c4 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,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 @@ -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 @@ -360,12 +363,21 @@ 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=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]) + 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?} @@ -585,6 +597,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 +678,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 — 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` +- `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 +729,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` | 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..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 @@ -32,9 +32,21 @@ 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" + __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) @@ -44,6 +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") class PhotoVerificationWhitelist(SQLModel, table=True): diff --git a/src/bot/database/service.py b/src/bot/database/service.py index 6a8187b..2d7b511 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,28 @@ 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 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( + self, user_id: int, group_id: int, warning_kind: str = "profile" + ) -> UserWarning: """ Get existing warning record or create a new one. @@ -90,6 +112,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 +122,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 +140,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 +162,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 +174,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 +186,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 +206,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 +218,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 +231,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 +252,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 +261,76 @@ 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 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: """ Clear bot restriction flag after user is unrestricted via DM. @@ -241,11 +340,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 +357,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 +372,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 +381,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 +404,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 +413,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 +649,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 +657,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 +666,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/captcha.py b/src/bot/handlers/captcha.py index bdd1a3d..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, @@ -31,6 +32,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 @@ -319,32 +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: - await unrestrict_user(context.bot, group_config.group_id, target_user_id) - logger.info(f"Unrestricted verified user {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) + 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/dm.py b/src/bot/handlers/dm.py index f7869ba..963a339 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, @@ -80,18 +81,31 @@ async def _unrestrict_in_groups( 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}") + logger.info(f"Processing unrestrict for user_id={user.id} 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 + 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 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 new file mode 100644 index 0000000..f8707bc --- /dev/null +++ b/src/bot/handlers/guest_bot.py @@ -0,0 +1,133 @@ +"""Guest bot message moderation for Telegram Guest Mode.""" + +import logging + +from telegram import Message, Update, User +from telegram.error import TelegramError +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.restriction_lock import restriction_lock +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 + return username.lower() in 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 TelegramError: + logger.error("Failed to delete guest bot message", exc_info=True) + + if 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 >= group_config.warning_threshold: + should_stop = False + final_count = record.message_count + 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: + 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 + 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) + 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_WARNING.format( + user_mention=user_mention, + warning_threshold=group_config.warning_threshold, + rules_link=group_config.rules_link, + ), + parse_mode="Markdown", + ) + 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") + + raise ApplicationHandlerStop diff --git a/src/bot/handlers/message.py b/src/bot/handlers/message.py index 3449fb6..f1cc759 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,27 @@ 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.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})" + ) + 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) @@ -177,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, @@ -192,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 16124b0..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 +from telegram.error import TelegramError from telegram.ext import ContextTypes from bot.constants import ( @@ -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, @@ -66,29 +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, - ) - - deleted_count = db.delete_user_warnings(target_user_id, group_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(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) - 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}" + unrestrict_failed = False + deleted_count = 0 + + 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 (TelegramError, 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" ) + if unrestrict_failed: + return UNRESTRICT_FAILED_MESSAGE.format( + user_id=target_user_id, group_id=group_id + ) + if deleted_count > 0 or did_unrestrict: try: user_info = await bot.get_chat(target_user_id) @@ -156,14 +178,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): - 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) + 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( f"Admin unrestricting user {target_user_id} in group {group_id}" ) diff --git a/src/bot/plugins/builtin/spam.py b/src/bot/plugins/builtin/spam.py index c445c64..f982848 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=0). + + Callback wrapped with ``guard_plugin(\"guest_bot_block\")``. Runs at + 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, 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 cc9617f..3a8dc31 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": 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"}, @@ -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/src/bot/services/captcha_recovery.py b/src/bot/services/captcha_recovery.py index 66e62ed..d47d48a 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__) @@ -45,18 +46,23 @@ 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 - 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): + 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) 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/restriction_lock.py b/src/bot/services/restriction_lock.py new file mode 100644 index 0000000..78caca9 --- /dev/null +++ b/src/bot/services/restriction_lock.py @@ -0,0 +1,57 @@ +"""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. ``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_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 = per_loop.get(key) + if lock is None: + lock = asyncio.Lock() + per_loop[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..7ee5e81 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,26 @@ 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): + fresh = db.get_active_user_warning(warning.user_id, warning.group_id) + 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})" + ) + continue + 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_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_database.py b/tests/test_database.py index dd9c92d..8785b6a 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -457,3 +457,235 @@ 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() + + 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_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_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..c5d32b1 --- /dev/null +++ b/tests/test_guest_bot.py @@ -0,0 +1,211 @@ +"""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_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") + 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): + 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 + 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..dc955a5 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=0.""" + 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"] == 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 145c0bc..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.""" @@ -144,6 +144,7 @@ def _expected_order() -> tuple[str, ...]: "captcha", "dm", "status", + "guest_bot_block", "inline_keyboard_spam", "contact_spam", "new_user_spam", diff --git a/tests/test_restriction_lock.py b/tests/test_restriction_lock.py new file mode 100644 index 0000000..a857bd9 --- /dev/null +++ b/tests/test_restriction_lock.py @@ -0,0 +1,51 @@ +"""Tests for the per-(group_id, user_id) restriction lock.""" + +import asyncio + +from bot.services.restriction_lock import _locks_by_loop, restriction_lock + + +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, scoped to the running loop.""" + key = (-300, 99) + async with restriction_lock(*key): + pass + 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 3a0d7d1..8e2daf3 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,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.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() @@ -197,6 +201,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 +281,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 +387,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 021e3b0..4a8222c 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 @@ -152,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) @@ -169,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) @@ -642,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) @@ -665,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() @@ -842,6 +845,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):