feat(moderation): implement global ban and chat registration features

MattiaFailla · MattiaFailla · commit d4e41dbdaf1f · 2026-02-22T16:26:00.000+01:00
- Added methods for registering active chats and managing global bans in the moderation service.
- Enhanced the in-memory and PostgreSQL repositories to support chat tracking and global ban functionality.
- Introduced command handlers for forcefully registering chats and handling global bans in moderation.
- Updated welcome message handling to automatically register chats and check for globally banned users.
diff --git a/src/python_italy_bot/db/base.py b/src/python_italy_bot/db/base.py
@@ -106,6 +106,36 @@ def mark_globally_verified(self, user_id: int) -> None:
         """Mark user as globally verified."""
         ...
 
+    @abstractmethod
+    def register_chat(self, chat_id: int) -> None:
+        """Track a chat where the bot is active."""
+        ...
+
+    @abstractmethod
+    def get_all_chats(self) -> list[int]:
+        """Get all tracked chat IDs."""
+        ...
+
+    @abstractmethod
+    def add_global_ban(
+        self,
+        user_id: int,
+        admin_id: int,
+        reason: str | None = None,
+    ) -> None:
+        """Add a global ban for a user."""
+        ...
+
+    @abstractmethod
+    def remove_global_ban(self, user_id: int) -> bool:
+        """Remove a global ban. Returns True if existed."""
+        ...
+
+    @abstractmethod
+    def is_globally_banned(self, user_id: int) -> bool:
+        """Check if user is globally banned."""
+        ...
+
 
 class AsyncRepository(ABC):
     """Abstract interface for data persistence (async)."""
@@ -210,6 +240,36 @@ async def mark_globally_verified(self, user_id: int) -> None:
         """Mark user as globally verified."""
         ...
 
+    @abstractmethod
+    async def register_chat(self, chat_id: int) -> None:
+        """Track a chat where the bot is active."""
+        ...
+
+    @abstractmethod
+    async def get_all_chats(self) -> list[int]:
+        """Get all tracked chat IDs."""
+        ...
+
+    @abstractmethod
+    async def add_global_ban(
+        self,
+        user_id: int,
+        admin_id: int,
+        reason: str | None = None,
+    ) -> None:
+        """Add a global ban for a user."""
+        ...
+
+    @abstractmethod
+    async def remove_global_ban(self, user_id: int) -> bool:
+        """Remove a global ban. Returns True if existed."""
+        ...
+
+    @abstractmethod
+    async def is_globally_banned(self, user_id: int) -> bool:
+        """Check if user is globally banned."""
+        ...
+
     async def close(self) -> None:
         """Close any resources (override if needed)."""
         pass
diff --git a/src/python_italy_bot/db/in_memory.py b/src/python_italy_bot/db/in_memory.py
@@ -17,6 +17,8 @@ def __init__(self) -> None:
         self._reports: list[Report] = []
         self._welcome_messages: dict[int, str] = {}
         self._globally_verified: set[int] = set()
+        self._bot_chats: set[int] = set()
+        self._global_bans: dict[int, tuple[int, str | None]] = {}
 
     async def add_pending_verification(self, user_id: int, chat_id: int) -> None:
         self._pending.add((user_id, chat_id))
@@ -134,3 +136,26 @@ async def is_globally_verified(self, user_id: int) -> bool:
 
     async def mark_globally_verified(self, user_id: int) -> None:
         self._globally_verified.add(user_id)
+
+    async def register_chat(self, chat_id: int) -> None:
+        self._bot_chats.add(chat_id)
+
+    async def get_all_chats(self) -> list[int]:
+        return list(self._bot_chats)
+
+    async def add_global_ban(
+        self,
+        user_id: int,
+        admin_id: int,
+        reason: str | None = None,
+    ) -> None:
+        self._global_bans[user_id] = (admin_id, reason)
+
+    async def remove_global_ban(self, user_id: int) -> bool:
+        if user_id in self._global_bans:
+            del self._global_bans[user_id]
+            return True
+        return False
+
+    async def is_globally_banned(self, user_id: int) -> bool:
+        return user_id in self._global_bans
diff --git a/src/python_italy_bot/db/postgres.py b/src/python_italy_bot/db/postgres.py
@@ -207,5 +207,59 @@ async def mark_globally_verified(self, user_id: int) -> None:
                 (user_id,),
             )
 
