diff --git a/flanabot/bots/penalty_bot.py b/flanabot/bots/penalty_bot.py index 0d229f8..1667b81 100644 --- a/flanabot/bots/penalty_bot.py +++ b/flanabot/bots/penalty_bot.py @@ -73,7 +73,7 @@ class PenaltyBot(MultiBot, ABC): 'group_id': message.chat.group_id }) punishment_seconds = (getattr(punishment, 'level', 0) + 2) ** constants.PUNISHMENT_INCREMENT_EXPONENT - await self.punish(message.author.id, message.chat.group_id, punishment_seconds, message) + await self.punish(message.author.id, message.chat.group_id, punishment_seconds, message, flood=True) await self.send(f'Castigado durante {TimeUnits(seconds=punishment_seconds).to_words()}.', message) async def _punish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): @@ -155,19 +155,11 @@ class PenaltyBot(MultiBot, ABC): if not punishment.until or now < punishment.until: continue - await self._remove_penalty(punishment, self._unpunish, delete=False) - if punishment.is_active: - punishment.is_active = False - punishment.last_update = now - punishment.save() + await self._remove_penalty(punishment, self._unpunish) 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() + punishment.level -= 1 + punishment.delete() async def is_punished(self, user: int | str | User, group_: int | str | Chat | Message) -> bool: pass @@ -177,16 +169,18 @@ class PenaltyBot(MultiBot, ABC): user: int | str | User, group_: int | str | Chat | Message, time: int | datetime.timedelta = None, - message: Message = None + message: Message = None, + flood=False ): # noinspection PyTypeChecker punishment = Punishment(self.platform, self.get_user_id(user), self.get_group_id(group_), time) punishment.pull_from_database(overwrite_fields=('level',), exclude_fields=('until',)) - punishment.level += 1 + if flood: + punishment.level += 1 await self._punish(punishment.user_id, punishment.group_id) punishment.save(pull_exclude_fields=('until',)) - await self._unpenalize_later(punishment, self._unpunish, message, delete=False) + await self._unpenalize_later(punishment, self._unpunish, message) async def unpunish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): # noinspection PyTypeChecker diff --git a/flanabot/models/punishment.py b/flanabot/models/punishment.py index b61ea93..3e2e95f 100644 --- a/flanabot/models/punishment.py +++ b/flanabot/models/punishment.py @@ -1,5 +1,6 @@ __all__ = ['Punishment'] +import datetime from dataclasses import dataclass from typing import Any @@ -16,3 +17,11 @@ class Punishment(Penalty): self_vars = super()._mongo_repr() self_vars['level'] = self.level return self_vars + + def delete(self, cascade=False): + if self.level == 0: + super().delete(cascade) + elif self.is_active: + self.is_active = False + self.last_update = datetime.datetime.now(datetime.timezone.utc) + self.save()