File size: 4,297 Bytes
ea508b7
 
 
 
 
 
51e1c40
 
 
 
ea508b7
 
51e1c40
ea508b7
 
 
 
 
 
 
518de8c
ea508b7
bb1bfd1
 
51e1c40
 
 
 
 
 
bb1bfd1
89ad488
ea508b7
89ad488
 
 
 
 
ea508b7
 
 
518de8c
ea508b7
bb1bfd1
 
89ad488
51e1c40
bb1bfd1
89ad488
ea508b7
51e1c40
ea508b7
 
 
 
6dcea66
ea508b7
518de8c
ea508b7
89ad488
 
 
 
 
 
 
 
 
ea508b7
89ad488
 
 
 
ea508b7
6cef7ec
89ad488
6cef7ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87ff2e2
6cef7ec
 
 
 
 
 
ea508b7
 
 
51e1c40
ea508b7
51e1c40
ea508b7
 
 
 
51e1c40
 
ea508b7
 
 
 
 
 
 
 
89ad488
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}))