+    async def register_chat(self, chat_id: int) -> None:
+        async with self._pool.connection() as conn:
+            await conn.execute(
+                """
+                INSERT INTO bot_chats (chat_id)
+                VALUES (%s)
+                ON CONFLICT (chat_id) DO NOTHING
+                """,
+                (chat_id,),
+            )
+
+    async def get_all_chats(self) -> list[int]:
+        async with self._pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute("SELECT chat_id FROM bot_chats")
+                rows = await cur.fetchall()
+                return [row[0] for row in rows]
+
+    async def add_global_ban(
+        self,
+        user_id: int,
+        admin_id: int,
+        reason: str | None = None,
+    ) -> None:
+        async with self._pool.connection() as conn:
+            await conn.execute(
+                """
+                INSERT INTO global_bans (user_id, admin_id, reason)
+                VALUES (%s, %s, %s)
+                ON CONFLICT (user_id)
+                DO UPDATE SET admin_id = EXCLUDED.admin_id,
+                              reason = EXCLUDED.reason,
+                              created_at = NOW()
+                """,
+                (user_id, admin_id, reason),
+            )
+
+    async def remove_global_ban(self, user_id: int) -> bool:
+        async with self._pool.connection() as conn:
+            result = await conn.execute(
+                "DELETE FROM global_bans WHERE user_id = %s",
+                (user_id,),
+            )
+            return result.rowcount > 0
+
+    async def is_globally_banned(self, user_id: int) -> bool:
+        async with self._pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT 1 FROM global_bans WHERE user_id = %s",
+                    (user_id,),
+                )
+                return await cur.fetchone() is not None
+
     async def close(self) -> None:
         await self._pool.close()
diff --git a/src/python_italy_bot/handlers/moderation.py b/src/python_italy_bot/handlers/moderation.py
@@ -34,11 +34,36 @@ def create_moderation_handlers(moderation_service: ModerationService) -> list:
         CommandHandler("mute", _handle_mute),
         CommandHandler("unmute", _handle_unmute),
         CommandHandler("report", _handle_report),
+        CommandHandler("forcegroupregistration", _handle_force_group_registration),
     ]
 
 
