File size: 4,620 Bytes
ea508b7
 
 
 
 
 
 
51e1c40
 
 
 
ea508b7
 
51e1c40
ea508b7
 
 
 
 
 
 
518de8c
ea508b7
bb1bfd1
 
51e1c40
 
 
 
 
 
bb1bfd1
51e1c40
ea508b7
 
 
 
51e1c40
ea508b7
51e1c40
ea508b7
 
 
 
518de8c
ea508b7
bb1bfd1
 
51e1c40
 
 
bb1bfd1
51e1c40
ea508b7
 
51e1c40
ea508b7
 
 
 
6dcea66
ea508b7
518de8c
ea508b7
51e1c40
ea508b7
 
6dcea66
 
 
 
ea508b7
6dcea66
 
 
 
 
 
 
 
 
 
518de8c
6dcea66
518de8c
6dcea66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ea508b7
 
 
51e1c40
ea508b7
51e1c40
ea508b7
 
 
 
51e1c40
 
ea508b7
 
 
 
 
 
 
 
51e1c40
ea508b7
 
 
 
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