Fixed error after major refactoring

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
+

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

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

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

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)})

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]

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()