File size: 4,620 Bytes
3c05ff4
 
 
 
 
 
 
058dd88
 
 
 
3c05ff4
 
058dd88
3c05ff4
 
 
 
 
 
 
ea44a4e
3c05ff4
36b5e1c
 
058dd88
 
 
 
 
 
36b5e1c
058dd88
3c05ff4
 
 
 
058dd88
3c05ff4
058dd88
3c05ff4
 
 
 
ea44a4e
3c05ff4
36b5e1c
 
058dd88
 
 
36b5e1c
058dd88
3c05ff4
 
058dd88
3c05ff4
 
 
 
3580517
3c05ff4
ea44a4e
3c05ff4
058dd88
3c05ff4
 
3580517
 
 
 
3c05ff4
3580517
 
 
 
 
 
 
 
 
 
ea44a4e
3580517
ea44a4e
3580517
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c05ff4
 
 
058dd88
3c05ff4
058dd88
3c05ff4
 
 
 
058dd88
 
3c05ff4
 
 
 
 
 
 
 
058dd88
3c05ff4
 
 
 
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
124
125
126
127
from threading import RLock

from Powers import LOGGER
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
        curr = self.find_one({"chat_id": chat, "locktype": locktype})
        if curr:
            return False
        else:
            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:
                curr = self.find_one({"chat_id": chat, "locktype": i})
                if curr:
                    self.delete_one({"chat_id": chat, "locktype": i})
            return True
        curr = self.find_one({"chat_id": chat, "locktype": locktype})
        if curr:
            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
        else:
            if locktype != "all":
                curr = self.find_one(
                    {"chat_id": chat, "locktype": locktype})
                return bool(curr)
            else:
                to_return = {
                    "anti_channel": False,
                    "anti_fwd": {
                        "user": False,
                        "chat": False
                    },
                    "anti_links": False,
                    "bot": False
                }
                curr = self.find_all({"chat_id": chat})
                if not curr:
                    return None
                else:
                    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["anti_links"] == "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
        """
        curr = self.find_one({"chat_id": chat, "locktype": locktype})
        if curr:
            return True
        else:
            return False