rewrite: Start SysAdmin module.
This commit is contained in:
226
bot/modules/bot_admin/exec_cog.py
Normal file
226
bot/modules/bot_admin/exec_cog.py
Normal file
@@ -0,0 +1,226 @@
|
||||
import io
|
||||
import ast
|
||||
import sys
|
||||
import types
|
||||
import traceback
|
||||
import builtins
|
||||
import inspect
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from typing import Callable, Any, Coroutine, List, Optional
|
||||
|
||||
from enum import Enum
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
from discord.ui import Modal, TextInput, View
|
||||
from discord.ui.button import button
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FastModal(Modal):
|
||||
def __init__(self, *items, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
for item in items:
|
||||
self.add_item(item)
|
||||
self._result: asyncio.Future[discord.Interaction] = asyncio.get_event_loop().create_future()
|
||||
self._waiters: List[Coroutine[discord.Interaction]] = []
|
||||
|
||||
async def wait_for(self, check=None, timeout=None):
|
||||
# Wait for _result or timeout
|
||||
# If we timeout, or the view times out, raise TimeoutError
|
||||
# Otherwise, return the Interaction
|
||||
# This allows multiple listeners and callbacks to wait on
|
||||
# TODO: Wait on the timeout as well
|
||||
while True:
|
||||
result = await asyncio.wait_for(asyncio.shield(self._result), timeout=timeout)
|
||||
if check is not None:
|
||||
if not check(result):
|
||||
continue
|
||||
return result
|
||||
|
||||
def submit_callback(self, timeout=None, check=None, once=False, pass_args=(), pass_kwargs={}):
|
||||
def wrapper(coro):
|
||||
async def wrapped_callback(interaction):
|
||||
if check is not None:
|
||||
if not check(interaction):
|
||||
return
|
||||
try:
|
||||
await coro(interaction, *pass_args, **pass_kwargs)
|
||||
except Exception:
|
||||
# TODO: Log exception
|
||||
...
|
||||
if once:
|
||||
self._waiters.remove(wrapped_callback)
|
||||
self._waiters.append(wrapped_callback)
|
||||
return wrapper
|
||||
|
||||
async def on_submit(self, interaction):
|
||||
old_result = self._result
|
||||
self._result = asyncio.get_event_loop().create_future()
|
||||
old_result.set_result(interaction)
|
||||
|
||||
for waiter in self._waiters:
|
||||
asyncio.create_task(waiter(interaction))
|
||||
|
||||
async def on_error(self, interaction, error):
|
||||
# This should never happen, since on_submit has its own error handling
|
||||
# TODO: Logging
|
||||
...
|
||||
|
||||
|
||||
class ExecModal(FastModal, title="Execute"):
|
||||
code: TextInput = TextInput(
|
||||
label="Code to execute",
|
||||
style=discord.TextStyle.long,
|
||||
required=True
|
||||
)
|
||||
|
||||
|
||||
class ExecStyle(Enum):
|
||||
EXEC = 'exec'
|
||||
EVAL = 'eval'
|
||||
|
||||
|
||||
class ExecUI(View):
|
||||
def __init__(self, ctx, code=None, style=ExecStyle.EXEC, ephemeral=False):
|
||||
super().__init__()
|
||||
|
||||
self.ctx: commands.Context = ctx
|
||||
self.interaction: Optional[discord.Interaction] = ctx.interaction
|
||||
self.code: Optional[str] = code
|
||||
self.style: ExecStyle = style
|
||||
self.ephemeral: bool = ephemeral
|
||||
|
||||
self._modal: Optional[ExecModal] = None
|
||||
self._msg: Optional[discord.Message] = None
|
||||
|
||||
async def run(self):
|
||||
if self.code is None:
|
||||
if (interaction := self.interaction) is not None:
|
||||
self.interaction = None
|
||||
await interaction.response.send_modal(self.get_modal())
|
||||
else:
|
||||
# Complain
|
||||
# TODO: error_reply
|
||||
await self.ctx.reply("Pls give code.")
|
||||
else:
|
||||
await self.interaction.response.defer(thinking=True)
|
||||
await self.compile()
|
||||
|
||||
@button(label="Recompile")
|
||||
async def recompile_button(self, interaction, butt):
|
||||
# Interaction response with modal
|
||||
await interaction.response.send_modal(self.get_modal())
|
||||
|
||||
def create_modal(self) -> ExecModal:
|
||||
modal = ExecModal()
|
||||
|
||||
@modal.submit_callback()
|
||||
async def exec_submit(interaction: discord.Interaction):
|
||||
if self.interaction is None:
|
||||
self.interaction = interaction
|
||||
await interaction.response.defer(thinking=True)
|
||||
else:
|
||||
await interaction.response.defer()
|
||||
|
||||
# Set code
|
||||
self.code = modal.code.value
|
||||
|
||||
# Call compile
|
||||
await self.compile()
|
||||
|
||||
return modal
|
||||
|
||||
def get_modal(self):
|
||||
if self._modal is None:
|
||||
# Create modal
|
||||
self._modal = self.create_modal()
|
||||
|
||||
self._modal.code.default = self.code
|
||||
return self._modal
|
||||
|
||||
async def compile(self):
|
||||
# Call _async
|
||||
result = await _async(self.ctx, self.code, style=self.style.value)
|
||||
|
||||
# Display output
|
||||
await self.show_output(result)
|
||||
|
||||
async def show_output(self, output):
|
||||
# Format output
|
||||
# If output message exists and not ephemeral, edit
|
||||
# Otherwise, send message, add buttons
|
||||
# TODO: File output
|
||||
# TODO: Check this handles ephemerals properly
|
||||
formatted = "```py\n{}```".format(output)
|
||||
if self._msg is None:
|
||||
if self.interaction is not None:
|
||||
msg = await self.interaction.edit_original_response(content=formatted, view=self)
|
||||
else:
|
||||
# Send new message
|
||||
msg = await self.ctx.reply(formatted, ephemeral=self.ephemeral, view=self)
|
||||
|
||||
if not self.ephemeral:
|
||||
self._msg = msg
|
||||
else:
|
||||
if self.interaction is not None:
|
||||
await self.interaction.edit_original_response(content=formatted, view=self)
|
||||
else:
|
||||
# Edit message
|
||||
await self._msg.edit(formatted)
|
||||
|
||||
|
||||
def mk_print(fp: io.StringIO) -> Callable[..., None]:
|
||||
def _print(*args, file: Any = fp, **kwargs):
|
||||
return print(*args, file=file, **kwargs)
|
||||
return _print
|
||||
|
||||
|
||||
async def _async(ctx: commands.Context, to_eval, style='exec'):
|
||||
output = io.StringIO()
|
||||
_print = mk_print(output)
|
||||
|
||||
scope = dict(sys.modules)
|
||||
scope['__builtins__'] = builtins
|
||||
scope.update(builtins.__dict__)
|
||||
scope['ctx'] = ctx
|
||||
scope['bot'] = ctx.bot
|
||||
scope['print'] = _print # type: ignore
|
||||
|
||||
try:
|
||||
code = compile(to_eval, f"<msg: {ctx.message.id}>", style, ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
|
||||
func = types.FunctionType(code, scope)
|
||||
|
||||
ret = func()
|
||||
if inspect.iscoroutine(ret):
|
||||
ret = await ret
|
||||
if ret is not None:
|
||||
_print(repr(ret))
|
||||
except Exception:
|
||||
_, exc, tb = sys.exc_info()
|
||||
_print("".join(traceback.format_tb(tb)))
|
||||
_print(repr(exc))
|
||||
|
||||
result = output.getvalue()
|
||||
logger.info(
|
||||
f"Exec complete, output:\n{result}",
|
||||
extra={'action': "Code Exec"}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class Exec(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
@commands.hybrid_command(name='async')
|
||||
async def async_cmd(self, ctx, *, string: str = None):
|
||||
await ExecUI(ctx, string, ExecStyle.EXEC).run()
|
||||
|
||||
@commands.hybrid_command(name='eval')
|
||||
async def eval_cmd(self, ctx, *, string: str):
|
||||
await ExecUI(ctx, string, ExecStyle.EVAL).run()
|
||||
Reference in New Issue
Block a user