Update punishments system

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)

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',

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 punishment in punishments:
-        for (_, _), sorted_punishments in punishment_groups:
+            now = datetime.datetime.now(datetime.timezone.utc)
-            if not (last_punishment := sorted_punishments[-1]).until or now < last_punishment.until:
+            if not punishment.until or now < punishment.until:
                 continue
 
-            if last_punishment.until + constants.PUNISHMENTS_RESET <= now:
+            await punishment.remove(unpunishment_method, delete=False)
-                for old_punishment in sorted_punishments:
+            if punishment.is_active:
-                    old_punishment.delete()
+                punishment.is_active = False
+                punishment.last_update = now
+                punishment.save()
 
-            if last_punishment.is_active:
+            if punishment.last_update + constants.PUNISHMENTS_RESET_TIME <= now:
-                await last_punishment.unpunish(unpunishment_method)
+                if punishment.level == 1:
-                last_punishment.is_active = False
+                    punishment.delete()
-                last_punishment.save()
+                else:
+                    punishment.level -= 1
+                    punishment.last_update = now
+                    punishment.save()