File size: 4,333 Bytes
ca4eb6d
 
11ae35a
 
ca4eb6d
 
 
 
 
 
 
 
 
 
 
 
 
6cef7ec
 
 
 
 
 
ca4eb6d
 
89ad488
6cef7ec
89ad488
ca4eb6d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89ad488
6cef7ec
89ad488
ca4eb6d
 
 
 
 
 
 
 
 
89ad488
ca4eb6d
 
 
89ad488
6cef7ec
89ad488
ca4eb6d
 
 
 
 
 
 
 
 
 
89ad488
ca4eb6d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89ad488
ca4eb6d
 
 
 
 
 
 
 
 
 
 
 
89ad488
ca4eb6d
 
 
 
89ad488
ca4eb6d
 
 
 
51e1c40
aef1161
51e1c40
aef1161
ca4eb6d
 
 
 
 
 
 
 
 
89ad488
ca4eb6d
 
 
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
128
129
130
131
132
133
134
from hashlib import md5
from threading import RLock
from time import time

from Powers.database import MongoDB
from Powers.utils.msg_types import Types

INSERTION_LOCK = RLock()


class Notes(MongoDB):
    db_name = "notes"

    def __init__(self) -> None:
        super().__init__(self.db_name)

    def save_note(
            self,
            chat_id: int,
            note_name: str,
            note_value: str,
            msgtype: int = Types.TEXT,
            fileid="",
    ):
        with INSERTION_LOCK:
            if curr := self.find_one(
                    {"chat_id": chat_id, "note_name": note_name},
            ):
                return False
            hash_gen = md5(
                (note_name + note_value + str(chat_id) + str(int(time()))).encode(),
            ).hexdigest()
            return self.insert_one(
                {
                    "chat_id": chat_id,
                    "note_name": note_name,
                    "note_value": note_value,
                    "hash": hash_gen,
                    "msgtype": msgtype,
                    "fileid": fileid,
                },
            )

    def get_note(self, chat_id: int, note_name: str):
        with INSERTION_LOCK:
            if curr := self.find_one(
                    {"chat_id": chat_id, "note_name": note_name},
            ):
                return curr
            return "Note does not exist!"

    def get_note_by_hash(self, note_hash: str):
        return self.find_one({"hash": note_hash})

    def get_all_notes(self, chat_id: int):
        with INSERTION_LOCK:
            curr = self.find_all({"chat_id": chat_id})
            return sorted([(note["note_name"], note["hash"]) for note in curr])

    def rm_note(self, chat_id: int, note_name: str):
        with INSERTION_LOCK:
            if curr := self.find_one(
                    {"chat_id": chat_id, "note_name": note_name},
            ):
                self.delete_one(curr)
                return True
            return False

    def rm_all_notes(self, chat_id: int):
        with INSERTION_LOCK:
            return self.delete_one({"chat_id": chat_id})

    def count_notes(self, chat_id: int):
        with INSERTION_LOCK:
            return len(curr) if (curr := self.find_all({"chat_id": chat_id})) else 0

    def count_notes_chats(self):
        with INSERTION_LOCK:
            notes = self.find_all()
            chats_ids = [chat["chat_id"] for chat in notes]
            return len(set(chats_ids))

    def count_all_notes(self):
        with INSERTION_LOCK:
            return self.count()

    def count_notes_type(self, ntype):
        with INSERTION_LOCK:
            return self.count({"msgtype": ntype})

    # Migrate if chat id changes!
    def migrate_chat(self, old_chat_id: int, new_chat_id: int):
        with INSERTION_LOCK:
            if old_chat_db := self.find_one({"_id": old_chat_id}):
                new_data = old_chat_db.update({"_id": new_chat_id})
                self.delete_one({"_id": old_chat_id})
                self.insert_one(new_data)


class NotesSettings(MongoDB):
    db_name = "notes_settings"

    def __init__(self) -> None:
        super().__init__(self.db_name)

    def set_privatenotes(self, chat_id: int, status: bool = False):
        if curr := self.find_one({"_id": chat_id}):
            return self.update({"_id": chat_id}, {"privatenotes": status})
        return self.insert_one({"_id": chat_id, "privatenotes": status})

    def get_privatenotes(self, chat_id: int):
        if curr := self.find_one({"_id": chat_id}):
            return curr["privatenotes"]
        self.update({"_id": chat_id}, {"privatenotes": False})
        return False

    def clean_notes(self, chat_id):
        with INSERTION_LOCK:
            return self.delete_one({"_id": chat_id})

    def list_chats(self):
        return self.find_all({"privatenotes": True})

    def count_chats(self):
        return len(self.find_all({"privatenotes": True}))

    # Migrate if chat id changes!
    def migrate_chat(self, old_chat_id: int, new_chat_id: int):
        with INSERTION_LOCK:
            if old_chat_db := self.find_one({"_id": old_chat_id}):
                new_data = old_chat_db.update({"_id": new_chat_id})
                self.delete_one({"_id": old_chat_id})
                self.insert_one(new_data)