+async def _handle_force_group_registration(
+    update: Update, context: ContextTypes.DEFAULT_TYPE
+) -> None:
+    """Force registration of current chat in bot_chats table. Admin only."""
+    moderation_service: ModerationService = context.bot_data["moderation_service"]
+    message = update.message
+    if message is None or message.from_user is None:
+        return
+
+    chat = update.effective_chat
+    if chat is None or chat.type == "private":
+        await message.reply_text("Questo comando funziona solo nei gruppi.")
+        return
+
+    if not await _is_admin(context, chat.id, message.from_user.id):
+        await message.reply_text(
+            "Solo gli amministratori possono usare questo comando."
+        )
+        return
+
+    await moderation_service.register_chat(chat.id)
+    await message.reply_text(f"Gruppo registrato. Chat ID: {chat.id}")
+
+
 async def _handle_ban(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """Ban a user. Usage: /ban @username or /ban user_id [reason] or reply to message with /ban [reason]"""
+    """Ban a user globally. Usage: /ban user_id [reason] or reply to message with /ban [reason]"""
     moderation_service: ModerationService = context.bot_data["moderation_service"]
     message = update.message
     if message is None or message.from_user is None:
@@ -60,30 +85,41 @@ async def _handle_ban(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non
     reason: str | None = None
     if message.reply_to_message and message.reply_to_message.from_user:
         user_id = message.reply_to_message.from_user.id
-        reason = args[0] if args else None
+        reason = " ".join(args) if args else None
     elif args:
         target = args[0]
         reason = args[1] if len(args) > 1 else None
         user_id = await _resolve_user_id(context, chat.id, target)
 
     if user_id is None:
         await message.reply_text(
-            "Uso: /ban @username, /ban user_id [motivo], o rispondi al messaggio con /ban [motivo]. "
-            "Per @username funziona solo con amministratori."
+            "Uso: /ban user_id [motivo], o rispondi al messaggio con /ban [motivo]."
         )
         return
 
-    try:
-        await context.bot.ban_chat_member(chat.id, user_id)
-        await moderation_service.add_ban(user_id, chat.id, message.from_user.id, reason)
-        await message.reply_text(f"Utente bannato. Motivo: {reason or 'Nessuno'}")
-    except Exception as e:
-        logger.warning("Ban failed: %s", e)
-        await message.reply_text("Impossibile bannare l'utente.")
+    chat_ids = await moderation_service.add_global_ban(
+        user_id, message.from_user.id, reason
+    )
+
+    success_count = 0
+    fail_count = 0
+    for cid in chat_ids:
+        try:
+            await context.bot.ban_chat_member(cid, user_id)
+            success_count += 1
+        except Exception as e:
+            logger.debug("Ban in chat %s failed: %s", cid, e)
+            fail_count += 1
+
+    msg = f"Utente bannato globalmente in {success_count} gruppi."
+    if fail_count > 0:
+        msg += f" ({fail_count} falliti)"
+    msg += f"\nMotivo: {reason or 'Nessuno'}"
+    await message.reply_text(msg)
 
 
 async def _handle_unban(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """Unban a user. Usage: /unban @username or /unban user_id"""
+    """Unban a user globally. Usage: /unban user_id or reply to message with /unban"""
     moderation_service: ModerationService = context.bot_data["moderation_service"]
     message = update.message
     if message is None or message.from_user is None:
@@ -108,19 +144,26 @@ async def _handle_unban(update: Update, context: ContextTypes.DEFAULT_TYPE) -> N
 
     if user_id is None:
         await message.reply_text(
-            "Uso: /unban @username, /unban user_id, o rispondi al messaggio"
+            "Uso: /unban user_id, o rispondi al messaggio con /unban"
         )
-    if user_id is None:
-        await message.reply_text("Utente non trovato.")
         return
 
-    try:
-        await context.bot.unban_chat_member(chat.id, user_id)
-        await moderation_service.remove_ban(user_id, chat.id)
-        await message.reply_text("Utente sbannato.")
-    except Exception as e:
-        logger.warning("Unban failed: %s", e)
-        await message.reply_text("Impossibile sbannare l'utente.")
+    chat_ids = await moderation_service.remove_global_ban(user_id)
+
+    success_count = 0
+    fail_count = 0
+    for cid in chat_ids:
+        try:
+            await context.bot.unban_chat_member(cid, user_id)
+            success_count += 1
+        except Exception as e:
+            logger.debug("Unban in chat %s failed: %s", cid, e)
+            fail_count += 1
+
+    msg = f"Utente sbannato globalmente da {success_count} gruppi."
+    if fail_count > 0:
+        msg += f" ({fail_count} falliti)"
+    await message.reply_text(msg)
 
 
 async def _handle_mute(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
diff --git a/src/python_italy_bot/handlers/welcome.py b/src/python_italy_bot/handlers/welcome.py
@@ -13,6 +13,7 @@
 )
 
 from ..services.captcha import CaptchaService
+from ..services.moderation import ModerationService
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +39,7 @@ async def _handle_new_member(
 ) -> None:
     """Handle new chat members: restrict and send welcome with captcha instructions."""
     captcha_service: CaptchaService = context.bot_data["captcha_service"]
+    moderation_service: ModerationService = context.bot_data["moderation_service"]
     result = update.chat_member
     if result is None:
         return
@@ -59,9 +61,19 @@ async def _handle_new_member(
     if user is None or chat is None:
         return
 
+    await moderation_service.register_chat(chat.id)
+
     if user.is_bot:
         return
 
+    if await moderation_service.is_globally_banned(user.id):
+        try:
+            await context.bot.ban_chat_member(chat.id, user.id)
+            logger.info("Kicked globally banned user %s from chat %s", user.id, chat.id)
+        except Exception as e:
+            logger.warning("Failed to kick globally banned user %s: %s", user.id, e)
+        return
+
     if await captcha_service.is_globally_verified(user.id):
         return
 
diff --git a/src/python_italy_bot/services/moderation.py b/src/python_italy_bot/services/moderation.py
@@ -39,21 +39,34 @@ async def is_muted(self, user_id: int, chat_id: int) -> bool:
         muted_users = await self._repo.get_muted_users(chat_id)
         return user_id in muted_users
 
-    async def add_ban(
+    async def add_global_ban(
         self,
         user_id: int,
-        chat_id: int,
         admin_id: int,
         reason: str | None = None,
-    ) -> None:
-        """Record a ban."""
-        await self._repo.add_ban(
-            user_id=user_id, chat_id=chat_id, admin_id=admin_id, reason=reason
+    ) -> list[int]:
+        """Record a global ban. Returns list of chat IDs to ban in."""
+        await self._repo.add_global_ban(
+            user_id=user_id, admin_id=admin_id, reason=reason
         )
+        return await self._repo.get_all_chats()
+
+    async def remove_global_ban(self, user_id: int) -> list[int]:
+        """Remove a global ban. Returns list of chat IDs to unban in."""
+        await self._repo.remove_global_ban(user_id)
+        return await self._repo.get_all_chats()
+
+    async def is_globally_banned(self, user_id: int) -> bool:
+        """Check if user is globally banned."""
+        return await self._repo.is_globally_banned(user_id)
+
+    async def register_chat(self, chat_id: int) -> None:
+        """Register a chat where the bot is active."""
+        await self._repo.register_chat(chat_id)
 
-    async def remove_ban(self, user_id: int, chat_id: int) -> bool:
-        """Remove a ban record."""
-        return await self._repo.remove_ban(user_id, chat_id)
+    async def get_all_chats(self) -> list[int]:
+        """Get all tracked chat IDs."""
+        return await self._repo.get_all_chats()
 
     async def add_mute(
         self,
