File size: 4,297 Bytes
3c05ff4
 
 
 
 
 
058dd88
 
 
 
3c05ff4
 
058dd88
3c05ff4
 
 
 
 
 
 
ea44a4e
3c05ff4
36b5e1c
 
058dd88
 
 
 
 
 
36b5e1c
3c9449b
3c05ff4
3c9449b
 
 
 
 
3c05ff4
 
 
ea44a4e
3c05ff4
36b5e1c
 
3c9449b
058dd88
36b5e1c
3c9449b
3c05ff4
058dd88
3c05ff4
 
 
 
3580517
3c05ff4
ea44a4e
3c05ff4
3c9449b
 
 
 
 
 
 
 
 
3c05ff4
3c9449b
 
 
 
3c05ff4
4056428
3c9449b
4056428
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83c52c0
4056428
 
 
 
 
 
3c05ff4
 
 
058dd88
3c05ff4
058dd88
3c05ff4
 
 
 
058dd88
 
3c05ff4
 
 
 
 
 
 
 
3c9449b
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
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}))