Add towel like mode for rolls

bot/skills/mute.py

@@ -4,7 +4,15 @@
 from typing import List
 
 from telegram import Update, User, ChatPermissions, TelegramError
-from telegram.ext import Updater, CommandHandler, CallbackContext
+from telegram.ext import (
+    Updater,
+    CommandHandler,
+    CallbackContext,
+    Filters,
+    MessageHandler,
+)
+
+from skills.towel_like_mode import catch_message, quarantine_user, unquarantine_user
 
 from filters import admin_filter
 from mode import cleanup_queue_update
@@ -28,6 +36,14 @@ def add_mute(upd: Updater, handlers_group: int):
         CommandHandler("unmute", unmute, filters=admin_filter, run_async=True),
         handlers_group,
     )
+    dp.add_handler(
+        MessageHandler(
+            Filters.chat_type.groups & ~Filters.status_update,
+            catch_message,
+            run_async=True,
+        ),
+        handlers_group,
+    )
 
 
 def _get_minutes(args: List[str]) -> timedelta:
@@ -65,6 +81,8 @@ def mute_user_for_time(
     except TelegramError as err:
         logger.error("can't mute user %s: %s", user, err)
         update.message.reply_text(f"😿 не вышло, потому что: \n\n{err}")
+        update.message.reply_text("Ну и что? Кинем полотенчик 🧻")
+        quarantine_user(user, mute_duration)
 
 
 def mute(update: Update, context: CallbackContext):
@@ -114,3 +132,4 @@ def unmute_user(update: Update, context: CallbackContext, user: User) -> None:
 def unmute(update: Update, context: CallbackContext) -> None:
     user = update.message.reply_to_message.from_user
     unmute_user(update, context, user)
+    unquarantine_user(user.id)

bot/skills/towel_like_mode.py

@@ -0,0 +1,86 @@
+import logging
+from datetime import datetime, timedelta
+from typing import Dict
+
+from pymongo.collection import Collection
+from telegram import Update, User
+from telegram.ext import (
+    CallbackContext,
+)
+
+from db.mongo import get_db
+from mode import Mode
+
+logger = logging.getLogger(__name__)
+
+
+# todo: extract maybe?
+class DB:
+    def __init__(self, db_name: str):
+        self._coll: Collection = get_db(db_name).quarantine
+
+    def add_user(self, user_id: str, ban_time: timedelta):
+        return (
+            self._coll.insert_one(
+                {
+                    "_id": user_id,
+                    "rel_messages": [],
+                    "datetime": datetime.now() + ban_time,
+                }
+            )
+            if self.find_user(user_id) is None
+            else None
+        )
+
+    def find_user(self, user_id: str):
+        return self._coll.find_one({"_id": user_id})
+
+    def find_all_users(self):
+        return self._coll.find({})
+
+    def add_user_rel_message(self, user_id: str, message_id: str):
+        self._coll.update_one(
+            {"_id": user_id}, {"$addToSet": {"rel_messages": message_id}}
+        )
+
+    def delete_user(self, user_id: str):
+        return self._coll.delete_one({"_id": user_id})
+
+    def delete_all_users(self):
+        return self._coll.delete_many({})
+
+
+db = DB("towel_like_mode")
+mode = Mode(
+    mode_name="towel_like_mode",
+    default=True,
+    off_callback=lambda _: db.delete_all_users(),
+)
+
+
+def _is_time_gone(user: Dict) -> bool:
+    return user["datetime"] < datetime.now()
+
+
+def unquarantine_user(user: User):
+    db.delete_user(user_id=user.id)
+
+
+def quarantine_user(user: User, ban_time: timedelta):
+    logger.info("put %s in quarantine", user)
+    db.add_user(user.id, ban_time)
+
+
+def catch_message(update: Update, context: CallbackContext):
+    # todo: cache it
+    user_id = update.effective_user.id
+    user = db.find_user(user_id)
+    if user is None:
+        return
+
+    if not _is_time_gone(user):
+        context.bot.delete_message(
+            update.effective_chat.id, update.effective_message.message_id, 10
+        )
+    else:
+        db.delete_user(user_id=user_id)