Add restart session cmd

- Fix some minor errors
- Confirm before restart

Signed-off-by: anasty17 <[email protected]>

README.md

@@ -517,6 +517,7 @@ shell - Run commands in Shell
 aexec - Execute async function
 exec - Execute sync function
 restart - Restart the Bot
+restartses - Restart Telegram Session(s)
 stats - Bot Usage Stats
 ping - Ping the Bot
 help - All cmds with description

bot/__main__.py

@@ -1,5 +1,7 @@
 from signal import signal, SIGINT
 from asyncio import gather
+from pyrogram.filters import regex
+from pyrogram.handlers import CallbackQueryHandler
 
 from .core.config_manager import Config
 
@@ -18,14 +20,41 @@
     update_variables,
 )
 from .helper.ext_utils.telegraph_helper import telegraph
-from .helper.ext_utils.bot_utils import sync_to_async, create_help_buttons
+from .helper.ext_utils.bot_utils import sync_to_async, create_help_buttons, new_task
 from .helper.ext_utils.files_utils import clean_all, exit_clean_up
 from .helper.ext_utils.jdownloader_booter import jdownloader
 from .helper.listeners.aria2_listener import start_aria2_listener
+from .helper.telegram_helper.filters import CustomFilters
 from .helper.mirror_leech_utils.rclone_utils.serve import rclone_serve_booter
+from .helper.telegram_helper.message_utils import (
+    send_message,
+    edit_message,
+    delete_message,
+)
 from .modules import initiate_search_tools, get_packages_version, restart_notification
 
 
+@new_task
+async def restart_sessions_confirm(_, query):
+    data = query.data.split()
+    message = query.message
+    if data[1] == "confirm":
+        reply_to = message.reply_to_message
+        restart_message = await send_message(reply_to, "Restarting Session(s)...")
+        await delete_message(message)
+        await TgClient.reload()
+        add_handlers()
+        TgClient.bot.add_handler(
+            CallbackQueryHandler(
+                restart_sessions_confirm,
+                filters=regex("^sessionrestart") & CustomFilters.sudo,
+            )
+        )
+        await edit_message(restart_message, "Session(s) Restarted Successfully!")
+    else:
+        await delete_message(message)
+
+
 async def main():
     await load_settings()
     await gather(TgClient.start_bot(), TgClient.start_user())
@@ -48,6 +77,12 @@ async def main():
     )
     create_help_buttons()
     add_handlers()
+    TgClient.bot.add_handler(
+        CallbackQueryHandler(
+            restart_sessions_confirm,
+            filters=regex("^sessionrestart") & CustomFilters.sudo,
+        )
+    )
     LOGGER.info("Bot Started!")
     signal(SIGINT, exit_clean_up)
 

bot/core/config_manager.py

@@ -87,16 +87,19 @@ def load(cls):
         for attr in dir(settings):
             if hasattr(cls, attr):
                 value = getattr(settings, attr)
-                if isinstance(value, str):
-                    value = value.strip()
                 if not value:
                     continue
-                if attr == "DEFAULT_UPLOAD" and value != "rc":
-                    value = "gd"
+                if isinstance(value, str):
+                    value = value.strip()
+                if attr == "DEFAULT_UPLOAD" and value != "gd":
+                    value = "rc"
                 elif attr == "DOWNLOAD_DIR" and not value.endswith("/"):
                     value = f"{value}/"
                 elif attr == "USENET_SERVERS":
-                    if not value[0].get("host"):
+                    try:
+                        if not value[0].get("host"):
+                            continue
+                    except:
                         continue
                 setattr(cls, attr, value)
         for key in ["BOT_TOKEN", "OWNER_ID", "TELEGRAM_API", "TELEGRAM_HASH"]:
@@ -110,14 +113,17 @@ def load(cls):
     def load_dict(cls, config_dict):
         for key, value in config_dict.items():
             if hasattr(cls, key):
-                if key == "DEFAULT_UPLOAD" and value != "rc":
-                    value = "gd"
+                if key == "DEFAULT_UPLOAD" and value != "gd":
+                    value = "rc"
                 elif key == "DOWNLOAD_DIR":
                     if not value.endswith("/"):
                         value = f"{value}/"
                 elif key == "USENET_SERVERS":
-                    if not value[0].get("host"):
-                        continue
+                    try:
+                        if not value[0].get("host"):
+                            value = []
+                    except:
+                        value = []
                 setattr(cls, key, value)
         for key in ["BOT_TOKEN", "OWNER_ID", "TELEGRAM_API", "TELEGRAM_HASH"]:
             value = getattr(cls, key)

bot/core/handlers.py

@@ -235,6 +235,18 @@ def add_handlers():
             & CustomFilters.sudo,
         )
     )
