diff --git a/flanabot/bots/flana_bot.py b/flanabot/bots/flana_bot.py index c81ba72..6e3495b 100644 --- a/flanabot/bots/flana_bot.py +++ b/flanabot/bots/flana_bot.py @@ -192,12 +192,12 @@ class FlanaBot(MultiBot, ABC): }) if len(last_2s_messages) >= constants.FLOOD_2s_LIMIT or len(last_7s_messages) >= constants.FLOOD_7s_LIMIT: - n_punishments = len(Punishment.find({ + punishment = Punishment.find_one({ 'platform': self.platform.value, 'user_id': message.author.id, 'group_id': message.chat.group_id - })) - punishment_seconds = (n_punishments + 2) ** constants.PUNISHMENT_INCREMENT_EXPONENT + }) + punishment_seconds = (getattr(punishment, 'level', 0) + 2) ** constants.PUNISHMENT_INCREMENT_EXPONENT try: await self.punish(message.author.id, message.chat.group_id, punishment_seconds, message) except BadRoleError as e: @@ -970,10 +970,10 @@ class FlanaBot(MultiBot, ABC): 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 | Message, time: int | datetime.timedelta, message: Message = None): + async def punish(self, user: int | str | User, group_: int | str | Chat | Message, time: int | datetime.timedelta = None, 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) + await punish.apply(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) @@ -1040,4 +1040,4 @@ class FlanaBot(MultiBot, ABC): 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) + await punish.remove(self._unpunish, message) diff --git a/flanabot/constants.py b/flanabot/constants.py index 2e51183..d0c90c2 100644 --- a/flanabot/constants.py +++ b/flanabot/constants.py @@ -1,7 +1,5 @@ import datetime -import flanautils - AUDIT_LOG_AGE = datetime.timedelta(hours=1) AUDIT_LOG_LIMIT = 5 AUTO_WEATHER_EVERY = datetime.timedelta(hours=6) @@ -12,9 +10,8 @@ HEAT_PERIOD_SECONDS = datetime.timedelta(minutes=15).total_seconds() INSULT_PROBABILITY = 0.00166666667 MAX_PLACE_QUERY_LENGTH = 50 PUNISHMENT_INCREMENT_EXPONENT = 6 -PUNISHMENTS_RESET = datetime.timedelta(weeks=6 * flanautils.WEEKS_IN_A_MONTH) +PUNISHMENTS_RESET_TIME = datetime.timedelta(weeks=2) RECOVERY_DELETED_MESSAGE_BEFORE = datetime.timedelta(hours=1) -SCRAPING_MESSAGE_WAITING_TIME = 0.1 BYE_PHRASES = ('Adiós.', 'adieu', 'adio', 'adioh', 'adios', 'adió', 'adiós', 'agur', 'bye', 'byyeeee', 'chao', 'hasta la vista', 'hasta luego', 'hasta nunca', ' hasta pronto', 'hasta la próxima', @@ -95,7 +92,7 @@ KEYWORDS = { 'cryptomoneda', 'cryptocurrency', 'currency', 'dinero', 'divisa', 'ethereum', 'inversion', 'moneda', 'pasta'), 'dice': ('dado', 'dice'), - 'poll': ('encuesta', 'poll', 'quiz'), + 'poll': ('encuesta', 'quiz'), 'punish': ('acaba', 'aprende', 'ataca', 'atalo', 'azota', 'boss', 'castiga', 'castigo', 'condena', 'controla', 'destroy', 'destroza', 'duro', 'ejecuta', 'enseña', 'escarmiento', 'execute', 'fuck', 'fusila', 'hell', 'humos', 'infierno', 'jefe', 'jode', 'learn', 'leccion', 'lesson', 'manda', 'purgatorio', 'sancion', diff --git a/flanabot/models/punishment.py b/flanabot/models/punishment.py index 2871645..3566815 100644 --- a/flanabot/models/punishment.py +++ b/flanabot/models/punishment.py @@ -2,31 +2,50 @@ __all__ = ['Punishment'] import datetime from dataclasses import dataclass -from typing import Callable +from typing import Any, Callable from multibot.models import Platform, PunishmentBase, db from flanabot import constants +from flanabot.models.message import Message @dataclass(eq=False) class Punishment(PunishmentBase): collection = db.punishment + level: int = 0 + + def _mongo_repr(self) -> Any: + self_vars = super()._mongo_repr() + self_vars['level'] = self.level + return self_vars + + async def apply(self, punishment_method: Callable, unpunishment_method: Callable, message: Message = None): + self.pull_from_database(overwrite_fields=('level',), exclude_fields=('until',)) + self.level += 1 + + await super().apply(punishment_method, unpunishment_method, message) + @classmethod async def check_olds(cls, unpunishment_method: Callable, platform: Platform): - punishment_groups = cls._get_grouped_punishments(platform) + punishments = cls.find({'platform': platform.value}) - 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: + for punishment in punishments: + now = datetime.datetime.now(datetime.timezone.utc) + if not punishment.until or now < punishment.until: continue - if last_punishment.until + constants.PUNISHMENTS_RESET <= now: - for old_punishment in sorted_punishments: - old_punishment.delete() + await punishment.remove(unpunishment_method, delete=False) + if punishment.is_active: + punishment.is_active = False + punishment.last_update = now + punishment.save() - if last_punishment.is_active: - await last_punishment.unpunish(unpunishment_method) - last_punishment.is_active = False - last_punishment.save() + if punishment.last_update + constants.PUNISHMENTS_RESET_TIME <= now: + if punishment.level == 1: + punishment.delete() + else: + punishment.level -= 1 + punishment.last_update = now + punishment.save()