Red-DiscordBot/redbot/core/dev_commands.py

import asyncio
import inspect
import io
import textwrap
import traceback
from contextlib import redirect_stdout

import discord
from discord.ext import commands
from . import checks
from .i18n import CogI18n
from .utils.chat_formatting import box, pagify
"""
Notice:

95% of the below code came from R.Danny which can be found here:

https://github.com/Rapptz/RoboDanny/blob/master/cogs/repl.py
"""

_ = CogI18n("Dev", __file__)


class Dev:
    """Various development focused utilities."""

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

    @staticmethod
    def cleanup_code(content):
        """Automatically removes code blocks from the code."""
        # remove ```py\n```
        if content.startswith('```') and content.endswith('```'):
            return '\n'.join(content.split('\n')[1:-1])

        # remove `foo`
        return content.strip('` \n')

    @staticmethod
    def get_syntax_error(e):
        """Format a syntax error to send to the user.

        Returns a string representation of the error formatted as a codeblock.
        """
        if e.text is None:
            return box('{0.__class__.__name__}: {0}'.format(e), lang="py")
        return box(
            '{0.text}{1:>{0.offset}}\n{2}: {0}'
            ''.format(e, '^', type(e).__name__),
            lang="py")

    @staticmethod
    def sanitize_output(ctx: commands.Context, input_: str) -> str:
        """Hides the bot's token from a string."""
        token = ctx.bot.http.token
        r = "[EXPUNGED]"
        result = input_.replace(token, r)
        result = result.replace(token.lower(), r)
        result = result.replace(token.upper(), r)
        return result

    @commands.command()
    @checks.is_owner()
    async def debug(self, ctx, *, code):
        """Evaluate a statement of python code.

        The bot will always respond with the return value of the code.
        If the return value of the code is a coroutine, it will be awaited,
        and the result of that will be the bot's response.

        Note: Only one statement may be evaluated. Using await, yield or
        similar restricted keywords will result in a syntax error. For multiple
        lines or asynchronous code, see [p]repl or [p]eval.

        Environment Variables:
            ctx      - command invokation context
            bot      - bot object
            channel  - the current channel object
            author   - command author's member object
            message  - the command's message object
            discord  - discord.py library
            commands - discord.py commands extension
            _        - The result of the last dev command.
        """
        env = {
            'bot': ctx.bot,
            'ctx': ctx,
            'channel': ctx.channel,
            'author': ctx.author,
            'guild': ctx.guild,
            'message': ctx.message,
            'discord': discord,
            'commands': commands,
            '_': self._last_result
        }

        code = self.cleanup_code(code)

        try:
            result = eval(code, env)
        except SyntaxError as e:
            await ctx.send(self.get_syntax_error(e))
            return
        except Exception as e:
            await ctx.send(
                box('{}: {!s}'.format(type(e).__name__, e), lang='py'))
            return

        if asyncio.iscoroutine(result):
            result = await result

        self._last_result = result

        result = self.sanitize_output(ctx, str(result))

        await ctx.send(box(result, lang="py"))

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

        This command wraps code into the body of an async function and then
        calls and awaits it. The bot will respond with anything printed to
        stdout, as well as the return value of the function.

        The code can be within a codeblock, inline code or neither, as long
        as they are not mixed and they are formatted correctly.

        Environment Variables:
            ctx      - command invokation context
            bot      - bot object
            channel  - the current channel object
            author   - command author's member object
            message  - the command's message object
            discord  - discord.py library
            commands - discord.py commands extension
            _        - The result of the last dev command.
        """
        env = {
            'bot': ctx.bot,
            'ctx': ctx,
            'channel': ctx.channel,
            'author': ctx.author,
            'guild': ctx.guild,
            'message': ctx.message,
            'discord': discord,
            'commands': commands,
            '_': self._last_result
        }

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

        to_compile = 'async def func():\n%s' % 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:
            value = stdout.getvalue()
            await ctx.send(
                box('\n{}{}'.format(value, traceback.format_exc()), lang="py"))
        else:
            value = stdout.getvalue()
            try:
                await ctx.message.add_reaction('\N{White Heavy Check Mark}')
            except:
                pass

            if ret is None:
                if value:
                    value = self.sanitize_output(ctx, str(value))
                    await ctx.send(box(value, lang="py"))
            else:
                self._last_result = ret
                ret = self.sanitize_output(ctx, str(ret))
                await ctx.send(box("{}{}".format(value, ret), lang="py"))

    @commands.command()
    @checks.is_owner()
    async def repl(self, ctx):
        """Open an interactive REPL.

        The REPL will only recognise code as messages which start with a
        backtick. This includes codeblocks, and as such multiple lines can be
        evaluated.

        You may not await any code in this REPL unless you define it inside an
        async function.
        """
        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:
            await ctx.send(_('Already running a REPL session in this channel. '
                             'Exit it with `quit`.'))
            return

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

        msg_check = lambda m: (m.author == ctx.author and
                               m.channel == ctx.channel and
                               m.content.startswith('`'))

        while True:
            response = await ctx.bot.wait_for("message", check=msg_check)

            cleaned = self.cleanup_code(response.content)

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

            executor = exec
            if cleaned.count('\n') == 0:
                # single statement, potentially 'eval'
                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

            stdout = io.StringIO()

            msg = None

            try:
                with redirect_stdout(stdout):
                    result = executor(code, variables)
                    if inspect.isawaitable(result):
                        result = await result
            except:
                value = stdout.getvalue()
                value = self.sanitize_output(ctx, value)
                msg = "{}{}".format(value, traceback.format_exc())
            else:
                value = stdout.getvalue()
                if result is not None:
                    msg = "{}{}".format(value, result)
                    variables['_'] = result
                elif value:
                    msg = "{}".format(value)

            try:
                for page in pagify(str(msg), shorten_by=12):
                    page = self.sanitize_output(ctx, page)
                    await ctx.send(box(page, "py"))
            except discord.Forbidden:
                pass
            except discord.HTTPException as e:
                await ctx.send(_('Unexpected error: `{}`').format(e))

    @commands.command()
    @checks.is_owner()
    async def mock(self, ctx, user: discord.Member, *, command):
        """Mock another user invoking a command.

        The prefix must not be entered.
        """
        # Since we have stateful objects now this might be pretty bad
        # Sorry Danny
        old_author = ctx.author
        old_content = ctx.message.content
        ctx.message.author = user
        ctx.message.content = ctx.prefix + command

        await ctx.bot.process_commands(ctx.message)

        ctx.message.author = old_author
        ctx.message.content = old_content

    @commands.command(name="mockmsg")
    @checks.is_owner()
    async def mock_msg(self, ctx, user: discord.Member, *, content: str):
        """Dispatch a message event as if it were sent by a different user.

        Only reads the raw content of the message. Attachments, embeds etc. are
        ignored.
        """
        old_author = ctx.author
        old_content = ctx.message.content
        ctx.message.author = user
        ctx.message.content = content

        ctx.bot.dispatch("message", ctx.message)

        # If we change the author and content back too quickly,
        # the bot won't process the mocked message in time.
        await asyncio.sleep(2)
        ctx.message.author = old_author
        ctx.message.content = old_content