diff --git a/bot/skills/mute.py b/bot/skills/mute.py index f70e398..37f16bc 100644 --- a/bot/skills/mute.py +++ b/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) diff --git a/bot/skills/towel_like_mode.py b/bot/skills/towel_like_mode.py new file mode 100644 index 0000000..aceb17d --- /dev/null +++ b/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)