refactor(sudo): simplify reboot and update commands, enhance restart logic

Author: HitaloM

src/korone/client.py

@@ -16,6 +16,7 @@
 from hydrogram.errors import MessageIdInvalid, MessageNotModified
 from hydrogram.raw.all import layer
 
+from korone import constants
 from korone.config import ConfigManager
 from korone.utils.backup import do_backup
 
@@ -96,7 +97,7 @@ async def start(self) -> None:
         if backups_chat:
             aiocron.crontab("0 * * * *", do_backup, loop=self.loop, args=(self, backups_chat))
 
-        reboot_data: dict[str, Any] | None = await cache.get("korone-reboot")
+        reboot_data: dict[str, Any] | None = await cache.get(constants.REBOOT_CACHE_KEY)
         if reboot_data:
             with suppress(MessageNotModified, MessageIdInvalid, KeyError):
                 chat_id = reboot_data["chat_id"]

src/korone/constants.py

@@ -98,3 +98,6 @@
 PRIVACY_POLICY_URL: str = f"{DOCS_URL}/en/latest/privacy.html"
 
 MESSAGE_LENGTH_LIMIT: int = 4096
+
+REBOOT_CACHE_KEY: str = "korone-reboot"
+REBOOT_CACHE_TTL: int = 300

src/korone/modules/sudo/handlers/reboot.py

@@ -1,35 +1,15 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2025 Hitalo M. <https://github.com/HitaloM>
 
-import os
-import sys
-import time
-from typing import TypedDict
-
 from hydrogram import Client
 from hydrogram.types import Message
 
 from korone.decorators import router
 from korone.filters import Command, IsSudo
-from korone.utils.caching import cache
-
-
-class RebootCache(TypedDict):
-    chat_id: int
-    message_id: int
-    time: float
+from korone.modules.sudo.utils import restart_bot
 
 
 @router.message(Command("reboot", disableable=False) & IsSudo)
 async def reboot_command(client: Client, message: Message) -> None:
-    cache_key = "korone-reboot"
-
-    if await cache.get(cache_key):
-        await cache.delete(cache_key)
-
-    sent = await message.reply("Rebooting...")
-
-    value: RebootCache = {"chat_id": message.chat.id, "message_id": sent.id, "time": time.time()}
-    await cache.set(cache_key, value=value, expire=300)
-
-    os.execv(sys.executable, [sys.executable, "-m", "korone"])
+    sent = await message.reply("Starting reboot...")
+    await restart_bot(sent, "Rebooting...")

src/korone/modules/sudo/handlers/updater.py

@@ -16,8 +16,8 @@
     generate_document,
     parse_commits,
     perform_update,
+    restart_bot,
 )
-from korone.utils.caching import cache
 
 
 @router.message(Command(commands=["update", "upgrade"], disableable=False) & IsSudo)
@@ -27,7 +27,7 @@ async def update_command(client: Client, message: Message) -> None:
     try:
         stdout = await fetch_updates()
     except Exception as e:
-        await sent.edit(f"An error occurred:\n<code>{e}</code>")
+        await sent.edit(f"An error occurred:\n<code>{html.escape(str(e))}</code>")
         return
 
     if not stdout.strip():
@@ -45,23 +45,29 @@ async def update_command(client: Client, message: Message) -> None:
 
 @router.callback_query(UpdateCallbackData.filter() & IsSudo)
 async def update_callback(client: Client, callback: CallbackQuery) -> None:
-    cache_key = "korone-reboot"
-    if await cache.get(cache_key):
-        await cache.delete(cache_key)
-
     message = callback.message
+
+    await callback.answer("Starting upgrade...")
     await message.edit_reply_markup()
     sent = await message.reply("Upgrading...")
 
     try:
         stdout = await perform_update()
     except Exception as e:
-        await sent.edit(f"An error occurred:\n<code>{e}</code>")
+        await sent.edit(f"An error occurred:\n<code>{html.escape(str(e))}</code>")
+        await callback.answer("Upgrade failed.", show_alert=True)
         return
 
     if len(stdout) > MESSAGE_LENGTH_LIMIT:
-        await sent.edit("Upgrade completed successfully. Reboot is required...")
         await generate_document(stdout, message)
+        final_text = "Upgrade completed successfully. Output sent as document. Rebooting..."
+        await restart_bot(sent, final_text)
         return
 
-    await sent.edit(f"<pre language='bash'>{html.escape(stdout)}</pre>")
+    formatted_stdout = html.escape(stdout) if stdout else "No output returned."
+    final_text = (
+        "Upgrade completed successfully.\n"
+        f"<pre language='bash'>{formatted_stdout}</pre>\n"
+        "Rebooting..."
+    )
+    await restart_bot(sent, final_text)

src/korone/modules/sudo/utils.py

@@ -3,17 +3,29 @@
 
 import html
 import io
+import os
+import sys
+import time
 from subprocess import PIPE
-from typing import Any, TypedDict
+from typing import Any, Never, TypedDict
 
 from anyio import run_process
 from hydrogram.types import Message
 
+from korone import constants
+from korone.utils.caching import cache
+
 
 class CommitInfo(TypedDict):
     title: str
 
 
+class RebootCache(TypedDict):
+    chat_id: int
+    message_id: int
+    time: float
+
+
 async def generate_document(output: Any, message: Message) -> None:
     with io.BytesIO(str.encode(str(output))) as file:
         file.name = "output.txt"
@@ -56,3 +68,34 @@ def generate_changelog(commits: dict[str, CommitInfo]) -> str:
 async def perform_update() -> str:
     commands = ["git reset --hard origin/main", "pybabel compile -d locales -D bot", "uv sync"]
     return "".join([await run_command(command) for command in commands])
+
+
+async def restart_bot(message: Message, text: str) -> Never:
+    await cache.delete(constants.REBOOT_CACHE_KEY)
+    await message.edit_text(text)
+
+    value: RebootCache = {
+        "chat_id": message.chat.id,
+        "message_id": message.id,
+        "time": time.time(),
+    }
+    await cache.set(constants.REBOOT_CACHE_KEY, value=value, expire=constants.REBOOT_CACHE_TTL)
+
+    main_module = sys.modules.get("__main__")
+    module_name = None
+
+    if main_module is not None:
+        if package_name := getattr(main_module, "__package__", None):
+            module_name = package_name
+        elif (spec := getattr(main_module, "__spec__", None)) and (
+            spec_name := getattr(spec, "name", None)
+        ):
+            module_name = spec_name.removesuffix(".__main__")
+
+    if module_name:
+        args = [sys.executable, "-m", module_name, *sys.argv[1:]]
+    else:
+        script_candidate = sys.argv[0] or getattr(main_module, "__file__", "")
+        args = [sys.executable, script_candidate, *sys.argv[1:]]
+
+    os.execv(sys.executable, args)