Powers/database/flood_db.py

from threading import RLock
from traceback import format_exc

from Powers import LOGGER
from Powers.database import MongoDB
from Powers.utils.msg_types import Types

INSERTION_LOCK = RLock()

class Floods(MongoDB):
    """Class to store flood limit and action of a chat"""
    db_name = "flood"

    def __init__(self):
        super().__init__(self.db_name)

    def save_flood(
        self,
        chat_id: int,
        limit: int,
        within: int,
        action: str,
    ):
        with INSERTION_LOCK:
            curr = self.find_one({"chat_id": chat_id})
            if curr:
                if not(limit == int(curr['limit']) and within == int(curr['within']) and action == str(curr['action'])):
                    return self.update(
                        {
                            "chat_id": chat_id
                        },
                        {
                            "limit": limit,
                            "within": within,
                            "action": action,
                        }
                    )
                else:
                    return 
            else:
                return self.insert_one(
                    {
                        "chat_id" : chat_id,
                        "limit": limit,
                        "within": within,
                        "action" : action
                    },
                )
    
    def is_chat(self, chat_id: int):
        with INSERTION_LOCK:
            curr = self.find_one({"chat_id": chat_id})
            if curr:
                action = [str(curr['limit']), str(curr['within']), str(curr['action'])]
                return action
            return False
    
    def get_action(self, chat_id: int):
        with INSERTION_LOCK:
            curr = self.find_one({"chat_id": chat_id})
            if curr:
                return curr['action']
            return "Flood haven't set"
    
    def rm_flood(self, chat_id: int):
        with INSERTION_LOCK:
            curr = self.find_one({"chat_id": chat_id})
            if curr:
                self.delete_one(curr)
                return True
            return False