rewrite: Restructure to include GUI.
This commit is contained in:
389
src/modules/sysadmin/exec_cog.py
Normal file
389
src/modules/sysadmin/exec_cog.py
Normal file
@@ -0,0 +1,389 @@
|
||||
import io
|
||||
import ast
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
import traceback
|
||||
import builtins
|
||||
import inspect
|
||||
import logging
|
||||
from io import StringIO
|
||||
|
||||
from typing import Callable, Any, Optional
|
||||
|
||||
from enum import Enum
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
from discord.ui import TextInput, View
|
||||
from discord.ui.button import button
|
||||
import discord.app_commands as appcmd
|
||||
|
||||
from meta.logger import logging_context
|
||||
from meta.app import shard_talk
|
||||
from meta import conf
|
||||
from meta.context import context, ctx_bot
|
||||
from meta.LionContext import LionContext
|
||||
from meta.LionCog import LionCog
|
||||
from meta.LionBot import LionBot
|
||||
|
||||
from utils.ui import FastModal, input
|
||||
|
||||
from babel.translator import LocalBabel
|
||||
|
||||
from wards import sys_admin
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_, _n, _p, _np = LocalBabel('exec').methods
|
||||
|
||||
|
||||
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=True) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.ctx: LionContext = 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 interaction_check(self, interaction: discord.Interaction):
|
||||
"""Only allow the original author to use this View"""
|
||||
if interaction.user.id != self.ctx.author.id:
|
||||
await interaction.response.send_message(
|
||||
("You cannot use this interface!"),
|
||||
ephemeral=True
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
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())
|
||||
await self.wait()
|
||||
else:
|
||||
# Complain
|
||||
# TODO: error_reply
|
||||
await self.ctx.reply("Pls give code.")
|
||||
else:
|
||||
await self.interaction.response.defer(thinking=True, ephemeral=self.ephemeral)
|
||||
await self.compile()
|
||||
await self.wait()
|
||||
|
||||
@button(label="Recompile")
|
||||
async def recompile_button(self, interaction, butt):
|
||||
# Interaction response with modal
|
||||
await interaction.response.send_modal(self.get_modal())
|
||||
|
||||
@button(label="Show Source")
|
||||
async def source_button(self, interaction, butt):
|
||||
if len(self.code) > 1900:
|
||||
# Send as file
|
||||
with StringIO(self.code) as fp:
|
||||
fp.seek(0)
|
||||
file = discord.File(fp, filename="source.py")
|
||||
await interaction.response.send_message(file=file, ephemeral=True)
|
||||
else:
|
||||
# Send as message
|
||||
await interaction.response.send_message(
|
||||
content=f"```py\n{self.code}```",
|
||||
ephemeral=True
|
||||
)
|
||||
|
||||
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.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
|
||||
if len(output) > 1900:
|
||||
# Send as file
|
||||
with StringIO(output) as fp:
|
||||
fp.seek(0)
|
||||
args = {
|
||||
'content': None,
|
||||
'attachments': [discord.File(fp, filename="output.md")]
|
||||
}
|
||||
else:
|
||||
args = {
|
||||
'content': f"```md\n{output}```",
|
||||
'attachments': []
|
||||
}
|
||||
|
||||
if self._msg is None:
|
||||
if self.interaction is not None:
|
||||
msg = await self.interaction.edit_original_response(**args, view=self)
|
||||
else:
|
||||
# Send new message
|
||||
if args['content'] is None:
|
||||
args['file'] = args.pop('attachments')[0]
|
||||
msg = await self.ctx.reply(**args, 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(**args, view=self)
|
||||
else:
|
||||
# Edit message
|
||||
await self._msg.edit(**args)
|
||||
|
||||
|
||||
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(to_eval: str, style='exec'):
|
||||
with logging_context(action="Code Exec"):
|
||||
newline = '\n' * ('\n' in to_eval)
|
||||
logger.info(
|
||||
f"Exec code with {style}: {newline}{to_eval}"
|
||||
)
|
||||
|
||||
output = io.StringIO()
|
||||
_print = mk_print(output)
|
||||
|
||||
scope: dict[str, Any] = dict(sys.modules)
|
||||
scope['__builtins__'] = builtins
|
||||
scope.update(builtins.__dict__)
|
||||
scope['ctx'] = ctx = context.get()
|
||||
scope['bot'] = ctx_bot.get()
|
||||
scope['print'] = _print # type: ignore
|
||||
|
||||
try:
|
||||
if ctx and ctx.message:
|
||||
source_str = f"<msg: {ctx.message.id}>"
|
||||
elif ctx and ctx.interaction:
|
||||
source_str = f"<iid: {ctx.interaction.id}>"
|
||||
else:
|
||||
source_str = "Unknown async"
|
||||
|
||||
code = compile(
|
||||
to_eval,
|
||||
source_str,
|
||||
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(f"{type(exc).__name__}: {exc}")
|
||||
|
||||
result = output.getvalue().strip()
|
||||
newline = '\n' * ('\n' in result)
|
||||
logger.info(
|
||||
f"Exec complete, output: {newline}{result}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class Exec(LionCog):
|
||||
guild_ids = conf.bot.getintlist('admin_guilds')
|
||||
|
||||
def __init__(self, bot: LionBot):
|
||||
self.bot = bot
|
||||
self.t = bot.translator.t
|
||||
|
||||
self.talk_async = shard_talk.register_route('exec')(_async)
|
||||
|
||||
async def cog_check(self, ctx: LionContext) -> bool: # type: ignore
|
||||
return await sys_admin(ctx)
|
||||
|
||||
@commands.hybrid_command(
|
||||
name=_('async'),
|
||||
description=_("Execute arbitrary code with Exec")
|
||||
)
|
||||
@appcmd.describe(
|
||||
string="Code to execute."
|
||||
)
|
||||
@appcmd.guilds(*guild_ids)
|
||||
async def async_cmd(self, ctx: LionContext, *, string: Optional[str] = None):
|
||||
await ExecUI(ctx, string, ExecStyle.EXEC).run()
|
||||
|
||||
@commands.hybrid_command(
|
||||
name=_p('command', 'eval'),
|
||||
description=_p('command:eval', 'Execute arbitrary code with Eval')
|
||||
)
|
||||
@appcmd.describe(
|
||||
string=_p('command:eval|param:string', "Code to evaluate.")
|
||||
)
|
||||
@appcmd.guilds(*guild_ids)
|
||||
async def eval_cmd(self, ctx: LionContext, *, string: str):
|
||||
await ExecUI(ctx, string, ExecStyle.EVAL).run()
|
||||
|
||||
@commands.hybrid_command(
|
||||
name=_p('command', 'asyncall'),
|
||||
description=_p('command:asyncall|desc', "Execute arbitrary code on all shards.")
|
||||
)
|
||||
@appcmd.describe(
|
||||
string=_p("command:asyncall|param:string", "Cross-shard code to execute. Cannot reference ctx!"),
|
||||
target=_p("command:asyncall|param:target", "Target shard app name, see autocomplete for options.")
|
||||
)
|
||||
@appcmd.guilds(*guild_ids)
|
||||
async def asyncall_cmd(self, ctx: LionContext, string: Optional[str] = None, target: Optional[str] = None):
|
||||
if string is None and ctx.interaction:
|
||||
try:
|
||||
ctx.interaction, string = await input(
|
||||
ctx.interaction, "Cross-shard execute", "Code to execute?",
|
||||
style=discord.TextStyle.long
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return
|
||||
if ctx.interaction:
|
||||
await ctx.interaction.response.defer(thinking=True, ephemeral=True)
|
||||
if target is not None:
|
||||
if target not in shard_talk.peers:
|
||||
embed = discord.Embed(description=f"Unknown peer {target}", colour=discord.Colour.red())
|
||||
if ctx.interaction:
|
||||
await ctx.interaction.edit_original_response(embed=embed)
|
||||
else:
|
||||
await ctx.reply(embed=embed)
|
||||
return
|
||||
else:
|
||||
result = await self.talk_async(string).send(target)
|
||||
results = {target: result}
|
||||
else:
|
||||
results = await self.talk_async(string).broadcast(except_self=False)
|
||||
|
||||
blocks = [f"# {appid}\n{result}" for appid, result in results.items()]
|
||||
output = "\n\n".join(blocks)
|
||||
if len(output) > 1900:
|
||||
# Send as file
|
||||
with StringIO(output) as fp:
|
||||
fp.seek(0)
|
||||
file = discord.File(fp, filename="output.md") # type: ignore
|
||||
await ctx.reply(file=file)
|
||||
else:
|
||||
# Send as message
|
||||
await ctx.reply(f"```md\n{output}```", ephemeral=True)
|
||||
|
||||
@asyncall_cmd.autocomplete('target')
|
||||
async def asyncall_target_acmpl(self, interaction: discord.Interaction, partial: str):
|
||||
appids = set(shard_talk.peers.keys())
|
||||
results = [
|
||||
appcmd.Choice(name=appid, value=appid)
|
||||
for appid in appids
|
||||
if partial.lower() in appid.lower()
|
||||
]
|
||||
if not results:
|
||||
results = [
|
||||
appcmd.Choice(name=f"No peers found matching {partial}", value="None")
|
||||
]
|
||||
return results
|
||||
|
||||
@commands.hybrid_command(
|
||||
name=_('reload'),
|
||||
description=_("Reload a given LionBot extension. Launches an ExecUI.")
|
||||
)
|
||||
@appcmd.describe(
|
||||
extension=_("Name of the extesion to reload. See autocomplete for options.")
|
||||
)
|
||||
@appcmd.guilds(*guild_ids)
|
||||
async def reload_cmd(self, ctx: LionContext, extension: str):
|
||||
"""
|
||||
This is essentially just a friendly wrapper to reload an extension.
|
||||
It is equivalent to running "await bot.reload_extension(extension)" in eval,
|
||||
with a slightly nicer interface through the autocomplete and error handling.
|
||||
"""
|
||||
if extension not in self.bot.extensions:
|
||||
embed = discord.Embed(description=f"Unknown extension {extension}", colour=discord.Colour.red())
|
||||
await ctx.reply(embed=embed)
|
||||
else:
|
||||
# Uses an ExecUI to simplify error handling and re-execution
|
||||
string = f"await bot.reload_extension('{extension}')"
|
||||
await ExecUI(ctx, string, ExecStyle.EVAL).run()
|
||||
|
||||
@reload_cmd.autocomplete('extension')
|
||||
async def reload_extension_acmpl(self, interaction: discord.Interaction, partial: str):
|
||||
keys = set(self.bot.extensions.keys())
|
||||
results = [
|
||||
appcmd.Choice(name=key, value=key)
|
||||
for key in keys
|
||||
if partial.lower() in key.lower()
|
||||
]
|
||||
if not results:
|
||||
results = [
|
||||
appcmd.Choice(name=f"No extensions found matching {partial}", value="None")
|
||||
]
|
||||
return results
|
||||
|
||||
@commands.hybrid_command(
|
||||
name=_('shutdown'),
|
||||
description=_("Shutdown (or restart) the client.")
|
||||
)
|
||||
@appcmd.guilds(*guild_ids)
|
||||
async def shutdown_cmd(self, ctx: LionContext):
|
||||
"""
|
||||
Shutdown the client.
|
||||
Maybe do something friendly here?
|
||||
"""
|
||||
logger.info("Shutting down on admin request.")
|
||||
await ctx.reply(
|
||||
embed=discord.Embed(
|
||||
description=f"Understood {ctx.author.mention}, cleaning up and shutting down!",
|
||||
colour=discord.Colour.orange()
|
||||
)
|
||||
)
|
||||
await self.bot.close()
|
||||
Reference in New Issue
Block a user