video-stream/program/updater.py

"""
Video + Music Stream Telegram Bot
Copyright (c) 2022-present levina=lab <https://github.com/levina-lab>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but without any warranty; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/licenses.html>
"""


import os

from git import Repo
from git.exc import InvalidGitRepositoryError, GitCommandError

from pyrogram.types import Message
from pyrogram import Client, filters

from program import LOGS
from config import UPSTREAM_BRANCH, UPSTREAM_REPO, BOT_USERNAME

from driver.filters import command
from driver.decorators import bot_creator
from driver.database.dbqueue import get_active_chats, remove_active_chat


async def install_requirements(cmd: str) -> Tuple[str, str, int, int]:
    args = shlex.split(cmd)
    process = await asyncio.create_subprocess_exec(
        *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await process.communicate()
    return (
        stdout.decode("utf-8", "replace").strip(),
        stderr.decode("utf-8", "replace").strip(),
        process.returncode,
        process.pid,
    )


@Client.on_message(command(["update", f"update@{BOT_USERNAME}"]) & ~filters.edited)
@bot_creator
async def update_bot(_, message: Message):
    chat_id = message.chat.id
    msg = await bot.send_message(chat_id, "❖ Checking updates...")
    try:
        repo = Repo()
    except GitCommandError:
        await msg.edit("Git command error !")
        return
    except InvalidGitRepositoryError:
        await msg.edit("❖ Checking Git updates...")
        repo = Repo.init()
        if "origin" in repo.remotes:
            origin = repo.remote("origin")
        else:
            origin = repo.create_remote("origin", UPSTREAM_REPO)
        origin.fetch()
        repo.create_head(UPSTREAM_BRANCH, origin.refs[UPSTREAM_BRANCH])
        repo.heads[UPSTREAM_BRANCH].set_tracking_branch(
            origin.refs[UPSTREAM_BRANCH]
        )
        repo.heads[UPSTREAM_BRANCH].checkout(True)
        try:
            repo.create_remote("origin", UPSTREAM_REPO)
        except BaseException:
            pass
        nrs = repo.remote("origin")
        nrs.fetch(UPSTREAM_BRANCH)
        try:
            nrs.pull(UPSTREAM_BRANCH)
        except GitCommandError:
            repo.git.reset("--hard", "FETCH_HEAD")
        await install_requirements(
            "pip3 install --no-cache-dir -r requirements.txt"
        )
        await msg.edit("Done, bot has updated !\n\n> now the bot will be rebooted to apply the update and will active again in 5-10 second(s).")
        return
    served_chats = []
    try:
        chats = await get_active_chats()
        for chat in chats:
            served_chats.append(int(chat["chat_id"]))
    except BaseException:
        pass
    for x in served_chats:
        try:
            await remove_active_chat(x)
        except BaseException:
            pass
        return
    os.system(f"kill -9 {os.getpid()} && python3 main.py")


@Client.on_message(command(["restart", f"restart@{BOT_USERNAME}"]) & ~filters.edited)
@bot_creator
async def restart_bot(_, message: Message):
    try:
        msg = await message.reply_text("❖ Restarting bot...")
        LOGS.info("[INFO]: >> BOT SERVER RESTARTED <<")
    except BaseException as err:
        LOGS.info(f"[ERROR]: {err}")
        return
    await msg.edit_text("✅ Bot has restarted !\n\n» back active again in 5-10 seconds.")
    os.system(f"kill -9 {os.getpid()} && python3 main.py")