PagerMaid-Pyro/pagermaid/modules/system.py

import io
import sys
import traceback

from sys import exit
from platform import node
from getpass import getuser

from pagermaid import bot
from pagermaid.listener import listener
from pagermaid.utils import attach_log, execute, Message, lang


@listener(is_plugin=False, command="sh",
          need_admin=True,
          description=lang('sh_des'),
          parameters=lang('sh_parameters'))
async def sh(message: Message):
    """ Use the command-line from Telegram. """
    user = getuser()
    command = message.arguments
    hostname = node()

    if not command:
        await message.edit(lang('arg_error'))
        return

    message = await message.edit(
        f"`{user}`@{hostname} ~"
        f"\n> `$` {command}"
    )

    result = await execute(command)

    if result:
        if len(result) > 4096:
            await attach_log(result, message.chat.id, "output.log", message.id)
            return

        await message.edit(
            f"`{user}`@{hostname} ~"
            f"\n> `#` {command}"
            f"\n`{result}`"
        )
    else:
        return


@listener(is_plugin=False, command="restart",
          need_admin=True,
          description=lang('restart_des'))
async def restart(message: Message):
    """ To re-execute PagerMaid. """
    if not message.text[0].isalpha():
        await message.edit(lang('restart_log'))
        exit(0)


@listener(is_plugin=False, command="eval",
          need_admin=True,
          description=lang('eval_des'),
          parameters=lang('eval_parameters'))
async def sh_eval(message: Message):
    """ Run python commands from Telegram. """
    try:
        cmd = message.text.split(" ", maxsplit=1)[1]
    except IndexError:
        return await message.edit(lang('eval_need_dev'))
    old_stderr = sys.stderr
    old_stdout = sys.stdout
    redirected_output = sys.stdout = io.StringIO()
    redirected_error = sys.stderr = io.StringIO()
    stdout, stderr, exc = None, None, None
    try:
        await aexec(cmd, message, bot)
    except Exception:  # noqa
        exc = traceback.format_exc()
    stdout = redirected_output.getvalue()
    stderr = redirected_error.getvalue()
    sys.stdout = old_stdout
    sys.stderr = old_stderr
    if exc:
        evaluation = exc
    elif stderr:
        evaluation = stderr
    elif stdout:
        evaluation = stdout
    else:
        evaluation = "Success"
    final_output = f"**>>>** ```{cmd}``` \n```{evaluation}```"
    if len(final_output) > 4096:
        message = await message.edit(f"**>>>** ```{cmd}```")
        await attach_log(evaluation, message.chat.id, "output.log", message.id)
    else:
        await message.edit(final_output)


async def aexec(code, event, client):
    exec(
        (
                (
                        ("async def __aexec(e, client): " + "\n msg = message = e")
                        + "\n reply = message.reply_to_message"
                )
                + "\n chat = e.chat"
        )
        + "".join(f"\n {x}" for x in code.split("\n"))
    )

    return await locals()["__aexec"](event, client)