From dc6bf7accb10531a23c4b282541bc20aca212318 Mon Sep 17 00:00:00 2001 From: AlberLC Date: Tue, 31 May 2022 22:11:16 +0200 Subject: [PATCH] Refactor punishments and optimize chat objects --- flanabot/bots/flana_bot.py | 95 +++++++++------------------------ flanabot/bots/flana_disc_bot.py | 20 +++---- flanabot/bots/flana_tele_bot.py | 4 +- flanabot/constants.py | 1 - flanabot/models/chat.py | 5 +- flanabot/models/punishment.py | 30 +++++++++-- 6 files changed, 66 insertions(+), 89 deletions(-) diff --git a/flanabot/bots/flana_bot.py b/flanabot/bots/flana_bot.py index a406435..54e5385 100644 --- a/flanabot/bots/flana_bot.py +++ b/flanabot/bots/flana_bot.py @@ -146,14 +146,14 @@ class FlanaBot(MultiBot, ABC): return last_2s_messages = Message.find({ - 'platform': self.bot_platform.value, + 'platform': self.platform.value, 'author': message.author.object_id, 'last_update': { '$gte': datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=2) } }) last_7s_messages = Message.find({ - 'platform': self.bot_platform.value, + 'platform': self.platform.value, 'author': message.author.object_id, 'last_update': { '$gte': datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=7), @@ -163,7 +163,7 @@ class FlanaBot(MultiBot, ABC): if len(last_2s_messages) >= 5 or len(last_7s_messages) >= 7: n_punishments = len(Punishment.find({ - 'platform': self.bot_platform.value, + 'platform': self.platform.value, 'user_id': message.author.id, 'group_id': message.chat.group_id })) @@ -175,23 +175,6 @@ class FlanaBot(MultiBot, ABC): else: await self.send(f'Castigado durante {TimeUnits(seconds=punishment_seconds).to_words()}.', message) - async def _check_punishments(self): - punishment_groups = self._get_grouped_punishments(Punishment) - - now = datetime.datetime.now(datetime.timezone.utc) - for (user_id, group_id), sorted_punishments in punishment_groups: - if now < (last_punishment := sorted_punishments[-1]).until: - continue - - if last_punishment.until + constants.PUNISHMENTS_RESET <= now: - for old_punishment in sorted_punishments: - old_punishment.delete() - - if last_punishment.is_active: - await self.unpunish(user_id, group_id) - last_punishment.is_active = False - last_punishment.save() - @staticmethod def _get_config_buttons(config: dict | Chat | Message) -> list[list[str]]: match config: @@ -261,7 +244,7 @@ class FlanaBot(MultiBot, ABC): new_line = ' ' if len(medias_sended_info) == 1 else '\n' return f'{new_line}{medias_sended_info_joined}:' - async def _punish(self, user: int | str | User, group_: int | str | Chat): + async def _punish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): pass async def _search_and_send_medias(self, message: Message, send_song_info=False) -> list[Message]: @@ -303,7 +286,7 @@ class FlanaBot(MultiBot, ABC): async def _show_config(self, config_name: str, message: Message): await self.send(f"{config_name} está {'activado ✔' if message.chat.config.get(config_name) else 'desactivado ❌'}", message) - async def _unpunish(self, user: int | str | User, group_: int | str | Chat): + async def _unpunish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): pass # ---------------------------------------------- # @@ -311,8 +294,6 @@ class FlanaBot(MultiBot, ABC): # ---------------------------------------------- # async def _on_button_press(self, message: Message): await self._accept_button_event(message) - if not message.button_pressed_user.is_admin: - return match message.button_pressed_text: case WeatherEmoji.ZOOM_IN.value: @@ -335,13 +316,13 @@ class FlanaBot(MultiBot, ABC): buttons_message_type = ButtonsMessageType.WEATHER trace_metadata_name = WeatherEmoji(emoji).name.lower() message.weather_chart.trace_metadatas[trace_metadata_name].show = not message.weather_chart.trace_metadatas[trace_metadata_name].show - case _ if 'auto_' in (config := message.button_pressed_text.split()[1]): + case _ if message.button_pressed_user.is_admin and 'auto_' in (config := message.button_pressed_text.split()[1]): buttons_message_type = ButtonsMessageType.CONFIG message.chat.config[config] = not message.chat.config[config] message.save() await self.edit('Estos son los ajustes del grupo:\n\n', self._get_config_buttons(message), message) case _: - buttons_message_type = None + return if buttons_message_type is ButtonsMessageType.WEATHER: message.weather_chart.apply_zoom() @@ -469,11 +450,12 @@ class FlanaBot(MultiBot, ABC): @group @admin(send_negative=True) async def _on_punish(self, message: Message): - await self._update_punishment(self.punish, message, time=flanautils.words_to_time(message.text)) + for user in await self._find_users_to_punish(message): + await self.punish(user, message, flanautils.words_to_time(message.text), message) async def _on_ready(self): await super()._on_ready() - await flanautils.do_every(constants.CHECK_PUNISHMENTS_EVERY_SECONDS, self._check_punishments) + await flanautils.do_every(constants.CHECK_PUNISHMENTS_EVERY_SECONDS, Punishment.check_olds, self._unpunish, self.platform) @inline(False) async def _on_recover_message(self, message: Message): @@ -489,7 +471,7 @@ class FlanaBot(MultiBot, ABC): return affected_object_ids = [affected_message_object_id for affected_message_object_id in message_deleted_bot_action.affected_objects] - deleted_messages: list[Message] = [affected_message for affected_object_id in affected_object_ids if (affected_message := Message.find_one({'platform': self.bot_platform.value, '_id': affected_object_id})).author.id != self.bot_id] + deleted_messages: list[Message] = [affected_message for affected_object_id in affected_object_ids if (affected_message := Message.find_one({'platform': self.platform.value, '_id': affected_object_id})).author.id != self.id] for deleted_message in deleted_messages: await self.send(deleted_message.text, message) @@ -563,7 +545,8 @@ class FlanaBot(MultiBot, ABC): @bot_mentioned @admin(send_negative=True) async def _on_unpunish(self, message: Message): - await self._update_punishment(self.unpunish, message) + for user in await self._find_users_to_punish(message): + await self.unpunish(user, message, message) async def _on_weather_chart(self, message: Message): bot_state_message: Message | None = None @@ -571,7 +554,7 @@ class FlanaBot(MultiBot, ABC): show_progress_state = False elif message.chat.is_group and not self.is_bot_mentioned(message): if message.chat.config['auto_weather_chart']: - if BotAction.find_one({'action': bytes(Action.AUTO_WEATHER_CHART), 'chat': message.chat.object_id, 'date': {'$gt': datetime.datetime.now(datetime.timezone.utc) - constants.AUTO_WEATHER_EVERY}}): + if BotAction.find_one({'action': Action.AUTO_WEATHER_CHART.value, 'chat': message.chat.object_id, 'date': {'$gt': datetime.datetime.now(datetime.timezone.utc) - constants.AUTO_WEATHER_EVERY}}): return show_progress_state = False else: @@ -584,7 +567,8 @@ class FlanaBot(MultiBot, ABC): possible_mentioned_ids.append(user.name.lower()) possible_mentioned_ids.append(user.name.split('#')[0].lower()) possible_mentioned_ids.append(f'@{user.id}') - for role in message.chat.roles: + + for role in await self.get_group_roles(message): possible_mentioned_ids.append(f'@{role.id}') original_text_words = flanautils.remove_accents(message.text.lower()) @@ -733,31 +717,13 @@ class FlanaBot(MultiBot, ABC): # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # # -------------------------------------------------------- # - async def is_punished(self, user: int | str | User, group_: int | str | Chat) -> bool: + async def is_punished(self, user: int | str | User, group_: int | str | Chat | Message) -> bool: pass - async def punish(self, user: int | str | User, group_: int | str | Chat, time: int | datetime.timedelta, message: Message = None): - user_id = self._get_user_id(user) - group_id = self._get_group_id(group_) - if isinstance(time, int): - time = datetime.timedelta(seconds=time) - - try: - await self._punish(user_id, group_id) - except BadRoleError as e: - if message and message.chat.original_object: - await self._manage_exceptions(e, message) - else: - raise e - else: - if time: - until = datetime.datetime.now(datetime.timezone.utc) + time - if datetime.timedelta() < time <= constants.TIME_THRESHOLD_TO_MANUAL_UNPUNISH: - await flanautils.do_later(time, self._check_punishments) - else: - until = None - # noinspection PyTypeChecker - Punishment(self.bot_platform, user_id, group_id, until=until).save() + async def punish(self, user: int | str | User, group_: int | str | Chat | Message, time: int | datetime.timedelta, message: Message = None): + # noinspection PyTypeChecker + punish = Punishment(self.platform, self.get_user_id(user), self.get_group_id(group_), time) + await punish.punish(self._punish, self._unpunish, message) async def send_bye(self, message: Message) -> multibot_constants.ORIGINAL_MESSAGE: return await self.send(random.choice((*constants.BYE_PHRASES, flanautils.CommonWords.random_time_greeting())), message) @@ -825,18 +791,7 @@ class FlanaBot(MultiBot, ABC): if song_info: await self.send(song_info, message) - async def unpunish(self, user: int | str | User, group_: int | str | Chat, message: Message = None): - user_id = self._get_user_id(user) - group_id = self._get_group_id(group_) - try: - await self._unpunish(user_id, group_id) - except BadRoleError as e: - if message and message.chat.original_object: - await self._manage_exceptions(e, message) - else: - raise e - else: - try: - Punishment.find_one({'platform': self.bot_platform.value, 'user_id': user_id, 'group_id': group_id, 'until': None}).delete() - except AttributeError: - pass + async def unpunish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): + # noinspection PyTypeChecker + punish = Punishment(self.platform, self.get_user_id(user), self.get_group_id(group_)) + await punish.unpunish(self._unpunish, message) diff --git a/flanabot/bots/flana_disc_bot.py b/flanabot/bots/flana_disc_bot.py index 12b5c9f..d503864 100644 --- a/flanabot/bots/flana_disc_bot.py +++ b/flanabot/bots/flana_disc_bot.py @@ -6,7 +6,7 @@ from multibot import BadRoleError, DiscordBot, User from flanabot import constants from flanabot.bots.flana_bot import FlanaBot -from flanabot.models import Chat, Punishment +from flanabot.models import Chat, Message, Punishment HEAT_NAMES = [ 'Canal Congelado', @@ -29,7 +29,7 @@ HOT_CHANNEL_ID = 493530483045564417 # ---------------------------------------------------------------------------------------------------- # class FlanaDiscBot(DiscordBot, FlanaBot): def __init__(self): - super().__init__(os.environ['DISCORD_BOT_TOKEN']) + super().__init__(os.environ['DISCORD_BOT_TOKEN2']) self.heating = False self.heat_level = 0 @@ -38,7 +38,7 @@ class FlanaDiscBot(DiscordBot, FlanaBot): # ----------------------------------------------------------- # def _add_handlers(self): super()._add_handlers() - self.bot_client.add_listener(self._on_voice_state_update, 'on_voice_state_update') + self.client.add_listener(self._on_voice_state_update, 'on_voice_state_update') async def _heat_channel(self, channel: discord.VoiceChannel): while True: @@ -57,16 +57,16 @@ class FlanaDiscBot(DiscordBot, FlanaBot): await channel.edit(name=HEAT_NAMES[int(self.heat_level)]) - async def _punish(self, user: int | str | User, group_: int | str | Chat): - user_id = self._get_user_id(user) + async def _punish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): + user_id = self.get_user_id(user) try: await self.add_role(user_id, group_, 'Castigado') await self.remove_role(user_id, group_, 'Persona') except AttributeError: raise BadRoleError(str(self._punish)) - async def _unpunish(self, user: int | str | User, group_: int | str | Chat): - user_id = self._get_user_id(user) + async def _unpunish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): + user_id = self.get_user_id(user) try: await self.add_role(user_id, group_, 'Persona') await self.remove_role(user_id, group_, 'Castigado') @@ -92,11 +92,11 @@ class FlanaDiscBot(DiscordBot, FlanaBot): # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # # -------------------------------------------------------- # - async def is_punished(self, user: int | str | User, group_: int | str | Chat): + async def is_punished(self, user: int | str | User, group_: int | str | Chat | Message): user = await self.get_user(user, group_) - group_id = self._get_group_id(group_) + group_id = self.get_group_id(group_) return group_id in {punishment.group_id for punishment in Punishment.find({ - 'platform': self.bot_platform.value, + 'platform': self.platform.value, 'user_id': user.id, 'group_id': group_id, 'is_active': True diff --git a/flanabot/bots/flana_tele_bot.py b/flanabot/bots/flana_tele_bot.py index 39aba97..21fb484 100644 --- a/flanabot/bots/flana_tele_bot.py +++ b/flanabot/bots/flana_tele_bot.py @@ -36,7 +36,7 @@ class FlanaTeleBot(TelegramBot, FlanaBot): super().__init__( api_id=os.environ['TELEGRAM_API_ID'], api_hash=os.environ['TELEGRAM_API_HASH'], - bot_session=os.environ['TELEGRAM_BOT_SESSION'], + bot_session=os.environ['TELEGRAM_BOT_SESSION2'], user_session=os.environ['TELEGRAM_USER_SESSION'] ) self.whitelist_ids = [] @@ -53,7 +53,7 @@ class FlanaTeleBot(TelegramBot, FlanaBot): @user_client async def _update_whitelist(self): - self.whitelist_ids = [self.owner_id, self.bot_id] + await self._get_contacts_ids() + self.whitelist_ids = [self.owner_id, self.id] + await self._get_contacts_ids() # ---------------------------------------------- # # HANDLERS # diff --git a/flanabot/constants.py b/flanabot/constants.py index 56cf724..8d694ad 100644 --- a/flanabot/constants.py +++ b/flanabot/constants.py @@ -10,7 +10,6 @@ MAX_PLACE_QUERY_LENGTH = 50 PUNISHMENT_INCREMENT_EXPONENT = 6 PUNISHMENTS_RESET = datetime.timedelta(weeks=6 * flanautils.WEEKS_IN_A_MONTH) RECOVERY_DELETED_MESSAGE_BEFORE = datetime.timedelta(hours=1) -TIME_THRESHOLD_TO_MANUAL_UNPUNISH = datetime.timedelta(days=3) SCRAPING_MESSAGE_WAITING_TIME = 0.1 BYE_PHRASES = ('Adiós.', 'adieu', 'adio', 'adioh', 'adios', 'adió', 'adiós', 'agur', 'bye', 'byyeeee', 'chao', diff --git a/flanabot/models/chat.py b/flanabot/models/chat.py index 6c0c792..f1a7b68 100644 --- a/flanabot/models/chat.py +++ b/flanabot/models/chat.py @@ -1,6 +1,6 @@ -from dataclasses import dataclass, field +from dataclasses import dataclass -from multibot.models import Chat as MultiBotChat, User +from multibot.models import Chat as MultiBotChat @dataclass(eq=False) @@ -11,7 +11,6 @@ class Chat(MultiBotChat): 'auto_insult': True, 'auto_scraping': True, 'auto_weather_chart': True} - users: list[User] = field(default_factory=list) def __post_init__(self): super().__post_init__() diff --git a/flanabot/models/punishment.py b/flanabot/models/punishment.py index 5526345..2871645 100644 --- a/flanabot/models/punishment.py +++ b/flanabot/models/punishment.py @@ -1,8 +1,32 @@ -from dataclasses import dataclass +__all__ = ['Punishment'] -from multibot.models import Mute, db +import datetime +from dataclasses import dataclass +from typing import Callable + +from multibot.models import Platform, PunishmentBase, db + +from flanabot import constants @dataclass(eq=False) -class Punishment(Mute): +class Punishment(PunishmentBase): collection = db.punishment + + @classmethod + async def check_olds(cls, unpunishment_method: Callable, platform: Platform): + punishment_groups = cls._get_grouped_punishments(platform) + + now = datetime.datetime.now(datetime.timezone.utc) + for (_, _), sorted_punishments in punishment_groups: + if not (last_punishment := sorted_punishments[-1]).until or now < last_punishment.until: + continue + + if last_punishment.until + constants.PUNISHMENTS_RESET <= now: + for old_punishment in sorted_punishments: + old_punishment.delete() + + if last_punishment.is_active: + await last_punishment.unpunish(unpunishment_method) + last_punishment.is_active = False + last_punishment.save()