Spaces:
Sleeping
Sleeping
File size: 4,800 Bytes
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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
from threading import RLock
from time import time
from Powers import LOGGER
from Powers.database import MongoDB
INSERTION_LOCK = RLock()
class Blacklist(MongoDB):
"""Class to manage database for blacklists for chats."""
# Database name to connect to to preform operations
db_name = "blacklists"
def __init__(self, chat_id: int) -> None:
super().__init__(self.db_name)
self.chat_id = chat_id
self.chat_info = self.__ensure_in_db()
def check_word_blacklist_status(self, word: str):
with INSERTION_LOCK:
bl_words = self.chat_info["triggers"]
return bool(word in bl_words)
def add_blacklist(self, trigger: str):
with INSERTION_LOCK:
if not self.check_word_blacklist_status(trigger):
return self.update(
{"_id": self.chat_id},
{
"_id": self.chat_id,
"triggers": self.chat_info["triggers"] + [trigger],
},
)
def remove_blacklist(self, trigger: str):
with INSERTION_LOCK:
if self.check_word_blacklist_status(trigger):
self.chat_info["triggers"].remove(trigger)
return self.update(
{"_id": self.chat_id},
{
"_id": self.chat_id,
"triggers": self.chat_info["triggers"],
},
)
def get_blacklists(self):
with INSERTION_LOCK:
return self.chat_info["triggers"]
@staticmethod
def count_blacklists_all():
with INSERTION_LOCK:
collection = MongoDB(Blacklist.db_name)
curr = collection.find_all()
return sum(len(chat["triggers"]) for chat in curr)
@staticmethod
def count_blackists_chats():
with INSERTION_LOCK:
collection = MongoDB(Blacklist.db_name)
curr = collection.find_all()
return sum(1 for chat in curr if chat["triggers"])
def set_action(self, action: str):
with INSERTION_LOCK:
return self.update(
{"_id": self.chat_id},
{"_id": self.chat_id, "action": action},
)
def get_action(self):
with INSERTION_LOCK:
return self.chat_info["action"]
def set_reason(self, reason: str):
with INSERTION_LOCK:
return self.update(
{"_id": self.chat_id},
{"_id": self.chat_id, "reason": reason},
)
def get_reason(self):
with INSERTION_LOCK:
return self.chat_info["reason"]
@staticmethod
def count_action_bl_all(action: str):
with INSERTION_LOCK:
collection = MongoDB(Blacklist.db_name)
all_data = collection.find_all({"action": action})
return sum(len(i["triggers"]) >= 1 for i in all_data)
def rm_all_blacklist(self):
with INSERTION_LOCK:
return self.update(
{"_id": self.chat_id},
{"triggers": []},
)
def __ensure_in_db(self):
chat_data = self.find_one({"_id": self.chat_id})
if not chat_data:
new_data = new_data = {
"_id": self.chat_id,
"triggers": [],
"action": "none",
"reason": "Automated blacklisted word: {{}}",
}
self.insert_one(new_data)
LOGGER.info(f"Initialized Blacklist Document for chat {self.chat_id}")
return new_data
return chat_data
# Migrate if chat id changes!
def migrate_chat(self, new_chat_id: int):
old_chat_db = self.find_one({"_id": self.chat_id})
new_data = old_chat_db.update({"_id": new_chat_id})
self.insert_one(new_data)
self.delete_one({"_id": self.chat_id})
@staticmethod
def repair_db(collection):
all_data = collection.find_all()
keys = {
"triggers": [],
"action": "none",
"reason": "Automated blacklisted word: {{}}",
}
for data in all_data:
for key, val in keys.items():
try:
_ = data[key]
except KeyError:
LOGGER.warning(
f"Repairing Blacklist Database - setting '{key}:{val}' for {data['_id']}",
)
collection.update({"_id": data["_id"]}, {key: val})
def __pre_req_blacklists():
start = time()
LOGGER.info("Starting Blacklists Database Repair...")
collection = MongoDB(Blacklist.db_name)
Blacklist.repair_db(collection)
LOGGER.info(f"Done in {round((time() - start), 3)}s!")
|