diff --git a/nomen/main.py b/nomen/main.py index 6962f10..35eb3de 100644 --- a/nomen/main.py +++ b/nomen/main.py @@ -1,6 +1,11 @@ +import io import logging import os +import pprint import sys +import textwrap +import traceback +from contextlib import redirect_stdout from disnake import Guild, Intents, Message from disnake.ext import commands @@ -10,6 +15,7 @@ from dotenv import find_dotenv, load_dotenv from .db import run_db_migrations, setup_db from .notifications import Notifications from .settings import Settings +from .utils import cleanup_code # Logger setup logger_disnake = logging.getLogger("disnake") @@ -122,12 +128,73 @@ async def on_command_error(ctx, error): await ctx.send(error) -@bot.command() +@bot.command(hidden=True) @commands.is_owner() async def echo(ctx, arg): await ctx.send(arg) +@bot.command(hidden=True) +@commands.is_owner() +async def sql(ctx, *, query): + params = { + "guild_id": ctx.guild.id, + } + + query = cleanup_code(query) + + cur = await ctx.bot.db.execute(query) + + result = await cur.fetchall() + + # TODO: Finish SQL evaluation command + + +@bot.command(hidden=True, name="eval") +@commands.is_owner() +async def admin_eval(self, ctx, *, body: str): + env = { + "bot": self.bot, + "ctx": ctx, + "channel": ctx.message.channel, + "author": ctx.message.author, + "guild": ctx.message.guild, + "message": ctx.message, + } + + env.update(globals()) + + body = cleanup_code(body) + stdout = io.StringIO() + + to_compile = "async def func():\n{}".format(textwrap.indent(body, " ")) + + try: + exec(to_compile, env) + except Exception as e: + return await ctx.send("```py\n{}: {}\n```".format(e.__class__.__name__, e)) + + func = env["func"] + try: + with redirect_stdout(stdout): + ret = await func() + except Exception: + value = stdout.getvalue() + await ctx.send("```py\n{}{}\n```".format(value, traceback.format_exc())) + else: + value = stdout.getvalue() + try: + await ctx.message.add_reaction("\u2705") + except: + pass + + if ret is None: + if value: + await ctx.send("```py\n{}\n```".format(value)) + else: + await ctx.send("```py\n{}{}\n```".format(value, ret)) + + @bot.command() async def ping(ctx): await ctx.send("Pong") diff --git a/nomen/utils.py b/nomen/utils.py index 8a22f3c..7253e0c 100644 --- a/nomen/utils.py +++ b/nomen/utils.py @@ -174,5 +174,13 @@ async def confirm(ctx, message, delete_msgs=False, response_check=get_positivity pass return reply_bool +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") # ===== End code borrowed from Avrae =====