Realigned to SirLancebot Structure

The code for rattlesnakes's internal eval was aligned
to Sir Lancebot's structure. It was mostly renaming
rattlesnake to bot and changing how some of the
imports were setup as.
It also included changing the __init__.py to match
the Sir Lancebot cog structure.

Additionally, the whitelist check has been significantly
simplified to only be a role check for the admin role.
The rattlesnake implementation had a more robust
`in_whitelist` decorator, so it may be worth investigating
adding that in if we see fit. For now, it's a simple
`with_role` decorator check.

The name of the cog file itself was changed
to include an underscore to sidestep what I think
was a namespace collision that would prevent
the setup function from properly running.

3 files changed, 405 insertions, 0 deletions

diff --git a/bot/exts/internal_eval/__init__.py b/bot/exts/internal_eval/__init__.py
index e69de29b..695fa74d 100644
--- a/bot/exts/internal_eval/__init__.py
+++ b/bot/exts/internal_eval/__init__.py
@@ -0,0 +1,10 @@
+from bot.bot import Bot
+
+
+def setup(bot: Bot) -> None:
+    """Set up the Internal Eval extension."""
+    # Import the Cog at runtime to prevent side effects like defining
+    # RedisCache instances too early.
+    from ._internal_eval import InternalEval
+
+    bot.add_cog(InternalEval(bot))
diff --git a/bot/exts/internal_eval/_helpers.py b/bot/exts/internal_eval/_helpers.py
new file mode 100644
index 00000000..5c602e4d
--- /dev/null
+++ b/bot/exts/internal_eval/_helpers.py
@@ -0,0 +1,243 @@
+import ast
+import collections
+import contextlib
+import functools
+import inspect
+import io
+import logging
+import sys
+import traceback
+import types
+import typing
+
+
+log = logging.getLogger("rattlesnake.exts.admin_tools.internal_eval")
+
+# A type alias to annotate the tuples returned from `sys.exc_info()`
+ExcInfo = typing.Tuple[typing.Type[Exception], Exception, types.TracebackType]
+Namespace = typing.Dict[str, typing.Any]
+
+# This will be used as an coroutine function wrapper for the code
+# to be evaluated. The wrapper contains one `pass` statement which
+# will be replaced with `ast` with the code that we want to have
+# evaluated.
+# The function redirects output and captures exceptions that were
+# raised in the code we evaluate. The latter is used to provide a
+# meaningful traceback to the end user.
+EVAL_WRAPPER = """
+async def _eval_wrapper_function():
+    try:
+        with contextlib.redirect_stdout(_eval_context.stdout):
+            pass
+        if '_value_last_expression' in locals():
+            if inspect.isawaitable(_value_last_expression):
+                _value_last_expression = await _value_last_expression
+            _eval_context._value_last_expression = _value_last_expression
+        else:
+            _eval_context._value_last_expression = None
+    except Exception:
+        _eval_context.exc_info = sys.exc_info()
+    finally:
+        _eval_context.locals = locals()
+_eval_context.function = _eval_wrapper_function
+"""
+
+
+def format_internal_eval_exception(exc_info: ExcInfo, code: str) -> str:
+    """Format an exception caught while evaluation code by inserting lines."""
+    exc_type, exc_value, tb = exc_info
+    stack_summary = traceback.StackSummary.extract(traceback.walk_tb(tb))
+    code = code.split("\n")
+
+    output = ["Traceback (most recent call last):"]
+    for frame in stack_summary:
+        if frame.filename == "<internal eval>":
+            line = code[frame.lineno - 1].lstrip()
+
+            if frame.name == "_eval_wrapper_function":
+                name = "<internal eval>"
+            else:
+                name = frame.name
+        else:
+            line = frame.line
+            name = frame.name
+
+        output.append(
+            f'  File "{frame.filename}", line {frame.lineno}, in {name}\n'
+            f"    {line}"
+        )
+
+    output.extend(traceback.format_exception_only(exc_type, exc_value))
+    return "\n".join(output)
+
+
+class EvalContext:
+    """
+    Represents the current `internal eval` context.
+    The context remembers names set during earlier runs of `internal eval`. To
+    clear the context, use the `?internal clear` command.
+    """
+
+    def __init__(self, context_vars: Namespace, local_vars: Namespace) -> None:
+        self._locals = dict(local_vars)
+        self.context_vars = dict(context_vars)
+
+        self.stdout = io.StringIO()
+        self._value_last_expression = None
+        self.exc_info = None
+        self.code = ""
+        self.function = None
+        self.eval_tree = None
+
+    @property
+    def dependencies(self) -> typing.Dict[str, typing.Any]:
+        """
+        Return a mapping of the dependencies for the wrapper function.
+        By using a property descriptor, the mapping can't be accidentally
+        mutated during evaluation. This ensures the dependencies are always
+        available.
+        """
+        return {
+            "print": functools.partial(print, file=self.stdout),
+            "contextlib": contextlib,
+            "inspect": inspect,
+            "sys": sys,
+            "_eval_context": self,
+            "_": self._value_last_expression,
+        }
+
+    @property
+    def locals(self) -> typing.Dict[str, typing.Any]:
+        """Return a mapping of names->values needed for evaluation."""
+        return {**collections.ChainMap(self.dependencies, self.context_vars, self._locals)}
+
+    @locals.setter
+    def locals(self, locals_: typing.Dict[str, typing.Any]) -> None:
+        """Update the contextual mapping of names to values."""
+        log.trace(f"Updating {self._locals} with {locals_}")
+        self._locals.update(locals_)
+
+    def prepare_eval(self, code: str) -> typing.Optional[str]:
+        """Prepare an evaluation by processing the code and setting up the context."""
+        self.code = code
+
+        if not self.code:
+            log.debug("No code was attached to the evaluation command")
+            return "[No code detected]"
+
+        try:
+            code_tree = ast.parse(code, filename="<internal eval>")
+        except SyntaxError:
+            log.debug("Got a SyntaxError while parsing the eval code")
+            return "".join(traceback.format_exception(*sys.exc_info(), limit=0))
+
+        log.trace("Parsing the AST to see if there's a trailing expression we need to capture")
+        code_tree = CaptureLastExpression(code_tree).capture()
+
+        log.trace("Wrapping the AST in the AST of the wrapper coroutine")
+        eval_tree = WrapEvalCodeTree(code_tree).wrap()
+
+        self.eval_tree = eval_tree
+        return None
+
+    async def run_eval(self) -> Namespace:
+        """Run the evaluation and return the updated locals."""
+        log.trace("Compiling the AST to bytecode using `exec` mode")
+        compiled_code = compile(self.eval_tree, filename="<internal eval>", mode="exec")
+
+        log.trace("Executing the compiled code with the desired namespace environment")
+        exec(compiled_code, self.locals)  # noqa: B102,S102
+
+        log.trace("Awaiting the created evaluation wrapper coroutine.")
+        await self.function()
+
+        log.trace("Returning the updated captured locals.")
+        return self._locals
+
+    def format_output(self) -> str:
+        """Format the output of the most recent evaluation."""
+        output = []
+
+        log.trace(f"Getting output from stdout `{id(self.stdout)}`")
+        stdout_text = self.stdout.getvalue()
+        if stdout_text:
+            log.trace("Appending output captured from stdout/print")
+            output.append(stdout_text)
+
+        if self._value_last_expression is not None:
+            log.trace("Appending the output of a captured trialing expression")
+            output.append(f"[Captured] {self._value_last_expression!r}")
+
+        if self.exc_info:
+            log.trace("Appending exception information")
+            output.append(format_internal_eval_exception(self.exc_info, self.code))
+
+        log.trace(f"Generated output: {output!r}")
+        return "\n".join(output) or "[No output]"
+
+
+class WrapEvalCodeTree(ast.NodeTransformer):
+    """Wraps the AST of eval code with the wrapper function."""
+
+    def __init__(self, eval_code_tree: ast.AST, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.eval_code_tree = eval_code_tree
+
+        # To avoid mutable aliasing, parse the WRAPPER_FUNC for each wrapping
+        self.wrapper = ast.parse(EVAL_WRAPPER, filename="<internal eval>")
+
+    def wrap(self) -> ast.AST:
+        """Wrap the tree of the code by the tree of the wrapper function."""
+        new_tree = self.visit(self.wrapper)
+        return ast.fix_missing_locations(new_tree)
+
+    def visit_Pass(self, node: ast.Pass) -> typing.List[ast.AST]:  # noqa: N802
+        """
+        Replace the `_ast.Pass` node in the wrapper function by the eval AST.
+        This method works on the assumption that there's a single `pass`
+        statement in the wrapper function.
+        """
+        return list(ast.iter_child_nodes(self.eval_code_tree))
+
+
+class CaptureLastExpression(ast.NodeTransformer):
+    """Captures the return value from a loose expression."""
+
+    def __init__(self, tree: ast.AST, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.tree = tree
+        self.last_node = list(ast.iter_child_nodes(tree))[-1]
+
+    def visit_Expr(self, node: ast.Expr) -> typing.Union[ast.Expr, ast.Assign]:  # noqa: N802
+        """
+        Replace the Expr node that is last child node of Module with an assignment.
+        We use an assignment to capture the value of the last node, if it's a loose
+        Expr node. Normally, the value of an Expr node is lost, meaning we don't get
+        the output of such a last "loose" expression. By assigning it a name, we can
+        retrieve it for our output.
+        """
+        if node is not self.last_node:
+            return node
+
+        log.trace("Found a trailing last expression in the evaluation code")
+
+        log.trace("Creating assignment statement with trailing expression as the right-hand side")
+        right_hand_side = list(ast.iter_child_nodes(node))[0]
+
+        assignment = ast.Assign(
+            targets=[ast.Name(id='_value_last_expression', ctx=ast.Store())],
+            value=right_hand_side,
+            lineno=node.lineno,
+            col_offset=0,
+        )
+        ast.fix_missing_locations(assignment)
+        return assignment
+
+    def capture(self) -> ast.AST:
+        """Capture the value of the last expression with an assignment."""
+        if not isinstance(self.last_node, ast.Expr):
+            # We only have to replace a node if the very last node is an Expr node
+            return self.tree
+
+        new_tree = self.visit(self.tree)
+        return ast.fix_missing_locations(new_tree)
diff --git a/bot/exts/internal_eval/_internal_eval.py b/bot/exts/internal_eval/_internal_eval.py
new file mode 100644
index 00000000..f7a0946b
--- /dev/null
+++ b/bot/exts/internal_eval/_internal_eval.py
@@ -0,0 +1,152 @@
+import logging
+import re
+import textwrap
+import typing
+
+import discord
+from discord.ext import commands
+
+from bot.bot import Bot
+from bot.constants import Roles
+from bot.utils.decorators import with_role
+from ._helpers import EvalContext
+
+__all__ = ["InternalEval"]
+
+log = logging.getLogger("rattlesnake.exts.admin_tools.internal_eval")
+
+CODEBLOCK_REGEX = re.compile(r"(^```(py(thon)?)?\n)|(```$)")
+
+
+class InternalEval(commands.Cog):
+    """Top secret code evaluation for admins and owners."""
+
+    def __init__(self, bot: Bot):
+        self.bot = bot
+        self.locals = {}
+
+    @staticmethod
+    def shorten_output(
+            output: str,
+            max_length: int = 1900,
+            placeholder: str = "\n[output truncated]"
+    ) -> str:
+        """
+        Shorten the `output` so it's shorter than `max_length`.
+        There are three tactics for this, tried in the following order:
+        - Shorten the output on a line-by-line basis
+        - Shorten the output on any whitespace character
+        - Shorten the output solely on character count
+        """
+        max_length = max_length - len(placeholder)
+
+        shortened_output = []
+        char_count = 0
+        for line in output.split("\n"):
+            if char_count + len(line) > max_length:
+                break
+            shortened_output.append(line)
+            char_count += len(line) + 1  # account for (possible) line ending
+
+        if shortened_output:
+            shortened_output.append(placeholder)
+            return "\n".join(shortened_output)
+
+        shortened_output = textwrap.shorten(output, width=max_length, placeholder=placeholder)
+
+        if shortened_output.strip() == placeholder.strip():
+            # `textwrap` was unable to find whitespace to shorten on, so it has
+            # reduced the output to just the placeholder. Let's shorten based on
+            # characters instead.
+            shortened_output = output[:max_length] + placeholder
+
+        return shortened_output
+
+    async def _upload_output(self, output: str) -> typing.Optional[str]:
+        """Upload `internal eval` output to our pastebin and return the url."""
+        try:
+            async with self.bot.http_session.post(
+                "https://paste.pythondiscord.com/documents", data=output, raise_for_status=True
+            ) as resp:
+                data = await resp.json()
+
+            if "key" in data:
+                return f"https://paste.pythondiscord.com/{data['key']}"
+        except Exception:
+            # 400 (Bad Request) means there are too many characters
+            log.exception("Failed to upload `internal eval` output to paste service!")
+
+    async def _send_output(self, ctx: commands.Context, output: str) -> None:
+        """Send the `internal eval` output to the command invocation context."""
+        upload_message = ""
+        if len(output) >= 1980:
+            # The output is too long, let's truncate it for in-channel output and
+            # upload the complete output to the paste service.
+            url = await self._upload_output(output)
+
+            if url:
+                upload_message = f"\nFull output here: {url}"
+            else:
+                upload_message = "\n:warning: Failed to upload full output!"
+
+            output = self.shorten_output(output)
+
+        await ctx.send(f"```py\n{output}```{upload_message}")
+
+    async def _eval(self, ctx: commands.Context, code: str) -> None:
+        """Evaluate the `code` in the current evaluation context."""
+        if code.startswith("exit"):
+            self.locals = {}
+            await ctx.send("The evaluation context was reset.")
+            return
+
+        context_vars = {
+            "message": ctx.message,
+            "author": ctx.message.author,
+            "channel": ctx.channel,
+            "guild": ctx.guild,
+            "ctx": ctx,
+            "self": self,
+            "bot": self.bot,
+            "discord": discord,
+        }
+
+        eval_context = EvalContext(context_vars, self.locals)
+
+        log.trace("Preparing the evaluation by parsing the AST of the code")
+        error = eval_context.prepare_eval(code)
+
+        if error:
+            log.trace("The code can't be evaluated due to an error")
+            await ctx.send(f"```py\n{error}\n```")
+            return
+
+        log.trace("Evaluate the AST we've generated for the evaluation")
+        new_locals = await eval_context.run_eval()
+
+        log.trace("Updating locals with those set during evaluation")
+        self.locals.update(new_locals)
+
+        log.trace("Sending the formatted output back to the context")
+        await self._send_output(ctx, eval_context.format_output())
+
+    @commands.group(name='internal', aliases=('int',))
+    @with_role(Roles.admin)
+    async def internal_group(self, ctx: commands.Context) -> None:
+        """Internal commands. Top secret!"""
+        if not ctx.invoked_subcommand:
+            await ctx.send_help(ctx.command)
+
+    @internal_group.command(name='eval', aliases=('e',))
+    @with_role(Roles.admin)
+    async def eval(self, ctx: commands.Context, *, code: str) -> None:
+        """Run eval in a REPL-like format."""
+        code = CODEBLOCK_REGEX.sub("", code.strip())
+        await self._eval(ctx, code)
+
+    @internal_group.command(name='reset', aliases=("clear", "exit", "r", "c"))
+    @with_role(Roles.admin)
+    async def reset(self, ctx: commands.Context) -> None:
+        """Run eval in a REPL-like format."""
+        self.locals = {}
+        await ctx.send("The evaluation context was reset.")