import time from peewee import CharField, IntegerField, CompositeKey from defence360agent.contracts.config import IContactMessageType from defence360agent.model import Model, instance from defence360agent.utils.common import DAY, WEEK THROTTLING_PERIOD = { IContactMessageType.MALWARE_FOUND: DAY, IContactMessageType.SCAN_NOT_SCHEDULED: WEEK, } class IContactThrottle(Model): class Meta: database = instance.db db_table = "icontact_throttle" primary_key = CompositeKey("message_type", "user") message_type = CharField() user = CharField(null=True) #: The last time we sent a notification about :attr:`message_type` timestamp = IntegerField(default=0) @classmethod def may_be_notified(cls, message_type, period_limit, user=None): obj, _ = cls.get_or_create(message_type=message_type, user=user) return (time.time() - obj.timestamp) > period_limit @classmethod def refresh(cls, message_type, user=None): cls.update(timestamp=time.time()).where( cls.message_type == message_type, cls.user.is_null(True) if user is None else cls.user == user, ).execute()