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})) | |