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""" try: 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'])): if limit != int(curr['limit']): return self.update( {"chat_id": chat_id,"within": within,"action": action}, {"limit": limit} ) elif within != int(curr['within']): return self.update( {"chat_id": chat_id,"limit": limit,"action": action}, {"within": within} ) elif action != str(curr['action']): return self.update( {"chat_id": chat_id,"limit": limit,"within": within}, {"action": action} ) 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, limit: int, within: int, action: str): with INSERTION_LOCK: curr = self.find_one({"chat_id": chat_id, "limit": limit, "within": within, "action": action}) if curr: return curr return "Flood haven't set" def rm_flood(self, chat_id: int, limit: int, within: int, action: str): with INSERTION_LOCK: curr = self.find_one({"chat_id": chat_id, "limit": limit, "within": within, "action": action}) if curr: self.delete_one(curr) return True return False except Exception as e: LOGGER.error(ef) LOGGER.error(format_exc())