diff --git a/bot/core/handlers.py b/bot/core/handlers.py index 9d82ad078..86cc3f468 100644 --- a/bot/core/handlers.py +++ b/bot/core/handlers.py @@ -185,6 +185,11 @@ def add_handlers(): BotCommands.SpeedTest, CustomFilters.authorized, ), + "broadcast": ( + broadcast, + BotCommands.BroadcastCommand, + CustomFilters.owner, + ), } for handler_func, command_name, custom_filter in command_filters.values(): diff --git a/bot/helper/ext_utils/db_handler.py b/bot/helper/ext_utils/db_handler.py index 7f1db6022..1d765275a 100644 --- a/bot/helper/ext_utils/db_handler.py +++ b/bot/helper/ext_utils/db_handler.py @@ -168,6 +168,23 @@ async def add_incomplete_task(self, cid, link, tag): {"_id": link, "cid": cid, "tag": tag}, ) + async def get_pm_uids(self): + if self._return: + return None + return [doc["_id"] async for doc in self._db.pm_users[TgClient.ID].find({})] + + async def update_pm_users(self, user_id): + if self._return: + return + if not bool(await self._db.pm_users[TgClient.ID].find_one({"_id": user_id})): + await self._db.pm_users[TgClient.ID].insert_one({"_id": user_id}) + LOGGER.info(f"New PM User Added : {user_id}") + + async def rm_pm_user(self, user_id): + if self._return: + return + await self._db.pm_users[TgClient.ID].delete_one({"_id": user_id}) + async def update_user_tdata(self, user_id, token, time): if self._return: return diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py index 9e11cff71..12bd62bb6 100644 --- a/bot/helper/telegram_helper/bot_commands.py +++ b/bot/helper/telegram_helper/bot_commands.py @@ -37,6 +37,7 @@ def __init__(self): self.BotSetCommand = f"botsettings{i}" self.UserSetCommand = f"settings{i}" self.SpeedTest = f"speedtest{i}" + self.BroadcastCommand = [f"broadcast{i}", "broadcastall"] self.SelectCommand = f"sel{i}" self.RssCommand = f"rss{i}" diff --git a/bot/modules/__init__.py b/bot/modules/__init__.py index caf11c9ec..3783280e6 100644 --- a/bot/modules/__init__.py +++ b/bot/modules/__init__.py @@ -1,4 +1,5 @@ from .bot_settings import edit_bot_settings, send_bot_settings +from .broadcast import broadcast from .cancel_task import cancel, cancel_all_buttons, cancel_all_update, cancel_multi from .chat_permission import add_sudo, authorize, remove_sudo, unauthorize from .clone import clone_node @@ -41,6 +42,7 @@ "authorize", "bot_help", "bot_stats", + "broadcast", "cancel", "cancel_all_buttons", "cancel_all_update", diff --git a/bot/modules/broadcast.py b/bot/modules/broadcast.py index 626dcdc58..1cdb7524d 100644 --- a/bot/modules/broadcast.py +++ b/bot/modules/broadcast.py @@ -7,7 +7,7 @@ from bot import bot from bot.helper.ext_utils.bot_utils import new_task -from bot.helper.ext_utils.db_handler import Database +from bot.helper.ext_utils.db_handler import database from bot.helper.ext_utils.status_utils import get_readable_time from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters @@ -28,7 +28,7 @@ async def broadcast(_, message): updater = time() broadcast_message = await send_message(message, "Broadcast in progress...") - for uid in await Database.get_pm_uids(): + for uid in await database.get_pm_uids(): try: await message.reply_to_message.copy(uid) successful += 1 @@ -37,7 +37,7 @@ async def broadcast(_, message): await message.reply_to_message.copy(uid) successful += 1 except (UserIsBlocked, InputUserDeactivated): - await Database.rm_pm_user(uid) + await database.rm_pm_user(uid) blocked += 1 except Exception: unsuccessful += 1 @@ -63,11 +63,3 @@ def generate_status(total, successful, blocked, unsuccessful, elapsed_time=""): if elapsed_time: status += f"\n\nElapsed Time: {elapsed_time}" return status - - -bot.add_handler( - MessageHandler( - broadcast, - filters=command(BotCommands.BroadcastCommand) & CustomFilters.owner, - ), -) diff --git a/bot/modules/mirror_leech.py b/bot/modules/mirror_leech.py index 279ed56f6..1ac3b0b6a 100644 --- a/bot/modules/mirror_leech.py +++ b/bot/modules/mirror_leech.py @@ -299,8 +299,15 @@ async def new_event(self): file_ = None if ( - self.link and (is_magnet(self.link) or self.link.endswith(".torrent")) - ) or (file_ and file_.file_name.endswith(".torrent")): + self.link + and ( + is_magnet(self.link) + or self.link.endswith(".torrent") + ) + ) or ( + file_ + and getattr(file_, "file_name", "").endswith(".torrent") + ): self.is_qbit = True if ( diff --git a/bot/modules/speedtest.py b/bot/modules/speedtest.py index c0b72c685..7512850bd 100644 --- a/bot/modules/speedtest.py +++ b/bot/modules/speedtest.py @@ -1,6 +1,7 @@ from speedtest import Speedtest from bot import LOGGER +from bot.core.aeon_client import TgClient from bot.helper.ext_utils.bot_utils import new_task from bot.helper.ext_utils.status_utils import get_readable_file_size from bot.helper.telegram_helper.message_utils import ( @@ -21,7 +22,7 @@ def get_speedtest_results(): test.upload() return test.results - result = get_speedtest_results() + result = await TgClient.bot.loop.run_in_executor(None, get_speedtest_results) if not result: await edit_message(speed, "Speedtest failed to complete.")