454 lines
15 KiB
Python
454 lines
15 KiB
Python
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.ext.commands.errors import CheckFailure
|
|
from discord.ui import TextInput, View
|
|
from discord.ui.button import button
|
|
import discord.app_commands as appcmd
|
|
|
|
from meta.logger import logging_context, log_wrap
|
|
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):
|
|
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
|
|
|
|
|
|
def mk_status_printer(bot, printer):
|
|
async def _status(details=False):
|
|
if details:
|
|
status = await bot.system_monitor.get_overview()
|
|
else:
|
|
status = await bot.system_monitor.get_summary()
|
|
printer(status)
|
|
return status
|
|
return _status
|
|
|
|
|
|
@log_wrap(action="Code Exec")
|
|
async def _async(to_eval: str, style='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
|
|
scope['print_status'] = mk_status_printer(scope['bot'], _print)
|
|
|
|
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
|
|
passed = await sys_admin(ctx.bot, ctx.author.id)
|
|
if passed:
|
|
return True
|
|
else:
|
|
raise CheckFailure(
|
|
ctx.bot.translator.t(_p(
|
|
'ward:sys_admin|failed',
|
|
"You must be a bot owner to do this!"
|
|
))
|
|
)
|
|
|
|
@commands.hybrid_command(
|
|
name=_('async'),
|
|
description=_("Execute arbitrary code with Exec")
|
|
)
|
|
@appcmd.describe(
|
|
string="Code to execute.",
|
|
target="Cross-shard peer to async on."
|
|
)
|
|
async def async_cmd(self, ctx: LionContext,
|
|
string: Optional[str] = None,
|
|
target: Optional[str] = None,
|
|
):
|
|
if target is not None:
|
|
if string is None:
|
|
try:
|
|
ctx.interaction, string = await input(
|
|
ctx.interaction, "Cross-shard async", "Code to execute?",
|
|
style=discord.TextStyle.long
|
|
)
|
|
except asyncio.TimeoutError:
|
|
return
|
|
await ctx.interaction.response.defer(thinking=True)
|
|
if target not in shard_talk.peers:
|
|
# Invalid target
|
|
embed = discord.Embed(
|
|
description="Unknown peer {target}",
|
|
colour=discord.Colour.brand_red(),
|
|
)
|
|
await ctx.interaction.edit_original_response(embed=embed)
|
|
else:
|
|
# Send to given target
|
|
result = await self.talk_async(string).send(target)
|
|
if len(result) > 1900:
|
|
# Send as file
|
|
with StringIO(result) as fp:
|
|
fp.seek(0)
|
|
file = discord.File(fp, filename=f"output-{target}.md")
|
|
await ctx.reply(file=file)
|
|
elif result:
|
|
await ctx.reply(f"```md\n{result}```")
|
|
else:
|
|
await ctx.reply("Command completed, and had no output.")
|
|
else:
|
|
await ExecUI(ctx, string, ExecStyle.EXEC, ephemeral=False).run()
|
|
|
|
async def _peer_acmpl(self, interaction: discord.Interaction, partial: str):
|
|
"""
|
|
Autocomplete utility for peer targets parameters.
|
|
"""
|
|
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=partial)
|
|
]
|
|
return results[:25]
|
|
|
|
async_cmd.autocomplete('target')(_peer_acmpl)
|
|
|
|
@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)
|
|
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}```")
|
|
|
|
asyncall_cmd.autocomplete('target')(_peer_acmpl)
|
|
|
|
@commands.hybrid_command(
|
|
name=_('reload'),
|
|
description=_("Reload a given LionBot extension. Launches an ExecUI.")
|
|
)
|
|
@appcmd.describe(
|
|
extension=_("Name of the extension to reload. See autocomplete for options."),
|
|
force=_("Whether to force an extension reload even if it doesn't exist.")
|
|
)
|
|
@appcmd.guilds(*guild_ids)
|
|
async def reload_cmd(self, ctx: LionContext, extension: str, force: Optional[bool] = False):
|
|
"""
|
|
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.
|
|
"""
|
|
exists = (extension in self.bot.extensions)
|
|
if not (force or exists):
|
|
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
|
|
if exists:
|
|
string = f"await bot.reload_extension('{extension}')"
|
|
else:
|
|
string = f"await bot.load_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[:25]
|
|
|
|
@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()
|