TuneBot/cogs/owner.py

import discord
from discord.ext import commands

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


class OwnerCog(commands.Cog):
    def __init__(self, bot):
        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)
    @commands.is_owner()
    async def _cog_load(self, ctx, *, 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)
    @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)
    @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)
    @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'])
    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
        }

        if "import os" in body and ctx.author.id == 190875175460405249:
            return await ctx.send("Ah ah ah, you didn't say the magic word.")

        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'])
    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):
    bot.add_cog(OwnerCog(bot))