From fc110099f8dc40361b8987defa81747a7a2989c7 Mon Sep 17 00:00:00 2001 From: nocturn9x Date: Thu, 25 Jun 2020 17:30:30 +0200 Subject: [PATCH] Fixed error after major refactoring --- BotBase/methods/safe_edit.py | 84 +++++++++++++++++++++++++ BotBase/methods/safe_send.py | 117 +++++++++++++++++++++++++++++++++++ BotBase/methods/various.py | 79 +++++++++++++++++++++++ BotBase/modules/admin.py | 6 +- BotBase/modules/antiflood.py | 47 ++++++-------- BotBase/modules/livechat.py | 4 +- bot.py | 15 ++--- 7 files changed, 308 insertions(+), 44 deletions(-) create mode 100755 BotBase/methods/safe_edit.py create mode 100755 BotBase/methods/safe_send.py create mode 100755 BotBase/methods/various.py diff --git a/BotBase/methods/safe_edit.py b/BotBase/methods/safe_edit.py new file mode 100755 index 0000000..768db3d --- /dev/null +++ b/BotBase/methods/safe_edit.py @@ -0,0 +1,84 @@ +from pyrogram.errors import RPCError, FloodWait +import time +import logging +from pyrogram import Client, CallbackQuery +from typing import Union + + +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 + object to call the method for + :type update: Union[Client, 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 + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ + + try: + return update.edit_message_text(*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 + + +def edit_message_caption(update: Union[CallbackQuery, Client], sleep: bool = True, *args, **kwargs): + """Edits a message caption in a way that never triggers exceptions and logs errors + + :param update: The pyrogram.Client instance or pyrogram.CallbackQuery + object to call the method for + :type update: Union[Client, 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 + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ + + try: + return update.edit_message_caption(*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 + + +def edit_message_media(update: Union[CallbackQuery, Client], sleep: bool = True, *args, **kwargs): + """Edits a message media in a way that never triggers exceptions and logs errors + + :param update: The pyrogram.Client instance or pyrogram.CallbackQuery + object to call the method for + :type update: Union[Client, 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 + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ + + try: + return update.edit_message_media(*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/methods/safe_send.py b/BotBase/methods/safe_send.py new file mode 100755 index 0000000..5e23ee8 --- /dev/null +++ b/BotBase/methods/safe_send.py @@ -0,0 +1,117 @@ +from pyrogram.errors import RPCError, FloodWait +from pyrogram import Client +import time +import logging + + +def send_message(client: Client, sleep: bool = True, *args, **kwargs): + """Sends a message 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 + """ + + try: + return client.send_message(*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 + + +def send_photo(client: Client, sleep: bool = True, *args, **kwargs): + """Sends a photo 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 + """ + + try: + return client.send_photo(*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 + + +def send_audio(client: Client, sleep: bool = True, *args, **kwargs): + """Sends an audio 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 + """ + + try: + return client.send_audio(*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 + + +def send_sticker(client: Client, sleep: bool = True, *args, **kwargs): + """Sends a sticker 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 + """ + try: + return client.send_sticker(*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 + + +def send_animation(client: Client, sleep: bool = True, *args, **kwargs): + """Sends an animation 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 + """ + try: + return client.send_animation(*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/methods/various.py b/BotBase/methods/various.py new file mode 100755 index 0000000..c9f17f4 --- /dev/null +++ b/BotBase/methods/various.py @@ -0,0 +1,79 @@ +from pyrogram.errors import RPCError, FloodWait +import time +from pyrogram import CallbackQuery +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 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 + immediately + :returns: Whatever the called pyrogram method returns, or an exception if + the method call caused an error + """ + + try: + return query.answer(*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 + + +def delete_messages(client, sleep: bool = True, *args, **kwargs): + """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) + 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 + + +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 dc50a6f..6b59e77 100644 --- a/BotBase/modules/admin.py +++ b/BotBase/modules/admin.py @@ -1,7 +1,7 @@ from ..config import ADMINS, USER_INFO, INVALID_SYNTAX, ERROR, NONNUMERIC_ID, USERS_COUNT, \ 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, \ - MARKED_BUSY, UNMARKED_BUSY, CACHE, bot, YES, NO, NAME_MISSING + MARKED_BUSY, UNMARKED_BUSY, CACHE, YES, NO, NAME_MISSING, bot from pyrogram import Client, Filters from ..database.query import get_user, get_users, update_name, ban_user, unban_user, get_user_by_name from .antiflood import BANNED_USERS @@ -15,7 +15,6 @@ from ..methods import MethodWrapper ADMINS_FILTER = Filters.user(list(ADMINS.keys())) wrapper = MethodWrapper(bot) - @Client.on_message(Filters.command("getranduser") & ADMINS_FILTER & ~BANNED_USERS & ~Filters.edited) def get_random_user(client, message): logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /getranduser") @@ -66,7 +65,8 @@ def get_user_info(client, message): @Client.on_message(Filters.command("userbyname") & ADMINS_FILTER & Filters.private & ~BANNED_USERS & ~Filters.edited) def get_user_by_uname(client, message): if len(message.command) == 2: - user = get_user_by_name(message.command[1]) + name = message.command[1].lstrip("@").lower() + user = get_user_by_name(name) if user: logging.warning(f"{ADMINS[message.from_user.id]} [{message.from_user.id}] sent /userbyname {message.command[1]}") _, uid, uname, date, banned = user diff --git a/BotBase/modules/antiflood.py b/BotBase/modules/antiflood.py index 1078bcf..572a0a0 100644 --- a/BotBase/modules/antiflood.py +++ b/BotBase/modules/antiflood.py @@ -1,6 +1,6 @@ -from pyrogram import Client, Filters, Message, CallbackQueryHandler +from pyrogram import Client, Filters from ..config import MAX_UPDATE_THRESHOLD, ANTIFLOOD_SENSIBILITY, BAN_TIME, ADMINS, BYPASS_FLOOD, FLOOD_NOTICE, \ - COUNT_CALLBACKS_SEPARATELY, FLOOD_PERCENTAGE, CACHE, PRIVATE_ONLY, DELETE_MESSAGES, bot + FLOOD_PERCENTAGE, CACHE, PRIVATE_ONLY, DELETE_MESSAGES, bot, user_banned from collections import defaultdict import logging import time @@ -12,7 +12,6 @@ from ..config import check_user_banned MESSAGES = defaultdict(list) # Internal variable for the antiflood module BANNED_USERS = Filters.user() # Filters where the antiflood will put banned users BYPASS_USERS = Filters.user(list(ADMINS.keys())) if BYPASS_FLOOD else Filters.user() -QUERIES = defaultdict(list) if COUNT_CALLBACKS_SEPARATELY else MESSAGES FILTER = Filters.private if PRIVATE_ONLY else ~Filters.user() wrapper = MethodWrapper(bot) @@ -27,34 +26,24 @@ def is_flood(updates: list): return sum(genexpr) >= int((len(genexpr) / 100) * FLOOD_PERCENTAGE) -@Client.on_message(FILTER & ~BYPASS_USERS, group=-1) +@Client.on_message(FILTER & ~BYPASS_USERS & ~user_banned(), group=-1) def anti_flood(client, update): """Anti flood module""" user_id = update.from_user.id - if isinstance(update, Message): - VAR = MESSAGES - chat = update.chat.id - date = update.date - message_id = update.message_id - else: - VAR = QUERIES - message_id = None - chat = user_id - if update.message: - date = update.message.date - else: - date = time.time() - if isinstance(VAR[user_id], tuple): - chat, date = VAR[user_id] + chat = update.chat.id + date = update.date + message_id = update.message_id + if isinstance(MESSAGES[user_id], tuple): + chat, date = MESSAGES[user_id] if time.time() - date >= BAN_TIME: - logging.warning(f"{user_id} has waited at least {BAN_TIME} seconds and can now text again") + logging.warning(f"{user_id} has waited at least {BAN_TIME} seconds in {chat} and can now text again") BANNED_USERS.remove(user_id) - del VAR[user_id] - elif len(VAR[user_id]) >= MAX_UPDATE_THRESHOLD: - VAR[user_id].append({chat: (date, message_id)}) + del MESSAGES[user_id] + elif len(MESSAGES[user_id]) >= MAX_UPDATE_THRESHOLD: + MESSAGES[user_id].append({chat: (date, message_id)}) logging.info(f"MAX_UPDATE_THRESHOLD ({MAX_UPDATE_THRESHOLD}) Reached for {user_id}") - user_data = VAR.pop(user_id) + user_data = MESSAGES.pop(user_id) timestamps = [list(*d.values())[0] for d in user_data] updates = [list(*d.values())[1] for d in user_data] if is_flood(timestamps): @@ -62,13 +51,13 @@ def anti_flood(client, update): if user_id in CACHE: del CACHE[user_id] BANNED_USERS.add(user_id) - VAR[user_id] = chat, time.time() + MESSAGES[user_id] = chat, time.time() if FLOOD_NOTICE: wrapper.send_message(user_id, FLOOD_NOTICE) - if DELETE_MESSAGES and any(updates): + if DELETE_MESSAGES: wrapper.delete_messages(chat, filter(bool, updates)) else: - if user_id in VAR: - del VAR[user_id] + if user_id in MESSAGES: + del MESSAGES[user_id] else: - VAR[user_id].append({chat: (date, message_id)}) + MESSAGES[user_id].append({chat: (date, message_id)}) diff --git a/BotBase/modules/livechat.py b/BotBase/modules/livechat.py index e455609..b65c3f9 100644 --- a/BotBase/modules/livechat.py +++ b/BotBase/modules/livechat.py @@ -2,7 +2,7 @@ from pyrogram import Client, Filters, InlineKeyboardMarkup, InlineKeyboardButton 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, NAME, bot, CANNOT_REQUEST_SUPPORT, YES, NO, user_banned, CLOSE_CHAT_BUTTON, BACK_BUTTON, UPDATE_BUTTON + TOO_FAST, CHAT_BUSY, LEAVE_CURRENT_CHAT, USER_JOINS_CHAT, NAME, CANNOT_REQUEST_SUPPORT, YES, NO, user_banned, CLOSE_CHAT_BUTTON, BACK_BUTTON, UPDATE_BUTTON, bot import time from ..database.query import get_user from .antiflood import BANNED_USERS @@ -105,7 +105,7 @@ def close_chat(_, query): del CACHE[user_id] wrapper.send_message(user_id, USER_CLOSES_CHAT.format(user_id=NAME.format(admin_id), user_name=admin_name)) - del CACHE[query.from_user.id][1] + del CACHE[query.from_user.id] else: data = CACHE[query.from_user.id][-1] diff --git a/bot.py b/bot.py index 7d903bd..743eb78 100644 --- a/bot.py +++ b/bot.py @@ -1,21 +1,16 @@ import logging -from pyrogram import CallbackQueryHandler +from pyrogram import Client from pyrogram.session import Session -import importlib +from BotBase.config import bot, LOGGING_LEVEL, LOGGING_FORMAT, DATE_FORMAT, DB_PATH, DB_CREATE +from BotBase.database.query import create_database if __name__ == "__main__": - MODULE_NAME = "BotBase" # Change this to match the FOLDER name that contains the config.py file - conf = importlib.import_module(f"{MODULE_NAME}.config") - bot = conf.bot - antiflood = importlib.import_module(f"{MODULE_NAME}.modules.antiflood") - dbmodule = importlib.import_module(f"{MODULE_NAME}.database.query") - logging.basicConfig(format=conf.LOGGING_FORMAT, datefmt=conf.DATE_FORMAT, level=conf.LOGGING_LEVEL) - bot.add_handler(CallbackQueryHandler(antiflood.anti_flood, ~antiflood.BYPASS_USERS), group=-1) + logging.basicConfig(format=LOGGING_FORMAT, datefmt=DATE_FORMAT, level=LOGGING_LEVEL) Session.notice_displayed = True try: logging.warning("Running create_database()") - dbmodule.create_database(conf.DB_PATH, conf.DB_CREATE) + create_database(DB_PATH, DB_CREATE) logging.warning("Database interaction complete") logging.warning("Starting bot") bot.start()