diff --git a/.gitignore b/.gitignore index b6e4761..565b46e 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,15 @@ dmypy.json # Pyre type checker .pyre/ + +# SQlite3 database and session files + +database.db +*.session +*.session-journal + +# Virtual Environment + +BotEnv +BotEnv/ +BotEnv/* \ No newline at end of file diff --git a/BotBase/database/query.py b/BotBase/database/query.py index a1ef9c6..16a6b1c 100644 --- a/BotBase/database/query.py +++ b/BotBase/database/query.py @@ -1,8 +1,7 @@ import sqlite3.dbapi2 as sqlite3 -from ..config import DB_GET_USERS, DB_GET_USER, DB_RELPATH, DB_SET_USER +from ..config import DB_GET_USERS, DB_GET_USER, DB_PATH, DB_SET_USER, DB_UPDATE_NAME, DB_BAN_USER, DB_UNBAN_USER import logging import time -from types import FunctionType import os @@ -24,9 +23,10 @@ def create_database(path: str, query: str): except sqlite3.Error as query_error: logging.info(f"An error has occurred while executing query: {query_error}") + def get_user(tg_id: int): try: - database = sqlite3.connect(DB_RELPATH) + database = sqlite3.connect(DB_PATH) except sqlite3.Error as connection_error: logging.error(f"An error has occurred while connecting to database: {connection_error}") else: @@ -39,9 +39,26 @@ def get_user(tg_id: int): logging.error(f"An error has occurred while executing DB_GET_USER query: {query_error}") return query_error + +def update_name(tg_id: int, name: str): + try: + database = sqlite3.connect(DB_PATH) + except sqlite3.Error as connection_error: + logging.error(f"An error has occurred while connecting to database: {connection_error}") + else: + try: + with database: + cursor = database.cursor() + query = cursor.execute(DB_UPDATE_NAME, (name, tg_id)) + return query.fetchone() + except sqlite3.Error as query_error: + logging.error(f"An error has occurred while executing DB_UPDATE_NAME query: {query_error}") + return query_error + + def get_users(): try: - database = sqlite3.connect(DB_RELPATH) + database = sqlite3.connect(DB_PATH) except sqlite3.Error as connection_error: logging.error(f"An error has occurred while connecting to database: {connection_error}") else: @@ -57,7 +74,7 @@ def get_users(): def set_user(tg_id: int, uname: str): try: - database = sqlite3.connect(DB_RELPATH) + database = sqlite3.connect(DB_PATH) except sqlite3.Error as connection_error: logging.error(f"An error has occurred while connecting to database: {connection_error}") else: @@ -70,3 +87,37 @@ def set_user(tg_id: int, uname: str): except sqlite3.Error as query_error: logging.error(f"An error has occurred while executing DB_GET_USERS query: {query_error}") return query_error + + +def ban_user(tg_id: int): + try: + database = sqlite3.connect(DB_PATH) + except sqlite3.Error as connection_error: + logging.error(f"An error has occurred while connecting to database: {connection_error}") + else: + try: + with database: + cursor = database.cursor() + cursor.execute(DB_BAN_USER, (tg_id, )) + cursor.close() + return True + except sqlite3.Error as query_error: + logging.error(f"An error has occurred while executing DB_BAN_USER query: {query_error}") + return query_error + + +def unban_user(tg_id: int): + try: + database = sqlite3.connect(DB_PATH) + except sqlite3.Error as connection_error: + logging.error(f"An error has occurred while connecting to database: {connection_error}") + else: + try: + with database: + cursor = database.cursor() + cursor.execute(DB_UNBAN_USER, (tg_id, )) + cursor.close() + return True + except sqlite3.Error as query_error: + logging.error(f"An error has occurred while executing DB_UNBAN_USER query: {query_error}") + return query_error diff --git a/BotBase/methods/safe_edit.py b/BotBase/methods/safe_edit.py index a794241..768db3d 100644 --- a/BotBase/methods/safe_edit.py +++ b/BotBase/methods/safe_edit.py @@ -5,7 +5,7 @@ from pyrogram import Client, CallbackQuery from typing import Union -def edit_message_text(update: Union[CallbackQuery, Client], sleep: bool = True *args, **kwargs): +def edit_message_text(update: Union[CallbackQuery, Client], sleep: bool = True, *args, **kwargs): """Edits a message in a way that never triggers exceptions and logs errors :param update: The pyrogram.Client instance or pyrogram.CallbackQuery diff --git a/BotBase/methods/safe_send.py b/BotBase/methods/safe_send.py index 87961e8..5e23ee8 100644 --- a/BotBase/methods/safe_send.py +++ b/BotBase/methods/safe_send.py @@ -20,8 +20,8 @@ def send_message(client: Client, sleep: bool = True, *args, **kwargs): except FloodWait as fw: logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") if sleep: - time.sleep(fw.x) - return fw + time.sleep(fw.x) + return fw except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error @@ -43,8 +43,8 @@ def send_photo(client: Client, sleep: bool = True, *args, **kwargs): except FloodWait as fw: logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") if sleep: - time.sleep(fw.x) - return fw + time.sleep(fw.x) + return fw except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error @@ -66,8 +66,8 @@ def send_audio(client: Client, sleep: bool = True, *args, **kwargs): except FloodWait as fw: logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") if sleep: - time.sleep(fw.x) - return fw + time.sleep(fw.x) + return fw except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error @@ -88,8 +88,8 @@ def send_sticker(client: Client, sleep: bool = True, *args, **kwargs): except FloodWait as fw: logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") if sleep: - time.sleep(fw.x) - return fw + time.sleep(fw.x) + return fw except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error @@ -110,11 +110,8 @@ def send_animation(client: Client, sleep: bool = True, *args, **kwargs): except FloodWait as fw: logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") if sleep: - time.sleep(fw.x) - return fw + time.sleep(fw.x) + return fw except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error - - - diff --git a/BotBase/methods/various.py b/BotBase/methods/various.py index 076e6ed..c9f17f4 100644 --- a/BotBase/methods/various.py +++ b/BotBase/methods/various.py @@ -7,8 +7,8 @@ import logging def answer(query: CallbackQuery, sleep: bool = True, *args, **kwargs): """Answers a query in a way that never triggers exceptions and logs errors - :param update: The pyrogram.CallbackQuery object to call the method for - :type update: class: CallbackQuery + :param query: The pyrogram.CallbackQuery object to call the method for + :type query: class: CallbackQuery :param sleep: If True, the default, the function will call time.sleep() in case of a FloodWait exception and return the exception object after the sleep is done, otherwise the ``FloodWait`` exception is returned @@ -30,7 +30,17 @@ def answer(query: CallbackQuery, sleep: bool = True, *args, **kwargs): def delete_messages(client, sleep: bool = True, *args, **kwargs): - """Deletes messages in a way that never triggers exceptions and logs errors""" + """Deletes messages in a way that never triggers exceptions and logs errors + + :param client: The pyrogram.Client instance to call the method for + :type client: class: Client + :param sleep: If True, the default, the function will call time.sleep() + in case of a FloodWait exception and return the exception object + after the sleep is done, otherwise the ``FloodWait`` exception is returned + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ try: return client.delete_messages(*args, **kwargs) @@ -42,3 +52,28 @@ def delete_messages(client, sleep: bool = True, *args, **kwargs): except RPCError as generic_error: logging.error(f"An exception occurred: {generic_error}") return generic_error + + +def get_users(client, sleep: bool = True, *args, **kwargs): + """Calls get_users in a way that never triggers exceptions and logs errors + + :param client: The pyrogram.Client instance to call the method for + :type client: class: Client + :param sleep: If True, the default, the function will call time.sleep() + in case of a FloodWait exception and return the exception object + after the sleep is done, otherwise the ``FloodWait`` exception is returned + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ + + try: + return client.get_users(*args, **kwargs) + except FloodWait as fw: + logging.warning(f"FloodWait! A wait of {fw.x} seconds is required") + if sleep: + time.sleep(fw.x) + return fw + except RPCError as generic_error: + logging.error(f"An exception occurred: {generic_error}") + return generic_error diff --git a/BotBase/modules/admin.py b/BotBase/modules/admin.py index 1d00f7e..09665e2 100644 --- a/BotBase/modules/admin.py +++ b/BotBase/modules/admin.py @@ -1,80 +1,196 @@ from ..config import ADMINS, USER_INFO, INVALID_SYNTAX, ERROR, NONNUMERIC_ID, USERS_COUNT, \ - NO_PARAMETERS, ID_MISSING, GLOBAL_MESSAGE_STATS + NO_PARAMETERS, ID_MISSING, GLOBAL_MESSAGE_STATS, NAME, WHISPER_FROM, USER_INFO_UPDATED, USER_INFO_UNCHANGED, \ + USER_BANNED, USER_UNBANNED, CANNOT_BAN_ADMIN, USER_ALREADY_BANNED, USER_NOT_BANNED, YOU_ARE_BANNED, YOU_ARE_UNBANNED from pyrogram import Client, Filters -from ..database.query import get_user, get_users +from ..database.query import get_user, get_users, update_name, ban_user, unban_user from .antiflood import BANNED_USERS import random from ..methods.safe_send import send_message +from ..methods.various import get_users as get_telegram_users import logging +import itertools +import re ADMINS_FILTER = Filters.user(list(ADMINS.keys())) -@Client.on_message(Filters.command("count") & ADMINS_FILTER & Filters.private & ~BANNED_USERS) +@Client.on_message(Filters.command("count") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def count_users(client, message): - logging.warning(f"Admin with id {message.from_user.id} sent /count") + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /count") count = len(get_users()) - send_message(client, message.chat.id, USERS_COUNT.format(count)) + send_message(client, True, message.chat.id, USERS_COUNT.format(count=count)) -@Client.on_message(Filters.command("getuser") & ADMINS_FILTER & Filters.private & ~BANNED_USERS) +@Client.on_message(Filters.command("getuser") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def get_user_info(client, message): if len(message.command) == 2: if message.command[1].isdigit(): user = get_user(message.command[1]) if user: - logging.warning(f"Admin with id {message.from_user.id} sent /getuser {message.command[1]}") + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /getuser {message.command[1]}") _, uid, uname, date, banned = user admin = uid in ADMINS - text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date, - status='✅' if banned else '❌', - admin='❌' if not admin else '✅'), - ) - send_message(client, message.chat.id, text) + text = USER_INFO.format(uid=uid, + uname='@' + uname if uname != 'null' else uname, + date=date, + status='✅' if banned else '❌', + admin='❌' if not admin else '✅') + send_message(client, True, message.chat.id, text) else: - send_message(client, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}") + send_message(client, True, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}") else: - send_message(client, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") + send_message(client, True, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") else: - send_message(client, message.chat.id, f"{INVALID_SYNTAX}: Use /getuser user_id") + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /getuser user_id") -@Client.on_message(Filters.command("getranduser") & ADMINS_FILTER & Filters.private & ~BANNED_USERS) +@Client.on_message(Filters.command("getranduser") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def get_random_user(client, message): - logging.warning(f"Admin with id {message.from_user.id} sent /getranduser") + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /getranduser") if len(message.command) > 1: - send_message(client, message.chat.id, f"{INVALID_SYNTAX}: {NO_PARAMETERS.format(command='/getranduser')}") + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: {NO_PARAMETERS.format(command='/getranduser')}") else: user = random.choice(get_users()) rowid, uid, uname, date, banned = get_user(*user) admin = uid in ADMINS - text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date, + text = USER_INFO.format(uid=uid, + uname='@' + uname if uname != 'null' else uname, + date=date, status='✅' if banned else '❌', - admin='❌' if not admin else '✅'), - ) - send_message(client, message.chat.id, text) + admin='❌' if not admin else '✅') + send_message(client, True, message.chat.id, text) -@Client.on_message(Filters.command("global") & ADMINS_FILTER & Filters.private & ~BANNED_USERS) +@Client.on_message(Filters.command("global") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def global_message(client, message): if len(message.command) > 1: msg = message.text.html[7:] - logging.warning(f"Admin with id {message.from_user.id} sent the following global message: {msg}") + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent the following global message: {msg}") missed = 0 count = 0 - for uid in get_users(): + for uid in itertools.chain(*get_users()): count += 1 - if isinstance(send_message(client, *uid, msg), Exception): + result = send_message(client, True, uid, msg) + if isinstance(result, Exception): + logging.error(f"Could not deliver the global message to {uid} because of {type(result).__name__}: {result}") missed += 1 - send_message(client, message.chat.id, GLOBAL_MESSAGE_STATS.format(count=count, success=(count - missed), msg=msg)) + logging.warning(f"{count - missed}/{count} global messages were successfully delivered") + send_message(client, True, message.chat.id, GLOBAL_MESSAGE_STATS.format(count=count, success=(count - missed), msg=msg)) else: - send_message(client, message.chat.id, f"{INVALID_SYNTAX}: Use /global message" - "\n🍮 Note that the /global command supports markdown and html styling") + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /global message" + f"\n🍮 Note that the /global command supports markdown and html styling") -@Client.on_message(Filters.command("whisper") & ADMINS_FILTER & Filters.private & ~BANNED_USERS) + +@Client.on_message(Filters.command("whisper") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def whisper(client, message): - if len(message.command) > 1: - msg = " ".join(message.command[1:]) - logging.warning(f"Admin with id {message.from_user.id} sent /whisper to {message.command} + if len(message.command) > 2: + msg = message.text.html[9:] + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent {message.text.html}") + if message.command[1].isdigit(): + msg = msg[re.search(message.command[1], msg).end():] + uid = int(message.command[1]) + if uid in itertools.chain(*get_users()): + result = send_message(client, True, uid, WHISPER_FROM.format(admin=f"[{ADMINS[message.from_user.id]}]({NAME.format(message.from_user.id)})", + msg=msg) + ) + if isinstance(result, Exception): + logging.error( + f"Could not whisper to {uid} because of {type(result).__name__}: {result}") + send_message(client, True, message.chat.id, f"{ERROR}: {type(result).__name__} -> {result}") + else: + send_message(client, True, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=uid)}") + else: + send_message(client, True, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") + else: + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /whisper ID message" + f"\n🍮 Note that the /whisper command supports markdown and html styling") + +@Client.on_message(Filters.command("update") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) +def update(client, message): + if len(message.command) == 2: + if message.command[1].isdigit(): + user = get_user(message.command[1]) + if user: + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /update {message.command[1]}") + _, uid, uname, date, banned = user + new = get_telegram_users(client, True, uid) + if isinstance(new, Exception): + logging.error(f"An error has occurred when calling get_users({uid}), {type(new).__name__}: {new}") + send_message(client, True, message.chat.id, f"{ERROR}: {type(new).__name__} -> {new}") + else: + if new.username is None: + new.username = "null" + if new.username != uname: + update_name(uid, new.username) + send_message(client, True, message.chat.id, USER_INFO_UPDATED) + else: + send_message(client, True, message.chat.id, USER_INFO_UNCHANGED) + else: + send_message(client, True, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}") + else: + send_message(client, True, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") + else: + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /update user_id") + + +@Client.on_message(Filters.command("ban") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) +def ban(client, message): + if len(message.command) == 2: + if message.command[1].isdigit(): + if int(message.command[1]) in ADMINS: + send_message(client, True, message.chat.id, CANNOT_BAN_ADMIN) + else: + user = get_user(message.command[1]) + if user: + if not user[-1]: + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /ban {message.command[1]}") + _, uid, uname, date, banned = user + res = ban_user(int(message.command[1])) + if isinstance(res, Exception): + logging.error(f"An error has occurred when calling ban_user({uid}), {type(res).__name__}: {res}") + send_message(client, True, message.chat.id, f"{ERROR}: {type(res).__name__} -> {res}") + else: + send_message(client, True, message.chat.id, USER_BANNED) + send_message(client, True, uid, YOU_ARE_BANNED) + BANNED_USERS.add(uid) + else: + send_message(client, True, message.chat.id, USER_ALREADY_BANNED) + else: + send_message(client, True, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}") + else: + send_message(client, True, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") + else: + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /ban user_id") + + +@Client.on_message(Filters.command("unban") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) +def unban(client, message): + if len(message.command) == 2: + if message.command[1].isdigit(): + if int(message.command[1]) in ADMINS: + send_message(client, True, message.chat.id, CANNOT_BAN_ADMIN) + else: + user = get_user(message.command[1]) + if user: + if user[-1]: + logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /unban {message.command[1]}") + _, uid, uname, date, banned = user + res = unban_user(int(message.command[1])) + if isinstance(res, Exception): + logging.error(f"An error has occurred when calling unban_user({uid}), {type(res).__name__}: {res}") + send_message(client, True, message.chat.id, f"{ERROR}: {type(res).__name__} -> {res}") + else: + send_message(client, True, message.chat.id, USER_UNBANNED) + if uid in BANNED_USERS: + BANNED_USERS.remove(uid) + send_message(client, True, uid, YOU_ARE_UNBANNED) + else: + send_message(client, True, message.chat.id, USER_NOT_BANNED) + else: + send_message(client, True, message.chat.id, f"{ERROR}: {ID_MISSING.format(uid=message.command[1])}") + else: + send_message(client, True, message.chat.id, f"{ERROR}: {NONNUMERIC_ID}") + else: + send_message(client, True, message.chat.id, f"{INVALID_SYNTAX}: Use /unban user_id") diff --git a/BotBase/modules/antiflood.py b/BotBase/modules/antiflood.py index fd2fcc5..a9471ec 100644 --- a/BotBase/modules/antiflood.py +++ b/BotBase/modules/antiflood.py @@ -29,7 +29,7 @@ def is_flood(updates: list): return False -@Client.on_callback_query(FILTER & ~BYPASS_USERS, group=-1) +@Client.on_callback_query(~BYPASS_USERS, group=-1) @Client.on_message(FILTER & ~BYPASS_USERS, group=-1) def anti_flood(client, update): """Anti flood module""" diff --git a/BotBase/modules/livechat.py b/BotBase/modules/livechat.py index dfe6ebe..472f789 100644 --- a/BotBase/modules/livechat.py +++ b/BotBase/modules/livechat.py @@ -2,48 +2,50 @@ from pyrogram import Client, Filters, InlineKeyboardMarkup, InlineKeyboardButton from ..methods.safe_send import send_message from ..methods.safe_edit import edit_message_text from ..methods.various import answer, delete_messages -from ..config import CACHE, ADMINS, STATUSES, ADMINS_LIST_UPDATE_DELAY, callback_regex, admin_is_chatting, \ +from ..config import CACHE, ADMINS, ADMINS_LIST_UPDATE_DELAY, callback_regex, admin_is_chatting, \ user_is_chatting, LIVE_CHAT_STATUSES, STATUS_BUSY, STATUS_FREE, SUPPORT_REQUEST_SENT, SUPPORT_NOTIFICATION, \ ADMIN_JOINS_CHAT, USER_CLOSES_CHAT, JOIN_CHAT_BUTTON, USER_INFO, USER_LEAVES_CHAT, ADMIN_MESSAGE, USER_MESSAGE, \ - TOO_FAST, CHAT_BUSY, LEAVE_CURRENT_CHAT, USER_JOINS_CHAT -from collections import defaultdict + TOO_FAST, CHAT_BUSY, LEAVE_CURRENT_CHAT, USER_JOINS_CHAT, NAME import time from ..database.query import get_user -from datetime import datetime from .antiflood import BANNED_USERS from .start import back_start +import logging + -CHATS = defaultdict(None) -NAME = "tg://user?id={}" ADMINS_FILTER = Filters.user(list(ADMINS.keys())) BUTTONS = InlineKeyboardMarkup( - [[InlineKeyboardButton("🔙 Back", "back_start")], [InlineKeyboardButton("🔄 Update", "update_admins_list")]]) + [ + [InlineKeyboardButton("🔙 Back", "back_start")], + [InlineKeyboardButton("🔄 Update", "update_admins_list")] + ]) @Client.on_callback_query(Filters.callback_data("sos") & ~BANNED_USERS) -def begin_chat(_, query): +def begin_chat(client, query): CACHE[query.from_user.id] = ["AWAITING_ADMIN", time.time()] queue = LIVE_CHAT_STATUSES - for admin_id, data in STATUSES.items(): - admin_name, status = data - if status == "free": + for admin_id, admin_name in ADMINS.items(): + status = CACHE[admin_id][0] + if status != "IN_CHAT": queue += f"- {STATUS_FREE}" else: queue += f"- {STATUS_BUSY}" queue += f"[{admin_name}]({NAME.format(admin_id)})\n" - msg = edit_message_text(query, SUPPORT_REQUEST_SENT.format(queue=queue, date=time.strftime('%d/%m/%Y %T')), + msg = edit_message_text(query, True, SUPPORT_REQUEST_SENT.format(queue=queue, date=time.strftime('%d/%m/%Y %T')), reply_markup=BUTTONS) join_chat_button = InlineKeyboardMarkup([[InlineKeyboardButton(JOIN_CHAT_BUTTON, f"join_{query.from_user.id}")]]) user = get_user(query.from_user.id) - _, uid, uname, date = user + _, uid, uname, date, banned = user admin = uid in ADMINS text = USER_INFO.format(uid=uid, uname='@' + uname if uname != 'null' else uname, date=date, - status='User' if not admin else 'Admin', - last_call=datetime.utcfromtimestamp(int(last_call)), imei=imei) + status='✅' if banned else '❌', + admin='❌' if not admin else '✅') CACHE[query.from_user.id].append([]) for admin in ADMINS: - if STATUSES[admin][0] == "free": - message = send_message(_, admin, SUPPORT_NOTIFICATION.format(uinfo=text), reply_markup=join_chat_button) + status = CACHE[admin][0] + if status != "IN_CHAT": + message = send_message(client, True, admin, SUPPORT_NOTIFICATION.format(uinfo=text), reply_markup=join_chat_button) CACHE[query.from_user.id][-1].append((message.chat.id, message.message_id)) CACHE[query.from_user.id][-1].append((msg.chat.id, msg.message_id)) @@ -54,19 +56,19 @@ def update_admins_list(_, query): if CACHE[query.from_user.id][0] == "AWAITING_ADMIN": CACHE[query.from_user.id] = ["AWAITING_ADMIN", time.time()] queue = LIVE_CHAT_STATUSES - for admin_id, data in STATUSES.items(): - admin_name, status = data - if status == "free": + for admin_id, admin_name in ADMINS.items(): + status = CACHE[admin_id][0] + if status != "IN_CHAT": queue += f"- {STATUS_FREE}" else: queue += f"- {STATUS_BUSY}" queue += f"[{admin_name}]({NAME.format(admin_id)})\n" - edit_message_text(query, SUPPORT_REQUEST_SENT.format(queue=queue, date=time.strftime('%d/%m/%Y %T')), + edit_message_text(query, True, SUPPORT_REQUEST_SENT.format(queue=queue, date=time.strftime('%d/%m/%Y %T')), reply_markup=BUTTONS) else: back_start(_, query) else: - answer(query, TOO_FAST, show_alert=True) + answer(query, True, TOO_FAST, show_alert=True) @Client.on_callback_query(callback_regex(r"close_chat_\d+") & ~BANNED_USERS) @@ -76,23 +78,26 @@ def close_chat(_, query): if query.from_user.id in ADMINS: data = CACHE[CACHE[query.from_user.id][1]][-1] if isinstance(data, list): + data.append((query.from_user.id, query.message.message_id)) for chatid, message_ids in data: - delete_messages(_, chatid, message_ids) - if STATUSES[query.from_user.id][1] != "free": - STATUSES[query.from_user.id][1] = "free" - admin_id, admin_name = query.from_user.id, STATUSES[query.from_user.id][0] - edit_message_text(query, USER_LEAVES_CHAT) + delete_messages(_, True, chatid, message_ids) + status = CACHE[query.from_user.id][0] + if status == "IN_CHAT": + del CACHE[query.from_user.id][1] + send_message(_, True, query.from_user.id, USER_LEAVES_CHAT) + admin_id, admin_name = query.from_user.id, ADMINS[query.from_user.id] if CACHE[user_id][1]: - send_message(_, user_id, + send_message(_, True, user_id, USER_CLOSES_CHAT.format(user_id=NAME.format(admin_id), user_name=admin_name)) if user_id in CACHE: del CACHE[user_id] + logging.warning(f"{ADMINS[admin_id]} [{admin_id}] has terminated the chat with user {CACHE[admin_id][1]}") del CACHE[admin_id] else: data = CACHE[query.from_user.id][-1] if isinstance(data, list): for chatid, message_ids in data: - delete_messages(_, chatid, message_ids) + delete_messages(_, True, chatid, message_ids) admin_id = CACHE[query.from_user.id][1] if CACHE[user_id][1]: if query.from_user.first_name: @@ -101,8 +106,10 @@ def close_chat(_, query): user_name = query.from_user.username else: user_name = "Anonymous" - edit_message_text(query, USER_LEAVES_CHAT) - send_message(_, CACHE[user_id][1], + logging.warning(f"{user_name} [{query.from_user.id}] has terminated the chat with admin {ADMINS[admin_id]} [{admin_id}]") + send_message(_, True, query.from_user.id, + USER_LEAVES_CHAT) + send_message(_, True, CACHE[user_id][1], USER_CLOSES_CHAT.format(user_id=NAME.format(query.from_user.id), user_name=user_name)) del CACHE[query.from_user.id] del CACHE[admin_id] @@ -110,11 +117,13 @@ def close_chat(_, query): back_start(_, query) -@Client.on_message(ADMINS_FILTER & Filters.private & admin_is_chatting() & Filters.text & ~BANNED_USERS) +@Client.on_message(admin_is_chatting() & Filters.text & ~BANNED_USERS) def forward_from_admin(client, message): - send_message(client, STATUSES[message.from_user.id][1], - ADMIN_MESSAGE.format(STATUSES[message.from_user.id][0], NAME.format(message.from_user.id), - message.text.html)) + logging.warning(f"Admin {ADMINS[message.from_user.id]} [{message.from_user.id}] says to {CACHE[message.from_user.id][1]}: {message.text.html}") + send_message(client, True, CACHE[message.from_user.id][1], + ADMIN_MESSAGE.format(user_name=ADMINS[message.from_user.id], + user_id=NAME.format(message.from_user.id), + message=message.text.html)) @Client.on_message(user_is_chatting() & Filters.text & ~BANNED_USERS) @@ -125,7 +134,8 @@ def forward_from_user(client, message): name = message.from_user.username else: name = "Anonymous" - send_message(client, CACHE[message.from_user.id][1], + logging.warning(f"User {name} [{message.from_user.id}] says to Admin {ADMINS[CACHE[message.from_user.id][1]]} [{CACHE[message.from_user.id][1]}]: {message.text.html}") + send_message(client, True, CACHE[message.from_user.id][1], USER_MESSAGE.format(user_name=name, user_id=NAME.format(message.from_user.id), message=message.text.html)) @@ -135,18 +145,18 @@ def join_chat(_, query): if CACHE[query.from_user.id][0] != "IN_CHAT": user_id = int(query.data.split("_")[1]) if CACHE[user_id][0] != "AWAITING_ADMIN": - answer(query, CHAT_BUSY) + answer(query, True, CHAT_BUSY) else: buttons = InlineKeyboardMarkup([[InlineKeyboardButton("❌ Close chat", f"close_chat_{user_id}")]]) - admin_id, admin_name = query.from_user.id, STATUSES[query.from_user.id][0] - STATUSES[query.from_user.id][1] = user_id + admin_id, admin_name = query.from_user.id, ADMINS[query.from_user.id] CACHE[user_id] = ["IN_CHAT", admin_id, CACHE[user_id][-1]] CACHE[query.from_user.id] = ["IN_CHAT", user_id] - message = send_message(_, query.from_user.id, USER_JOINS_CHAT, reply_markup=buttons) - send_message(_, user_id, ADMIN_JOINS_CHAT.format(admin_name=admin_name, admin_id=NAME.format(admin_id)), - reply_markup=buttons) + message = send_message(_, True, query.from_user.id, USER_JOINS_CHAT, reply_markup=buttons) + admin_joins = send_message(_, True, user_id, ADMIN_JOINS_CHAT.format(admin_name=admin_name, admin_id=NAME.format(admin_id)), + reply_markup=buttons) for chatid, message_ids in CACHE[CACHE[query.from_user.id][1]][-1]: - delete_messages(_, chatid, message_ids) + delete_messages(_, True, chatid, message_ids) CACHE[user_id][-1].append((message.chat.id, message.message_id)) + CACHE[user_id][-1].append((admin_joins.chat.id, admin_joins.message_id)) else: - answer(query, LEAVE_CURRENT_CHAT) + answer(query, True, LEAVE_CURRENT_CHAT) diff --git a/BotBase/modules/start.py b/BotBase/modules/start.py index 8252f2b..96b9ec8 100644 --- a/BotBase/modules/start.py +++ b/BotBase/modules/start.py @@ -1,11 +1,23 @@ from pyrogram import Client, Filters, InlineKeyboardButton, InlineKeyboardMarkup from .antiflood import BANNED_USERS -from ..config import GREET, BUTTONS, CREDITS -from ..database.query import get_users, set_user +from ..config import GREET, BUTTONS, CREDITS, CACHE, YOU_ARE_BANNED +from ..database.query import get_users, set_user, get_user import logging import itertools from ..methods.safe_send import send_message from ..methods.safe_edit import edit_message_text +from ..methods.various import delete_messages + + +def check_user_banned(tg_id: int): + res = get_user(tg_id) + if isinstance(res, Exception): + return False + else: + if res[-1]: + return True + else: + return False @Client.on_message(Filters.command("start") & ~BANNED_USERS & Filters.private) @@ -19,21 +31,47 @@ def start_handler(client, message): name = message.from_user.username else: name = "Anonymous" - if message.from_user.id not in itertools.chain(*get_users()): - logging.warning(f"New user detected ({message.from_user.id}), adding to database") - set_user(message.from_user.id, None if not message.from_user.username else message.from_user.username) - if GREET: - send_message(client, - message.chat.id, - GREET.format(mention=f"[{name}](tg://user?id={message.from_user.id})", - id=message.from_user.id, - username=message.from_user.username - ), - reply_markup=BUTTONS - ) + if check_user_banned(message.from_user.id): + BANNED_USERS.add(message.from_user.id) + send_message(client, True, message.from_user.id, YOU_ARE_BANNED) + else: + if message.from_user.id not in itertools.chain(*get_users()): + logging.warning(f"New user detected ({message.from_user.id}), adding to database") + set_user(message.from_user.id, message.from_user.username) + if GREET: + send_message(client, + True, + message.from_user.id, + GREET.format(mention=f"[{name}](tg://user?id={message.from_user.id})", + id=message.from_user.id, + username=message.from_user.username + ), + reply_markup=BUTTONS + ) @Client.on_callback_query(Filters.callback_data("info")) def bot_info(_, query): buttons = InlineKeyboardMarkup([[InlineKeyboardButton("🔙 Back", "back_start")]]) - edit_message_text(query, CREDITS, reply_markup=buttons) + edit_message_text(query, True, CREDITS.format(), reply_markup=buttons) + + +@Client.on_callback_query(Filters.callback_data("back_start")) +def back_start(_, query): + if query.from_user.first_name: + name = query.from_user.first_name + elif query.from_user.username: + name = query.from_user.username + else: + name = "Anonymous" + if CACHE[query.from_user.id][0] == "AWAITING_ADMIN": + data = CACHE[query.from_user.id][-1] + if isinstance(data, list): + for chatid, message_ids in data: + delete_messages(_, True, chatid, message_ids) + start_handler(_, query) + else: + edit_message_text(query, True, + GREET.format(mention=f"[{name}](tg://user?id={query.from_user.id})", id=query.from_user.id, + username=query.from_user.username), + reply_markup=BUTTONS) diff --git a/DATABASE.md b/DATABASE.md index bc6e884..3412e2f 100644 --- a/DATABASE.md +++ b/DATABASE.md @@ -36,9 +36,11 @@ of tuples. Each tuple contains a user ID as stored in the database - `set_user()` -> Saves an ID/username pair (in this order) to the database. The username parameter can be `None` -- `ban_user()` -> Bans the user with the given user ID (Coming soon) +- `ban_user()` -> Bans the user with the given user ID -- `unban_user()` -> Unbans a user with the given ID (Coming soon) +- `unban_user()` -> Unbans a user with the given ID + +- `update_user` -> Updates a user's username with the given ID # I need MySQL/other DBMS! diff --git a/METHODS.md b/METHODS.md index f362705..eb46344 100644 --- a/METHODS.md +++ b/METHODS.md @@ -4,7 +4,7 @@ BotBase has a builtin collection of wrappers around Pyrogram methods that make it even easier to use them properly. **DISCLAIMER**: These methods are just wrappers around Pyrogram's ones and behave -exactly the same. Every method listed here takes 2 extra positional arguments, +exactly the same. Every method listed here takes 2 extra arguments, namely a `Client`/`CallbackQuery` instance and a boolean parameter (read below) All other arguments, including keyword ones, are passed to pyrogram directly. @@ -44,4 +44,5 @@ List of the available functions in `BotBase.methods.various` - `answer` (for `CallbackQuery` objects) - `delete_messages` +- `get_users` diff --git a/README.md b/README.md index c3197ef..f9fa723 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,10 @@ # BotBase -BotBase is a collection of plugins that use [Pyrogram's](https://github.com/pyrogram/pyrogram) API to make bot development extremely easy. +BotBase is a collection of plugins that use [Pyrogram's](https://github.com/pyrogram/pyrogram) API to Telegram make bot development extremely easy. ### Disclaimer -BotBase requires a solid knowledge of pyrogram and of the Telegram MTProto API itself, you can check pyrogram' docs [here](https://docs.pyrogram.org) +BotBase requires a solid knowledge of pyrogram and of the Telegram MTProto API itself, you can check pyrogram's docs [here](https://docs.pyrogram.org) Also, you need to know how to host a bot yourself. I mean, I coded all of this for you already, make some effort! @@ -56,20 +56,20 @@ This is the administrative module for the bot, and it also has its own section i To configure this plugin, go to the appropriate section in `config.py` and change the default values in the `ADMINS` dictionary. Do **NOT** change the value of `ADMINS`, just update the dictionary as follows: -- Use the admin's ID as a key -- As a value you can set the name that will be displayed to users in the admin's queue +- Use the admin's Telegram ID as a key +- As a value choose the name that the users will see when that admin joins a chat The available commands are: - `/getuser ID`: Fetches user information by its Telegram ID - `/getranduser`: Fetches a random user from the database -- `/ban ID`: Bans a user from using the bot, permanently (Coming soon) -- `/unban ID`: Unbans a user from using the bot (Coming soon) +- `/ban ID`: Bans a user from using the bot, permanently +- `/unban ID`: Unbans a user from using the bot - `/count`: Shows the current number of registered users -- `/global msg`: Broadcast whatever comes after `/global` to all users, supports HTML and markdown formatting -- `/whisper ID msg`: Send `msg` to a specific user given its ID. HTML and markdown formatting supported (Coming soon) -- `/update ID`: Updates the user's info in the database, if they've changed (Coming soon) +- `/global msg`: Broadcast `msg` to all users, supports HTML and markdown formatting +- `/whisper ID msg`: Send `msg` to a specific user given its ID. HTML and markdown formatting supported +- `/update ID`: Updates the user's info in the database, if they've changed ### Plugins - Antiflood @@ -92,7 +92,8 @@ If you don't know what a smart plugin is, check [this link](https://docs.pyrogra There are some things to keep in mind, though: -- If you want to protect your plugin from flood, import the `BotBase.modules.antiflood.BANNED_USERS` filter (basically a `Filters.user()` object) and use it like this: `~BANNED_USERS`. This will restrict banned users from reaching your handler at all +- If you want to protect your plugin from flood, import the `BotBase.modules.antiflood.BANNED_USERS` filter (basically a `Filters.user()` object) and use it like this: `~BANNED_USERS`. This will restrict banned users from reaching your handler at all. +Please note that users banned with the `/ban` command will be put in that filter, too. - To avoid repetition with try/except blocks, BotBase also implements some wrappers around `pyrogram.Client` and `pyrogram.CallbackQuery` (and many more soon) that perform automatic exception handling and log to the console automatically, check the `METHODS.md` file in this repo to know more -- Nothing restricts you from changing how the default plugins work, but this is not advised. The default plugins have been design to cooperate together and breaking this might lead to obscure tracebacks and errors that are hard to debug +- Nothing restricts you from changing how the default plugins work, but this is not advised. The default plugins have been designed to cooperate together and breaking this might lead to obscure tracebacks and errors that are hard to debug - BotBase also has many default methods to handle database interaction, check the `DATABASE.md` file in this repo to know more diff --git a/bot.py b/bot.py index 4d84d82..aaacc91 100644 --- a/bot.py +++ b/bot.py @@ -6,14 +6,15 @@ import importlib if __name__ == "__main__": MODULE_NAME = "BotBase" # Change this to match the FOLDER name that contains the config.py file - conf = importlib.import_module(MODULE_NAME) + conf = importlib.import_module(f"{MODULE_NAME}.config") + dbmodule = importlib.import_module(f"{MODULE_NAME}.database.query") logging.basicConfig(format=conf.LOGGING_FORMAT, datefmt=conf.DATE_FORMAT, level=conf.LOGGING_LEVEL) bot = Client(api_id=conf.API_ID, api_hash=conf.API_HASH, bot_token=conf.BOT_TOKEN, plugins=conf.PLUGINS_ROOT, session_name=conf.SESSION_NAME, workers=conf.WORKERS_NUM) Session.notice_displayed = True try: logging.warning("Running create_database()") - conf.create_database(conf.DB_RELPATH, conf.DB_CREATE) + dbmodule.create_database(conf.DB_PATH, conf.DB_CREATE) logging.warning("Database interaction complete") logging.warning("Starting bot") bot.start()