+    TgClient.bot.add_handler(
+        CallbackQueryHandler(
+            confirm_restart, filters=regex("^botrestart") & CustomFilters.sudo
+        )
+    )
+    TgClient.bot.add_handler(
+        MessageHandler(
+            restart_sessions,
+            filters=command(BotCommands.RestartSessionsCommand, case_sensitive=True)
+            & CustomFilters.sudo,
+        )
+    )
     TgClient.bot.add_handler(
         MessageHandler(
             ping,

bot/core/mltb_client.py

@@ -60,12 +60,12 @@ async def stop(cls):
                 await cls.bot.stop()
             if cls.user:
                 await cls.user.stop()
-            LOGGER.info("Client stopped")
+            LOGGER.info("Client(s) stopped")
 
     @classmethod
     async def reload(cls):
         async with cls._lock:
             await cls.bot.restart()
             if cls.user:
                 await cls.user.restart()
-            LOGGER.info("Client restarted")
+            LOGGER.info("Client(s) restarted")

bot/core/startup.py

@@ -212,9 +212,6 @@ async def update_variables():
                 else:
                     index_urls.append("")
 
-    if not await aiopath.exists("accounts"):
-        Config.USE_SERVICE_ACCOUNTS = False
-
 
 async def load_configurations():
 
@@ -251,3 +248,6 @@ async def load_configurations():
         ).wait()
         await (await create_subprocess_exec("chmod", "-R", "777", "accounts")).wait()
         await remove("accounts.zip")
+
+    if not await aiopath.exists("accounts"):
+        Config.USE_SERVICE_ACCOUNTS = False

bot/helper/mirror_leech_utils/download_utils/yt_dlp_download.py

@@ -143,9 +143,9 @@ def _extract_meta_data(self):
                     if not entry:
                         continue
                     elif "filesize_approx" in entry:
-                        self._listener.size += entry["filesize_approx"]
+                        self._listener.size += entry.get("filesize_approx", 0)
                     elif "filesize" in entry:
-                        self._listener.size += entry["filesize"]
+                        self._listener.size += entry.get("filesize", 0)
                     if not self._listener.name:
                         outtmpl_ = "%(series,playlist_title,channel)s%(season_number& |)s%(season_number&S|)s%(season_number|)02d.%(ext)s"
                         self._listener.name, ext = ospath.splitext(

bot/helper/mirror_leech_utils/status_utils/jdownloader_status.py

@@ -23,7 +23,7 @@ def _get_combined_info(result):
         bytesTotal += res.get("bytesTotal", 0)
         speed += res.get("speed", 0)
     if len(status) == 0:
-        status = "UnknownError Check Web Interface"
+        status = "UnknownError Check WebUI"
     try:
         eta = (bytesTotal - bytesLoaded) / speed
     except:
@@ -92,10 +92,10 @@ def eta(self):
 
     def status(self):
         async_to_sync(self._update)
-        state = self._info.get("status", "jdlimit")
+        state = self._info.get("status", "jdlimit").capitalize()
         if len(state) == 0:
-            return "UnknownError Check Web Interface"
-        return MirrorStatus.STATUS_QUEUEDL if state == "jdlimit" else state
+            return "UnknownError Check WebUI"
+        return MirrorStatus.STATUS_QUEUEDL if state == "Jdlimit" else state
 
     def task(self):
         return self
@@ -106,7 +106,9 @@ def gid(self):
     async def cancel_task(self):
         self.listener.is_cancelled = True
         LOGGER.info(f"Cancelling Download: {self.name()}")
-        await jdownloader.device.downloads.remove_links(package_ids=jd_downloads[self._gid]["ids"])
+        await jdownloader.device.downloads.remove_links(
+            package_ids=jd_downloads[self._gid]["ids"]
+        )
         async with jd_lock:
             del jd_downloads[self._gid]
         await self.listener.on_download_error("Download cancelled by user!")

bot/helper/telegram_helper/bot_commands.py

@@ -30,6 +30,7 @@ def __init__(self):
         self.RmSudoCommand = f"rmsudo{Config.CMD_SUFFIX}"
         self.PingCommand = f"ping{Config.CMD_SUFFIX}"
         self.RestartCommand = f"restart{Config.CMD_SUFFIX}"
+        self.RestartSessionsCommand = f"restartses{Config.CMD_SUFFIX}"
         self.StatsCommand = f"stats{Config.CMD_SUFFIX}"
         self.HelpCommand = f"help{Config.CMD_SUFFIX}"
         self.LogCommand = f"log{Config.CMD_SUFFIX}"

bot/modules/__init__.py

@@ -19,7 +19,12 @@
     nzb_leech,
     nzb_mirror,
 )
-from .restart import restart_bot, restart_notification
+from .restart import (
+    restart_bot,
+    restart_notification,
+    confirm_restart,
+    restart_sessions,
+)
 from .rss import get_rss_menu, rss_listener
 from .search import torrent_search, torrent_search_update, initiate_search_tools
 from .services import start, ping, log
@@ -62,6 +67,8 @@
     "nzb_mirror",
     "restart_bot",
     "restart_notification",
+    "confirm_restart",
+    "restart_sessions",
     "get_rss_menu",
     "rss_listener",
     "torrent_search",

bot/modules/bot_settings.py

@@ -901,6 +901,9 @@ async def load_config():
     index_urls.clear()
     await update_variables()
 
+    if not await aiopath.exists("accounts"):
+        Config.USE_SERVICE_ACCOUNTS = False
+
     if len(task_dict) != 0 and (st := intervals["status"]):
         for key, intvl in list(st.items()):
             intvl.cancel()