class GBan:
"""Class for managing Gbans in bot."""
def __init__(self) -> None:
self.collection = MongoDB("gbans")
def check_gban(self, user_id: int):
with INSERTION_LOCK:
if user_id in ANTISPAM_BANNED:
return True
return bool(self.collection.find_one({"_id": user_id}))
def add_gban(self, user_id: int, reason: str, by_user: int):
global ANTISPAM_BANNED
with INSERTION_LOCK:
# Check if user is already gbanned or not
if self.collection.find_one({"_id": user_id}):
return self.update_gban_reason(user_id, reason)
# If not already gbanned, then add to gban
time_rn = datetime.now()
ANTISPAM_BANNED.add(user_id)
return self.collection.insert_one(
{
"_id": user_id,
"reason": reason,
"by": by_user,
"time": time_rn,
}, )
def remove_gban(self, user_id: int):
global ANTISPAM_BANNED
with INSERTION_LOCK:
# Check if user is already gbanned or not
if self.collection.find_one({"_id": user_id}):
ANTISPAM_BANNED.remove(user_id)
return self.collection.delete_one({"_id": user_id})
return "User not gbanned!"
def get_gban(self, user_id: int):
if self.check_gban(user_id):
curr = self.collection.find_one({"_id": user_id})
if curr:
return True, curr["reason"]
return False, ""
def update_gban_reason(self, user_id: int, reason: str):
with INSERTION_LOCK:
return self.collection.update(
{"_id": user_id},
{"reason": reason},
)
def count_gbans(self):
with INSERTION_LOCK:
try:
return len(ANTISPAM_BANNED)
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count()
def load_from_db(self):
with INSERTION_LOCK:
return self.collection.find_all()
def list_gbans(self):
with INSERTION_LOCK:
try:
return list(ANTISPAM_BANNED)
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.find_all()
class Notes:
def __init__(self) -> None:
self.collection = MongoDB("notes")
def save_note(
self,
chat_id: int,
note_name: str,
note_value: str,
msgtype: int = Types.TEXT,
fileid="",
):
with INSERTION_LOCK:
curr = self.collection.find_one(
{"chat_id": chat_id, "note_name": note_name},
)
if curr:
return False
hash_gen = md5(
(note_name + note_value + str(chat_id) + str(int(time()))).encode(),
).hexdigest()
return self.collection.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:
curr = self.collection.find_one(
{"chat_id": chat_id, "note_name": note_name},
)
if curr:
return curr
return "Note does not exist!"
def get_note_by_hash(self, note_hash: str):
return self.collection.find_one({"hash": note_hash})
def get_all_notes(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_all({"chat_id": chat_id})
note_list = []
for note in curr:
note_list.append((note["note_name"], note["hash"]))
note_list.sort()
return note_list
def rm_note(self, chat_id: int, note_name: str):
with INSERTION_LOCK:
curr = self.collection.find_one(
{"chat_id": chat_id, "note_name": note_name},
)
if curr:
self.collection.delete_one(curr)
return True
return False
def rm_all_notes(self, chat_id: int):
with INSERTION_LOCK:
return self.collection.delete_one({"chat_id": chat_id})
def count_notes(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_all({"chat_id": chat_id})
if curr:
return len(curr)
return 0
def count_notes_chats(self):
with INSERTION_LOCK:
notes = self.collection.find_all()
chats_ids = []
for chat in notes:
chats_ids.append(chat["chat_id"])
return len(set(chats_ids))
def count_all_notes(self):
with INSERTION_LOCK:
return self.collection.count()
def count_notes_type(self, ntype):
with INSERTION_LOCK:
return self.collection.count({"msgtype": ntype})
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
with INSERTION_LOCK:
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
def count_grouprules_chats():
with INSERTION_LOCK:
collection = MongoDB(Rules.db_name)
return collection.count({"privrules": False})
class Filters:
def __init__(self) -> None:
self.collection = MongoDB("chat_filters")
def save_filter(
self,
chat_id: int,
keyword: str,
filter_reply: str,
msgtype: int = Types.TEXT,
fileid="",
):
global FILTER_CACHE
with INSERTION_LOCK:
# local dict update
try:
curr_filters = FILTER_CACHE[chat_id]
except KeyError:
curr_filters = []
try:
keywords = {i["keyword"] for i in curr_filters}
except KeyError:
keywords = set()
if keyword not in keywords:
curr_filters.append(
{
"chat_id": chat_id,
"keyword": keyword,
"filter_reply": filter_reply,
"msgtype": msgtype,
"fileid": fileid,
}, )
FILTER_CACHE[chat_id] = curr_filters
# Database update
curr = self.collection.find_one(
{
"chat_id": chat_id,
"keyword": keyword
}, )
if curr:
return False
return self.collection.insert_one(
{
"chat_id": chat_id,
"keyword": keyword,
"filter_reply": filter_reply,
"msgtype": msgtype,
"fileid": fileid,
}, )
def get_filter(self, chat_id: int, keyword: str):
with INSERTION_LOCK:
try:
curr = next(i for i in FILTER_CACHE[chat_id]
if keyword in i["keyword"].split("|"))
if curr:
return curr
except (KeyError, StopIteration):
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_one(
{
"chat_id": chat_id,
"keyword": {
"$regex": fr"\|?{keyword}\|?"
}
}, )
if curr:
return curr
return "Filter does not exist!"
def get_all_filters(self, chat_id: int):
with INSERTION_LOCK:
try:
return [i["keyword"] for i in FILTER_CACHE[chat_id]]
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_all({"chat_id": chat_id})
if curr:
filter_list = {i["keyword"] for i in curr}
return list(filter_list)
return []
def rm_filter(self, chat_id: int, keyword: str):
global FILTER_CACHE
with INSERTION_LOCK:
try:
FILTER_CACHE[chat_id].remove(
next(i for i in FILTER_CACHE[chat_id]
if keyword in i["keyword"].split("|")), )
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_one(
{
"chat_id": chat_id,
"keyword": {
"$regex": fr"\|?{keyword}\|?"
}
}, )
if curr:
self.collection.delete_one(curr)
return True
return False
def rm_all_filters(self, chat_id: int):
global FILTER_CACHE
with INSERTION_LOCK:
try:
del FILTER_CACHE[chat_id]
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.delete_one({"chat_id": chat_id})
def count_filters_all(self):
with INSERTION_LOCK:
try:
return len([
j for i in ((i["keyword"] for i in FILTER_CACHE[chat_id])
for chat_id in set(FILTER_CACHE.keys()))
for j in i
], )
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_all()
if curr:
return len(curr)
return 0
def count_filter_aliases(self):
with INSERTION_LOCK:
try:
return len([
i for j in [
j.split("|")
for i in ((i["keyword"] for i in FILTER_CACHE[chat_id])
for chat_id in set(FILTER_CACHE.keys()))
for j in i if len(j.split("|")) >= 2
] for i in j
], )
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_all()
if curr:
return len([
z for z in (i["keyword"].split("|")
for i in curr) if len(z) >= 2
], )
return 0
def count_filters_chats(self):
with INSERTION_LOCK:
try:
return len(set(FILTER_CACHE.keys()))
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
filters = self.collection.find_all()
chats_ids = {i["chat_id"] for i in filters}
return len(chats_ids)
def count_all_filters(self):
with INSERTION_LOCK:
try:
return len([(i["keyword"] for i in FILTER_CACHE[chat_id])
for chat_id in set(FILTER_CACHE.keys())], )
except KeyError:
pass
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count()
def count_filter_type(self, ntype):
with INSERTION_LOCK:
return self.collection.count({"msgtype": ntype})
def load_from_db(self):
with INSERTION_LOCK:
return self.collection.find_all()
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
with INSERTION_LOCK:
# Update locally
try:
old_db_local = FILTER_CACHE[old_chat_id]
del FILTER_CACHE[old_chat_id]
FILTER_CACHE[new_chat_id] = old_db_local
except KeyError:
pass
# Update in db
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
class Rules:
"""Class for rules for chats in bot."""
def __init__(self) -> None:
self.collection = MongoDB("rules")
def get_rules(self, chat_id: int):
global RULES_CACHE
with INSERTION_LOCK:
if (chat_id in set(
RULES_CACHE.keys())) and (RULES_CACHE[chat_id]["rules"]):
return RULES_CACHE[chat_id]["rules"]
rules = self.collection.find_one({"_id": chat_id})
if rules:
return rules["rules"]
return None
def set_rules(self, chat_id: int, rules: str, privrules: bool = False):
global RULES_CACHE
with INSERTION_LOCK:
if chat_id in set(RULES_CACHE.keys()):
RULES_CACHE[chat_id]["rules"] = rules
curr_rules = self.collection.find_one({"_id": chat_id})
if curr_rules:
return self.collection.update(
{"_id": chat_id},
{"rules": rules},
)
RULES_CACHE[chat_id] = {"rules": rules, "privrules": privrules}
return self.collection.insert_one(
{
"_id": chat_id,
"rules": rules,
"privrules": privrules
}, )
def get_privrules(self, chat_id: int):
global RULES_CACHE
with INSERTION_LOCK:
if (chat_id in set(RULES_CACHE.keys())) and (
RULES_CACHE[chat_id]["privrules"]):
return RULES_CACHE[chat_id]["privrules"]
curr_rules = self.collection.find_one({"_id": chat_id})
if curr_rules:
return curr_rules["privrules"]
RULES_CACHE[chat_id] = {"privrules": False, "rules": ""}
return self.collection.insert_one(
{
"_id": chat_id,
"rules": "",
"privrules": False
}, )
def set_privrules(self, chat_id: int, privrules: bool):
global RULES_CACHE
with INSERTION_LOCK:
if chat_id in set(RULES_CACHE.keys()):
RULES_CACHE[chat_id]["privrules"] = privrules
curr_rules = self.collection.find_one({"_id": chat_id})
if curr_rules:
return self.collection.update(
{"_id": chat_id},
{"privrules": privrules},
)
RULES_CACHE[chat_id] = {"rules": "", "privrules": privrules}
return self.collection.insert_one(
{
"_id": chat_id,
"rules": "",
"privrules": privrules
}, )
def clear_rules(self, chat_id: int):
global RULES_CACHE
with INSERTION_LOCK:
if chat_id in set(RULES_CACHE.keys()):
del RULES_CACHE[chat_id]
curr_rules = self.collection.find_one({"_id": chat_id})
if curr_rules:
return self.collection.delete_one({"_id": chat_id})
return "Rules not found!"
def count_chats(self):
with INSERTION_LOCK:
try:
return len([i for i in RULES_CACHE if RULES_CACHE[i]["rules"]])
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count({"rules": {"$regex": ".*"}})
def count_privrules_chats(self):
with INSERTION_LOCK:
try:
return len(
[i for i in RULES_CACHE if RULES_CACHE[i]["privrules"]])
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count({"privrules": True})
def count_grouprules_chats(self):
with INSERTION_LOCK:
try:
return len([
i for i in RULES_CACHE if not RULES_CACHE[i]["privrules"]
])
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count({"privrules": True})
def load_from_db(self):
return self.collection.find_all()
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
global RULES_CACHE
with INSERTION_LOCK:
# Update locally
try:
old_db_local = RULES_CACHE[old_chat_id]
del RULES_CACHE[old_chat_id]
RULES_CACHE[new_chat_id] = old_db_local
except KeyError:
pass
# Update in db
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
def count_chats_with_rules():
with INSERTION_LOCK:
collection = MongoDB(Rules.db_name)
return collection.count({"rules": {"$regex": ".*"}})
class Chats:
"""Class to manage users for bot."""
def __init__(self) -> None:
self.collection = MongoDB("chats")
def remove_chat(self, chat_id: int):
with INSERTION_LOCK:
self.collection.delete_one({"_id": chat_id})
def update_chat(self, chat_id: int, chat_name: str, user_id: int):
global CHATS_CACHE
with INSERTION_LOCK:
# Local Cache
try:
chat = CHATS_CACHE[chat_id]
users_old = chat["users"]
if user_id in set(users_old):
# If user_id already exists, return
return "user already exists in chat users"
users_old.append(user_id)
users = list(set(users_old))
CHATS_CACHE[chat_id] = {
"chat_name": chat_name,
"users": users,
}
except KeyError:
pass
# Databse Cache
curr = self.collection.find_one({"_id": chat_id})
if curr:
users_old = curr["users"]
users_old.append(user_id)
users = list(set(users_old))
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"chat_name": chat_name,
"users": users,
},
)
CHATS_CACHE[chat_id] = {
"chat_name": chat_name,
"users": [user_id],
}
return self.collection.insert_one(
{
"_id": chat_id,
"chat_name": chat_name,
"users": [user_id],
}, )
def count_chat_users(self, chat_id: int):
with INSERTION_LOCK:
try:
return len(CHATS_CACHE[chat_id]["users"])
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_one({"_id": chat_id})
if curr:
return len(curr["users"])
return 0
def chat_members(self, chat_id: int):
with INSERTION_LOCK:
try:
return CHATS_CACHE[chat_id]["users"]
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["users"]
return []
def count_chats(self):
with INSERTION_LOCK:
try:
return len(CHATS_CACHE)
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count() or 0
def list_chats(self):
with INSERTION_LOCK:
try:
return list(CHATS_CACHE.keys())
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
chats = self.collection.find_all()
chat_list = {i["_id"] for i in chats}
return list(chat_list)
def get_all_chats(self):
with INSERTION_LOCK:
try:
return CHATS_CACHE
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.find_all()
def get_chat_info(self, chat_id: int):
with INSERTION_LOCK:
try:
return CHATS_CACHE[chat_id]
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.find_one({"_id": chat_id})
def load_from_db(self):
with INSERTION_LOCK:
return self.collection.find_all()
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
global CHATS_CACHE
with INSERTION_LOCK:
# Update locally
try:
old_db_local = CHATS_CACHE[old_chat_id]
del CHATS_CACHE[old_chat_id]
CHATS_CACHE[new_chat_id] = old_db_local
except KeyError:
pass
# Update in db
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
class Pins:
"""Class for managing antichannelpins in chats."""
def __init__(self) -> None:
self.collection = MongoDB("antichannelpin")
def check_status(self, chat_id: int, atype: str):
with INSERTION_LOCK:
if chat_id in (PINS_CACHE[atype]):
return True
curr = self.collection.find_one({"_id": chat_id, atype: True})
if curr:
return True
return False
def get_current_stngs(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr
curr = {
"_id": chat_id,
"antichannelpin": False,
"cleanlinked": False
}
self.collection.insert_one(curr)
return curr
def set_on(self, chat_id: int, atype: str):
global PINS_CACHE
with INSERTION_LOCK:
otype = "cleanlinked" if atype == "antichannelpin" else "antichannelpin"
if chat_id not in (PINS_CACHE[atype]):
(PINS_CACHE[atype]).add(chat_id)
try:
return self.collection.insert_one(
{
"_id": chat_id,
atype: True,
otype: False
}, )
except DuplicateKeyError:
return self.collection.update(
{"_id": chat_id},
{
atype: True,
otype: False
},
)
return "Already exists"
def set_off(self, chat_id: int, atype: str):
global PINS_CACHE
with INSERTION_LOCK:
if chat_id in (PINS_CACHE[atype]):
(PINS_CACHE[atype]).remove(chat_id)
return self.collection.update({"_id": chat_id}, {atype: False})
return f"{atype} not enabled"
def count_chats(self, atype):
with INSERTION_LOCK:
try:
return len(PINS_CACHE[atype])
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.count({atype: True})
def load_chats_from_db(self, query=None):
with INSERTION_LOCK:
if query is None:
query = {}
return self.collection.find_all(query)
def list_chats(self, query):
with INSERTION_LOCK:
try:
return PINS_CACHE[query]
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return self.collection.find_all({query: True})
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
global PINS_CACHE
with INSERTION_LOCK:
# Update locally
if old_chat_id in (PINS_CACHE["antichannelpin"]):
(PINS_CACHE["antichannelpin"]).remove(old_chat_id)
(PINS_CACHE["antichannelpin"]).add(new_chat_id)
if old_chat_id in (PINS_CACHE["cleanlinked"]):
(PINS_CACHE["cleanlinked"]).remove(old_chat_id)
(PINS_CACHE["cleanlinked"]).add(new_chat_id)
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
class Approve:
"""Class for managing Approves in Chats in Bot."""
def __init__(self) -> None:
self.collection = MongoDB("approve")
def check_approve(self, chat_id: int, user_id: int):
with INSERTION_LOCK:
try:
users = list(APPROVE_CACHE[chat_id])
except KeyError:
return True
if user_id in {i[0] for i in users}:
return True
curr_approve = self.collection.find_one({"_id": chat_id}, )
if curr_approve:
try:
return next(user for user in curr_approve["users"]
if user[0] == user_id)
except StopIteration:
return False
return False
def add_approve(self, chat_id: int, user_id: int, user_name: str):
global APPROVE_CACHE
with INSERTION_LOCK:
try:
users = list(APPROVE_CACHE[chat_id])
except KeyError:
return True
if user_id in {i[0] for i in users}:
return True
users_old = APPROVE_CACHE[chat_id]
users_old.add((user_id, user_name))
APPROVE_CACHE[chat_id] = users_old
curr = self.collection.find_one({"_id": chat_id})
if curr:
users_old = curr["users"]
users_old.append((user_id, user_name))
users = list(set(users_old)) # Remove duplicates
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"users": users,
},
)
APPROVE_CACHE[chat_id] = {user_id, user_name}
return self.collection.insert_one(
{
"_id": chat_id,
"users": [(user_id, user_name)],
}, )
def remove_approve(self, chat_id: int, user_id: int):
global APPROVE_CACHE
with INSERTION_LOCK:
try:
users = list(APPROVE_CACHE[chat_id])
except KeyError:
return True
try:
user = next(user for user in users if user[0] == user_id)
users.remove(user)
ADMIN_CACHE[chat_id] = users
except StopIteration:
pass
curr = self.collection.find_one({"_id": chat_id})
if curr:
users = curr["users"]
try:
user = next(user for user in users if user[0] == user_id)
except StopIteration:
return "Not Approved"
users.remove(user)
# If the list is emptied, then delete it
if not users:
return self.collection.delete_one({"_id": chat_id}, )
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"users": users,
},
)
return "Not approved"
def unapprove_all(self, chat_id: int):
global APPROVE_CACHE
with INSERTION_LOCK:
del APPROVE_CACHE[chat_id]
return self.collection.delete_one({"_id": chat_id}, )
def list_approved(self, chat_id: int):
with INSERTION_LOCK:
try:
return APPROVE_CACHE[chat_id]
except KeyError:
pass
except Exception as ef:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["users"]
LOGGER.error(ef)
LOGGER.error(format_exc())
return []
def count_all_approved(self):
with INSERTION_LOCK:
try:
return len(set(list(ADMIN_CACHE.keys())))
except KeyError:
pass
except Exception as ef:
num = 0
curr = self.collection.find_all()
if curr:
for chat in curr:
users = chat["users"]
num += len(users)
LOGGER.error(ef)
LOGGER.error(format_exc())
return num
def count_approved_chats(self):
with INSERTION_LOCK:
try:
return len(list(APPROVE_CACHE.keys()))
except Exception as ef:
LOGGER.error(ef)
LOGGER.error(format_exc())
return (self.collection.count()) or 0
def count_approved(self, chat_id: int):
with INSERTION_LOCK:
try:
return len(APPROVE_CACHE[chat_id])
except KeyError:
pass
except Exception as ef:
all_app = self.collection.find_one({"_id": chat_id})
if all_app:
return len(all_app["users"]) or 0
LOGGER.error(ef)
LOGGER.error(format_exc())
return 0
def load_from_db(self):
return self.collection.find_all()
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
global APPROVE_CACHE
with INSERTION_LOCK:
# Update locally
try:
old_db_local = APPROVE_CACHE[old_chat_id]
del APPROVE_CACHE[old_chat_id]
APPROVE_CACHE[new_chat_id] = old_db_local
except KeyError:
pass
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
class WarnSettings:
def __init__(self) -> None:
self.collection = MongoDB("chat_warn_settings")
def get_warnings_settings(self, chat_id: int):
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr
curr = {"_id": chat_id, "warn_mode": "kick", "warn_limit": 3}
self.collection.insert_one(curr)
return curr
def set_warnmode(self, chat_id: int, warn_mode: str = "kick"):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
self.collection.update(
{"_id": chat_id},
{"warn_mode": warn_mode},
)
return warn_mode
self.collection.insert_one(
{
"_id": chat_id,
"warn_mode": warn_mode,
"warn_limit": 3
}, )
return warn_mode
def get_warnmode(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["warn_mode"]
self.collection.insert_one(
{
"_id": chat_id,
"warn_mode": "kick",
"warn_limit": 3
}, )
return "kick"
def set_warnlimit(self, chat_id: int, warn_limit: int = 3):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
self.collection.update(
{"_id": chat_id},
{"warn_limit": warn_limit},
)
return warn_limit
self.collection.insert_one(
{
"_id": chat_id,
"warn_mode": "kick",
"warn_limit": warn_limit
}, )
return warn_limit
def get_warnlimit(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["warn_limit"]
self.collection.insert_one(
{
"_id": chat_id,
"warn_mode": "kick",
"warn_limit": 3
}, )
return 3
def count_action_chats(self, mode: str):
return self.collection.count({"warn_mode": mode})
def count_chats():
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
return collection.count() or 0
class Blacklist:
"""Class to manage database for blacklists for chats."""
def __init__(self) -> None:
self.collection = MongoDB("blacklists")
def add_blacklist(self, chat_id: int, trigger: str):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
triggers_old = curr["triggers"]
triggers_old.append(trigger)
triggers = list(set(triggers_old))
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"triggers": triggers,
},
)
return self.collection.insert_one(
{
"_id": chat_id,
"triggers": [trigger],
"action": "none",
"reason": "Automated blacklisted word",
}, )
def remove_blacklist(self, chat_id: int, trigger: str):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
triggers_old = curr["triggers"]
try:
triggers_old.remove(trigger)
except ValueError:
return "Trigger not found"
triggers = list(set(triggers_old))
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"triggers": triggers,
},
)
def get_blacklists(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["triggers"]
return []
def count_blacklists_all(self):
with INSERTION_LOCK:
curr = self.collection.find_all()
num = 0
for chat in curr:
num += len(chat["triggers"])
return num
def count_blackists_chats(self):
with INSERTION_LOCK:
curr = self.collection.find_all()
num = 0
for chat in curr:
if chat["triggers"]:
num += 1
return num
def set_action(self, chat_id: int, action: str):
with INSERTION_LOCK:
if action not in ("kick", "mute", "ban", "warn", "none"):
return "invalid action"
curr = self.collection.find_one({"_id": chat_id})
if curr:
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"action": action
},
)
return self.collection.insert_one(
{
"_id": chat_id,
"triggers": [],
"action": action,
"reason": "Automated blacklisted word",
}, )
def get_action(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["action"] or "none"
self.collection.insert_one(
{
"_id": chat_id,
"triggers": [],
"action": "none",
"reason": "Automated blacklisted word",
}, )
return "Automated blacklisted word"
def set_reason(self, chat_id: int, reason: str):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return self.collection.update(
{"_id": chat_id},
{
"_id": chat_id,
"reason": reason
},
)
return self.collection.insert_one(
{
"_id": chat_id,
"triggers": [],
"action": "none",
"reason": reason,
}, )
def get_reason(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
return curr["reason"] or "none"
self.collection.insert_one(
{
"_id": chat_id,
"triggers": [],
"action": "none",
"reason": "Automated blacklistwd word",
}, )
return "Automated blacklisted word"
def count_action_bl_all(self, action: str):
return self.collection.count({"action": action})
def rm_all_blacklist(self, chat_id: int):
with INSERTION_LOCK:
curr = self.collection.find_one({"_id": chat_id})
if curr:
self.collection.update(
{"_id": chat_id},
{"triggers": []},
)
return False
# Migrate if chat id changes!
def migrate_chat(self, old_chat_id: int, new_chat_id: int):
with INSERTION_LOCK:
old_chat_db = self.collection.find_one({"_id": old_chat_id})
if old_chat_db:
new_data = old_chat_db.update({"_id": new_chat_id})
self.collection.delete_one({"_id": old_chat_id})
self.collection.insert_one(new_data)
def count_users():
with INSERTION_LOCK:
collection = MongoDB(Users.db_name)
return collection.count()
def count_chats(query: str):
with INSERTION_LOCK:
collection = MongoDB(Greetings.db_name)
return collection.count({query: True})
def count_chats(atype: str):
with INSERTION_LOCK:
collection = MongoDB(Pins.db_name)
return collection.count({atype: True})
