Spaces:
Sleeping
Sleeping
| from threading import RLock | |
| from Powers.database import MongoDB | |
| INSERTION_LOCK = RLock() | |
| lock_t = ["bot", "anti_c_send", "anti_fwd", | |
| "anti_fwd_u", "anti_fwd_c", "anti_links"] | |
| class LOCKS(MongoDB): | |
| """Class to store locks""" | |
| db_name = "locks" | |
| def __init__(self) -> None: | |
| super().__init__(self.db_name) | |
| def insert_lock_channel(self, chat: int, locktype: str): | |
| """ | |
| locktypes: all, bot, anti_c_send, anti_fwd, anti_fwd_u, anti_fwd_c, anti_links | |
| """ | |
| if locktype == "all": | |
| for i in lock_t: | |
| curr = self.find_one({"chat_id": chat, "locktype": i}) | |
| if curr: | |
| continue | |
| if i in ["anti_fwd_u", "anti_fwd_c"]: | |
| continue | |
| self.insert_one({"chat_id": chat, "locktype": i}) | |
| return True | |
| if curr := self.find_one({"chat_id": chat, "locktype": locktype}): | |
| return False | |
| with INSERTION_LOCK: | |
| hmm = self.merge_u_and_c(chat, locktype) | |
| if not hmm: | |
| self.insert_one({"chat_id": chat, "locktype": locktype}) | |
| return True | |
| def remove_lock_channel(self, chat: int, locktype: str): | |
| """ | |
| locktypes: all, bot, anti_c_send, anti_fwd, anti_fwd_u, anti_fwd_c, anti_links | |
| """ | |
| if locktype == "all": | |
| for i in lock_t: | |
| if curr := self.find_one({"chat_id": chat, "locktype": i}): | |
| self.delete_one({"chat_id": chat, "locktype": i}) | |
| return True | |
| if curr := self.find_one({"chat_id": chat, "locktype": locktype}): | |
| with INSERTION_LOCK: | |
| self.delete_one({"chat_id": chat, "locktype": locktype}) | |
| return True | |
| else: | |
| return False | |
| def get_lock_channel(self, chat: int, locktype: str = "all"): | |
| """ | |
| locktypes: anti_c_send, anti_fwd, anti_fwd_u, anti_fwd_c, anti_links, bot | |
| """ | |
| if locktype not in [ | |
| "anti_c_send", | |
| "anti_fwd", | |
| "anti_fwd_u", | |
| "anti_fwd_c", | |
| "anti_links", | |
| "bot", | |
| "all", | |
| ]: | |
| return False | |
| if locktype != "all": | |
| curr = self.find_one( | |
| {"chat_id": chat, "locktype": locktype}) | |
| return bool(curr) | |
| else: | |
| if not (curr := self.find_all({"chat_id": chat})): | |
| return None | |
| to_return = { | |
| "anti_channel": False, | |
| "anti_fwd": { | |
| "user": False, | |
| "chat": False | |
| }, | |
| "anti_links": False, | |
| "bot": False | |
| } | |
| for i in list(curr): | |
| if i["locktype"] == "anti_c_send": | |
| to_return["anti_channel"] = True | |
| elif i["locktype"] == "anti_fwd": | |
| to_return["anti_fwd"]["user"] = to_return["anti_fwd"]["chat"] = True | |
| elif i["locktype"] == "anti_fwd_u": | |
| to_return["anti_fwd"]["user"] = True | |
| elif i["locktype"] == "anti_fwd_c": | |
| to_return["anti_fwd"]["chat"] = True | |
| elif i["locktype"] == "anti_links": | |
| to_return["anti_links"] = True | |
| elif i["locktype"] == "bot": | |
| to_return["bot"] = True | |
| else: | |
| continue | |
| return to_return | |
| def merge_u_and_c(self, chat: int, locktype: str): | |
| if locktype == "anti_fwd_u": | |
| curr = self.find_one({"chat_id": chat, "locktype": "anti_fwd_c"}) | |
| elif locktype == "anti_fwd_c": | |
| curr = self.find_one({"chat_id": chat, "locktype": "anti_fwd_u"}) | |
| else: | |
| return False | |
| if curr: | |
| self.delete_one({"chat_id": chat, "locktype": locktype}) | |
| self.insert_one({"chat_id": chat, "locktype": "anti_fwd"}) | |
| return True | |
| else: | |
| return False | |
| def is_particular_lock(self, chat: int, locktype: str): | |
| """ | |
| locktypes: anti_c_send, anti_fwd, anti_fwd_u, anti_fwd_c, anti_links | |
| """ | |
| return bool(curr := self.find_one({"chat_id": chat, "locktype": locktype})) | |