<filename>cogs/owner.py
import io
import time
import asyncio
import inspect
import textwrap
import traceback
from contextlib import redirect_stdout

import discord
from discord.ext import commands

import utils


class Owner:
    """Owner-only commands used to maintain the bot."""

    def __init__(self):
        self.sessions = set()
        self._last_result = None

    async def __local_check(self, ctx):
        return await ctx.bot.is_owner(ctx.author)

    async def __error(self, ctx, exception):
        if isinstance(exception, commands.CheckFailure):
            await ctx.send('Only the bot owner may use this command.')

    def cleanup_code(self, content):
        if content.startswith('```') and content.endswith('```'):
            return '\n'.join(content.split('\n')[1:-1])

        return content.strip('` \n')

    def get_syntax_error(self, exception):
        if exception.text is None:
            return '```py\n{0.__class__.__name__}: {0}\n```'.format(exception)
        return '```py\n{0.text}{1:>{0.offset}}\n{0.__class__.__name__}: {0}```'.format(exception, '^')

    @commands.command()
    async def logout(self, ctx):
        """Logs out of the bot."""

        await ctx.bot.logout()

    @commands.command()
    async def load(self, ctx, *, module: str):
        """Loads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.load_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    async def unload(self, ctx, *, module: str):
        """Unloads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.unload_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    async def reload(self, ctx, *, module: str):
        """Reloads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.unload_extension(module)
            ctx.bot.load_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command(name='eval')
    async def _eval(self, ctx, *, body: str):
        """Evaluates code."""

        env = {
            'bot': ctx.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, "  ")}'

        try:
            exec(to_compile, env)
        except SyntaxError as e:
            return await ctx.send(self.get_syntax_error(e))

        func = env['func']
        try:
            with redirect_stdout(stdout):
                ret = await func()
        except Exception as e:
            value = stdout.getvalue()
            await ctx.send(f'```py\n{value}{traceback.format_exc(chain=False)}\n```')
        else:
            value = stdout.getvalue()
            try:
                await ctx.message.add_reaction('\u2705')
            except:
                pass

            if ret is None:
                if value:
                    await ctx.send(f'```py\n{value}\n```')
            else:
                self._last_result = ret
                await ctx.send(f'```py\n{value}{ret}\n```')

    @commands.command()
    async def repl(self, ctx):
        """Launches an interactive REPL session."""

        variables = {
            'ctx': ctx,
            'bot': ctx.bot,
            'message': ctx.message,
            'guild': ctx.guild,
            'channel': ctx.channel,
            'author': ctx.author,
            '_': None,
        }

        if ctx.channel.id in self.sessions:
            return await ctx.send('Already running a REPL session in this channel. Exit it with `quit`')

        self.sessions.add(ctx.channel.id)
        await ctx.send('Enter code to execute or evaluate. `exit()` or `quit` to exit.')

        def check(m):
            return m.author == ctx.author and \
                    m.channel == ctx.channel and \
                    m.content.startswith('`')

        while True:
            try:
                response = await ctx.bot.wait_for('message', check=check, timeout=10.0 * 60.0)
            except asyncio.TimeoutError:
                await ctx.send('Exiting REPL session.')
                self.sessions.remove(ctx.channel.id)
                break

            cleaned = self.cleanup_code(response.content)

            if cleaned in ('quit', 'exit', 'exit()'):
                await ctx.send('Exiting.')
                self.sessions.remove(ctx.channel.id)
                break

            executor = exec
            if cleaned.count('\n') == 0:
                try:
                    code = compile(cleaned, '<repl session>', 'eval')
                except SyntaxError:
                    pass
                else:
                    executor = eval

            if executor is exec:
                try:
                    code = compile(cleaned, '<repl session>', 'exec')
                except SyntaxError as e:
                    await ctx.send(self.get_syntax_error(e))
                    continue

            variables['message'] = response

            fmt = None
            stdout = io.StringIO()

            try:
                with redirect_stdout(stdout):
                    result = executor(code, variables)
                    if inspect.isawaitable(result):
                        result = await result
            except Exception as e:
                value = stdout.getvalue()
                fmt = f'```py\n{value}{traceback.format_exc(chain=False)}\n```'
            else:
                value = stdout.getvalue()
                if result is not None:
                    fmt = f'```py\n{value}{result}\n```'
                    variables['_'] = result
                elif value:
                    fmt = f'```py\n{value}\n```'

            try:
                if fmt is not None:
                    if len(fmt) > 2000:
                        await ctx.send('Content too big to be printed.')
                    else:
                        await ctx.send(fmt)
            except discord.Forbidden:
                pass
            except discord.HTTPException as e:
                await ctx.send(f'Unexpected error: `{e}`')


def setup(bot):
    bot.add_cog(Owner())

<filename>cogs/owner.py
import io
import time
import asyncio
import inspect
import textwrap
import traceback
from contextlib import redirect_stdout

import discord
from discord.ext import commands

import utils


class Owner:
    """Owner-only commands used to maintain the bot."""

    def __init__(self):
        self.sessions = set()
        self._last_result = None

    async def __local_check(self, ctx):
        return await ctx.bot.is_owner(ctx.author)

    async def __error(self, ctx, exception):
        if isinstance(exception, commands.CheckFailure):
            await ctx.send('Only the bot owner may use this command.')

    def cleanup_code(self, content):
        if content.startswith('```') and content.endswith('```'):
            return '\n'.join(content.split('\n')[1:-1])

        return content.strip('` \n')

    def get_syntax_error(self, exception):
        if exception.text is None:
            return '```py\n{0.__class__.__name__}: {0}\n```'.format(exception)
        return '```py\n{0.text}{1:>{0.offset}}\n{0.__class__.__name__}: {0}```'.format(exception, '^')

    @commands.command()
    async def logout(self, ctx):
        """Logs out of the bot."""

        await ctx.bot.logout()

    @commands.command()
    async def load(self, ctx, *, module: str):
        """Loads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.load_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    async def unload(self, ctx, *, module: str):
        """Unloads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.unload_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    async def reload(self, ctx, *, module: str):
        """Reloads a module."""

        module = f'cogs.{module}'
        try:
            ctx.bot.unload_extension(module)
            ctx.bot.load_extension(module)
        except Exception as e:
            await ctx.send(f'```py\n{type(e).__name__}: {e}\n```')
        else:
            await ctx.send('\N{OK HAND SIGN}')

    @commands.command(name='eval')
    async def _eval(self, ctx, *, body: str):
        """Evaluates code."""

        env = {
            'bot': ctx.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, "  ")}'

        try:
            exec(to_compile, env)
        except SyntaxError as e:
            return await ctx.send(self.get_syntax_error(e))

        func = env['func']
        try:
            with redirect_stdout(stdout):
                ret = await func()
        except Exception as e:
            value = stdout.getvalue()
            await ctx.send(f'```py\n{value}{traceback.format_exc(chain=False)}\n```')
        else:
            value = stdout.getvalue()
            try:
                await ctx.message.add_reaction('\u2705')
            except:
                pass

            if ret is None:
                if value:
                    await ctx.send(f'```py\n{value}\n```')
            else:
                self._last_result = ret
                await ctx.send(f'```py\n{value}{ret}\n```')

    @commands.command()
    async def repl(self, ctx):
        """Launches an interactive REPL session."""

        variables = {
            'ctx': ctx,
            'bot': ctx.bot,
            'message': ctx.message,
            'guild': ctx.guild,
            'channel': ctx.channel,
            'author': ctx.author,
            '_': None,
        }

        if ctx.channel.id in self.sessions:
            return await ctx.send('Already running a REPL session in this channel. Exit it with `quit`')

        self.sessions.add(ctx.channel.id)
        await ctx.send('Enter code to execute or evaluate. `exit()` or `quit` to exit.')

        def check(m):
            return m.author == ctx.author and \
                    m.channel == ctx.channel and \
                    m.content.startswith('`')

        while True:
            try:
                response = await ctx.bot.wait_for('message', check=check, timeout=10.0 * 60.0)
            except asyncio.TimeoutError:
                await ctx.send('Exiting REPL session.')
                self.sessions.remove(ctx.channel.id)
                break

            cleaned = self.cleanup_code(response.content)

            if cleaned in ('quit', 'exit', 'exit()'):
                await ctx.send('Exiting.')
                self.sessions.remove(ctx.channel.id)
                break

            executor = exec
            if cleaned.count('\n') == 0:
                try:
                    code = compile(cleaned, '<repl session>', 'eval')
                except SyntaxError:
                    pass
                else:
                    executor = eval

            if executor is exec:
                try:
                    code = compile(cleaned, '<repl session>', 'exec')
                except SyntaxError as e:
                    await ctx.send(self.get_syntax_error(e))
                    continue

            variables['message'] = response

            fmt = None
            stdout = io.StringIO()

            try:
                with redirect_stdout(stdout):
                    result = executor(code, variables)
                    if inspect.isawaitable(result):
                        result = await result
            except Exception as e:
                value = stdout.getvalue()
                fmt = f'```py\n{value}{traceback.format_exc(chain=False)}\n```'
            else:
                value = stdout.getvalue()
                if result is not None:
                    fmt = f'```py\n{value}{result}\n```'
                    variables['_'] = result
                elif value:
                    fmt = f'```py\n{value}\n```'

            try:
                if fmt is not None:
                    if len(fmt) > 2000:
                        await ctx.send('Content too big to be printed.')
                    else:
                        await ctx.send(fmt)
            except discord.Forbidden:
                pass
            except discord.HTTPException as e:
                await ctx.send(f'Unexpected error: `{e}`')


def setup(bot):
    bot.add_cog(Owner())

Owner-only commands used to maintain the bot. Logs out of the bot. Loads a module. Unloads a module. Reloads a module. Evaluates code. Launches an interactive REPL session.