TuneBot/cogs/owner.py

import asyncio
import io
import textwrap
import time
import traceback
from asyncio.subprocess import PIPE
from contextlib import redirect_stdout
from io import BytesIO
from platform import python_version

import discord
from discord.ext import commands
from discord.ext.commands.context import Context

from bot import ChristmasBot


class OwnerCog(commands.Cog):
    def __init__(self, bot: ChristmasBot):
        self.bot = bot
        self._last_result = None

    @staticmethod
    def cleanup_code(content):
        if content.startswith("```") and content.endswith("```"):
            return "\n".join(content.split("\n")[1:-1])
        return content.strip("` \n")

    # Hidden means it won't show up on the default help.
    @commands.command(name="load", hidden=True, slash_command=False)
    @commands.is_owner()
    async def _cog_load(self, ctx: Context, *, cog: str):
        """Command which Loads a Module."""

        try:
            self.bot.load_extension(cog)
        except Exception as e:
            await ctx.send(f"**`ERROR:`** {type(e).__name__} - {e}")
        else:
            await ctx.send("**`SUCCESS`**")

    @commands.command(name="unload", hidden=True, slash_command=False)
    @commands.is_owner()
    async def _cog_unload(self, ctx, *, cog: str):
        """Command which Unloads a Module."""

        try:
            self.bot.unload_extension(cog)
        except Exception as e:
            await ctx.send(f"**`ERROR:`** {type(e).__name__} - {e}")
        else:
            await ctx.send("**`SUCCESS`**")

    @commands.command(name="reload", hidden=True, slash_command=False)
    @commands.is_owner()
    async def _cog_reload(self, ctx, *, cog: str):
        """Command which Reloads a Module."""

        try:
            self.bot.unload_extension(cog)
            self.bot.load_extension(cog)
        except Exception as e:
            await ctx.send(f"**`ERROR:`** {type(e).__name__} - {e}")
        else:
            await ctx.send("**`SUCCESS`**")

    @commands.command(name="shutdown", hidden=True, slash_command=False)
    @commands.is_owner()
    async def shutdown(self, ctx):
        """Command which shutdowns the bot."""
        await ctx.bot.logout()

    @commands.is_owner()
    @commands.command(
        pass_context=True,
        hidden=True,
        name="eval",
        aliases=["evaluate"],
        slash_command=False,
    )
    async def _eval(self, ctx, *, body: str):
        env = {
            "bot": self.bot,
            "ctx": ctx,
            "channel": ctx.channel,
            "author": ctx.author,
            "guild": ctx.guild,
            "message": ctx.message,
            "_": self._last_result,
        }

        env.update(globals())

        body = self.cleanup_code(body)
        stdout = io.StringIO()

        to_compile = f'async def func():\n{textwrap.indent(body, "  ")}'
        # await ctx.message.add_reaction('a:loading:452489773396000778')

        try:
            exec(to_compile, env)
        except Exception as e:
            # await ctx.message.add_reaction('naokoerror:447495055603662849')
            fooem = discord.Embed(color=0xFF0000)
            fooem.add_field(
                name="Code evaluation was not successful.",
                value=f"```\n{e.__class__.__name__}: {e}\n```",
            )
            fooem.set_footer(
                text=f"Evaluated using Python {python_version()}",
                icon_url="http://i.imgur.com/9EftiVK.png",
            )
            await ctx.send(embed=fooem)
            # await ctx.message.remove_reaction('a:loading:452489773396000778', member=ctx.me)

        func = env["func"]
        try:
            with redirect_stdout(stdout):
                ret = await func()
        except Exception as e:
            value = stdout.getvalue()
            # await ctx.message.add_reaction('naokoerror:447495055603662849')
            fooem = discord.Embed(color=0xFF0000)
            fooem.add_field(
                name="Code evaluation was not successful.",
                value=f"```\n{value}{traceback.format_exc()}\n```",
            )
            fooem.set_footer(
                text=f"Evaluated using Python {python_version()}",
                icon_url="http://i.imgur.com/9EftiVK.png",
            )
            await ctx.send(embed=fooem)
            try:
                # await ctx.message.remove_reaction('a:loading:452489773396000778', member=ctx.me)
                # await ctx.message.add_reaction('naokotick:447494238872141827')
                pass
            except Exception:
                pass
        else:
            value = stdout.getvalue()
            try:
                await ctx.message.remove_reaction(
                    "a:loading:452489773396000778", member=ctx.me
                )
                # await ctx.message.add_reaction('naokotick:447494238872141827')
            except Exception:
                pass

            if ret is None:
                if value:
                    sfooem = discord.Embed(color=0x170041)
                    sfooem.add_field(
                        name="Code evaluation was successful!",
                        value=f"```\n{value}\n```",
                    )
                    sfooem.set_footer(
                        text=f"Evaluated using Python {python_version()}",
                        icon_url="http://i.imgur.com/9EftiVK.png",
                    )
                    await ctx.send(embed=sfooem)
            else:
                self._last_result = ret
                ssfooem = discord.Embed(color=0x170041)
                ssfooem.add_field(
                    name="Code evaluation was successful!",
                    value=f"```\n{value}{ret}\n```",
                )
                ssfooem.set_footer(
                    text=f"Evaluated using Python {python_version()}",
                    icon_url="http://i.imgur.com/9EftiVK.png",
                )
                await ctx.send(embed=ssfooem)

    @commands.is_owner()
    @commands.command(hidden=True, aliases=["exec"], slash_command=False)
    async def execute(self, ctx, *, text: str):
        """Do a shell command."""
        message = await ctx.send(f"Loading...")
        proc = await asyncio.create_subprocess_shell(
            text, stdin=None, stderr=PIPE, stdout=PIPE
        )
        out = (await proc.stdout.read()).decode("utf-8").strip()
        err = (await proc.stderr.read()).decode("utf-8").strip()

        if not out and not err:
            await message.delete()
            return await ctx.message.add_reaction("👌")

        content = ""

        if err:
            content += f"Error:\r\n{err}\r\n{'-' * 30}\r\n"
        if out:
            content += out

        if len(content) > 1500:
            try:
                data = BytesIO(content.encode("utf-8"))
                await message.delete()
                await ctx.send(
                    content=f"The result was a bit too long..",
                    file=discord.File(data, filename=f"result_{int(time.time())}.txt"),
                )
            except asyncio.TimeoutError as e:
                await message.delete()
                return await ctx.send(e)
        else:
            await message.edit(content=f"```fix\n{content}\n```")


def setup(bot: ChristmasBot):
    bot.add_cog(OwnerCog(bot))