Add restart session cmd

- Fix some minor errors
- Confirm before restart

Signed-off-by: anasty17 <e.anastayyar@gmail.com>
This commit is contained in:
anasty17 2024-12-23 20:54:08 +02:00
parent f2e3be4bf9
commit 03f3276dc3
No known key found for this signature in database
GPG Key ID: EAAC8A2DF2861DE4
12 changed files with 180 additions and 73 deletions

View File

@ -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

View File

@ -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 @@ from .core.startup import (
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)

View File

@ -87,16 +87,19 @@ class Config:
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 @@ class Config:
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)

View File

@ -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,

View File

@ -60,7 +60,7 @@ class TgClient:
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):
@ -68,4 +68,4 @@ class TgClient:
await cls.bot.restart()
if cls.user:
await cls.user.restart()
LOGGER.info("Client restarted")
LOGGER.info("Client(s) restarted")

View File

@ -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

View File

@ -143,9 +143,9 @@ class YoutubeDLHelper:
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(

View File

@ -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 @@ class JDownloaderStatus:
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 @@ class JDownloaderStatus:
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!")

View File

@ -30,6 +30,7 @@ class _BotCommands:
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}"

View File

@ -19,7 +19,12 @@ from .mirror_leech import (
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 @@ __all__ = [
"nzb_mirror",
"restart_bot",
"restart_notification",
"confirm_restart",
"restart_sessions",
"get_rss_menu",
"rss_listener",
"torrent_search",

View File

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

View File

@ -6,46 +6,57 @@ from os import execl as osexecl
from .. import intervals, scheduler, sabnzbd_client, LOGGER
from ..helper.ext_utils.bot_utils import new_task, sync_to_async
from ..helper.telegram_helper.message_utils import send_message
from ..helper.telegram_helper.message_utils import (
send_message,
delete_message,
edit_message,
)
from ..helper.ext_utils.db_handler import database
from ..helper.ext_utils.files_utils import clean_all
from ..helper.telegram_helper import button_build
from ..core.mltb_client import TgClient
from ..core.config_manager import Config
@new_task
async def restart_bot(_, message):
intervals["stopAll"] = True
restart_message = await send_message(message, "Restarting...")
if scheduler.running:
scheduler.shutdown(wait=False)
if qb := intervals["qb"]:
qb.cancel()
if jd := intervals["jd"]:
jd.cancel()
if nzb := intervals["nzb"]:
nzb.cancel()
if st := intervals["status"]:
for intvl in list(st.values()):
intvl.cancel()
await sync_to_async(clean_all)
if sabnzbd_client.LOGGED_IN:
await gather(
sabnzbd_client.pause_all(),
sabnzbd_client.purge_all(True),
sabnzbd_client.delete_history("all", delete_files=True),
)
proc1 = await create_subprocess_exec(
"pkill",
"-9",
"-f",
"gunicorn|aria2c|qbittorrent-nox|ffmpeg|rclone|java|sabnzbdplus",
buttons = button_build.ButtonMaker()
buttons.data_button("Yes!", "botrestart confirm")
buttons.data_button("Cancel", "restart cancel")
button = buttons.build_menu(2)
await send_message(message, "Are you sure you want to restart the bot ?!", button)
@new_task
async def restart_sessions(_, message):
buttons = button_build.ButtonMaker()
buttons.data_button("Yes!", "sessionrestart confirm")
buttons.data_button("Cancel", "restart cancel")
button = buttons.build_menu(2)
await send_message(
message, "Are you sure you want to restart the session(s) ?!", button
)
proc2 = await create_subprocess_exec("python3", "update.py")
await gather(proc1.wait(), proc2.wait())
async with aiopen(".restartmsg", "w") as f:
await f.write(f"{restart_message.chat.id}\n{restart_message.id}\n")
osexecl(executable, executable, "-m", "bot")
async def send_incomplete_task_message(cid, msg_id, msg):
try:
if msg.startswith("Restarted Successfully!"):
await TgClient.bot.edit_message_text(
chat_id=cid,
message_id=msg_id,
text=msg,
disable_web_page_preview=True,
)
await remove(".restartmsg")
else:
await TgClient.bot.send_message(
chat_id=cid,
text=msg,
disable_web_page_preview=True,
disable_notification=True,
)
except Exception as e:
LOGGER.error(e)
async def restart_notification():
@ -55,23 +66,6 @@ async def restart_notification():
else:
chat_id, msg_id = 0, 0
async def send_incomplete_task_message(cid, msg):
try:
if msg.startswith("Restarted Successfully!"):
await TgClient.bot.edit_message_text(
chat_id=chat_id, message_id=msg_id, text=msg
)
await remove(".restartmsg")
else:
await TgClient.bot.send_message(
chat_id=cid,
text=msg,
disable_web_page_preview=True,
disable_notification=True,
)
except Exception as e:
LOGGER.error(e)
if Config.INCOMPLETE_TASK_NOTIFIER and Config.DATABASE_URL:
if notifier_dict := await database.get_incomplete_tasks():
for cid, data in notifier_dict.items():
@ -81,10 +75,10 @@ async def restart_notification():
for index, link in enumerate(links, start=1):
msg += f" <a href='{link}'>{index}</a> |"
if len(msg.encode()) > 4000:
await send_incomplete_task_message(cid, msg)
await send_incomplete_task_message(cid, msg_id, msg)
msg = ""
if msg:
await send_incomplete_task_message(cid, msg)
await send_incomplete_task_message(cid, msg_id, msg)
if await aiopath.isfile(".restartmsg"):
try:
@ -94,3 +88,49 @@ async def restart_notification():
except:
pass
await remove(".restartmsg")
@new_task
async def confirm_restart(_, query):
await query.answer()
data = query.data.split()
message = query.message
await delete_message(message)
if data[1] == "confirm":
reply_to = message.reply_to_message
intervals["stopAll"] = True
restart_message = await send_message(reply_to, "Restarting...")
await delete_message(message)
await TgClient.stop()
if scheduler.running:
scheduler.shutdown(wait=False)
if qb := intervals["qb"]:
qb.cancel()
if jd := intervals["jd"]:
jd.cancel()
if nzb := intervals["nzb"]:
nzb.cancel()
if st := intervals["status"]:
for intvl in list(st.values()):
intvl.cancel()
await sync_to_async(clean_all)
if sabnzbd_client.LOGGED_IN:
await gather(
sabnzbd_client.pause_all(),
sabnzbd_client.purge_all(True),
sabnzbd_client.delete_history("all", delete_files=True),
)
proc1 = await create_subprocess_exec(
"pkill",
"-9",
"-f",
"gunicorn|aria2c|qbittorrent-nox|ffmpeg|rclone|java|sabnzbdplus",
)
proc2 = await create_subprocess_exec("python3", "update.py")
await gather(proc1.wait(), proc2.wait())
async with aiopen(".restartmsg", "w") as f:
await f.write(f"{restart_message.chat.id}\n{restart_message.id}\n")
osexecl(executable, executable, "-m", "bot")
else:
await delete_message(message)