discordbot/cogs/eval_cog.py

132 lines
4.7 KiB
Python

import discord
from discord.ext import commands
import io
import textwrap
import traceback
import contextlib
from discord import app_commands
from discord.ui import Modal, TextInput
class EvalModal(Modal, title="Evaluate Python Code"):
code_input = TextInput(
label="Code",
style=discord.TextStyle.paragraph,
placeholder="Enter your Python code here...",
required=True,
max_length=1900 # Discord modal input limit is 2000 characters
)
async def on_submit(self, interaction: discord.Interaction):
await interaction.response.defer(ephemeral=True) # Defer the interaction to prevent timeout
# Access the bot and cleanup_code method from the cog instance
cog = interaction.client.get_cog("EvalCog")
if not cog:
await interaction.followup.send("EvalCog not found. Bot might be restarting or not loaded correctly.", ephemeral=True)
return
env = {
'bot': cog.bot,
'ctx': interaction, # Use interaction as ctx for slash commands
'channel': interaction.channel,
'author': interaction.user,
'guild': interaction.guild,
'message': None, # No message object for slash commands
'discord': discord,
'commands': commands
}
env.update(globals())
body = cog.cleanup_code(self.code_input.value)
stdout = io.StringIO()
to_compile = f'async def func():\n _ctx = ctx\n{textwrap.indent(body, " ")}'
try:
exec(to_compile, env)
except Exception as e:
return await interaction.followup.send(f'```py\n{e.__class__.__name__}: {e}\n```', ephemeral=True)
func = env['func']
try:
with contextlib.redirect_stdout(stdout):
ret = await func()
except Exception as e:
value = stdout.getvalue()
await interaction.followup.send(f'```py\n{value}{traceback.format_exc()}\n```', ephemeral=True)
else:
value = stdout.getvalue()
if ret is None:
if value:
await interaction.followup.send(f'```py\n{value}\n```', ephemeral=True)
else:
await interaction.followup.send(f'```py\n{value}{ret}\n```', ephemeral=True)
class EvalCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
def cleanup_code(self, 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])
return content.strip('` \n')
@commands.command(name='evalpy', hidden=True)
@commands.is_owner()
async def _eval(self, ctx, *, body: str):
"""Evaluates a code snippet."""
env = {
'bot': self.bot,
'ctx': ctx,
'channel': ctx.channel,
'author': ctx.author,
'guild': ctx.guild,
'message': ctx.message,
'discord': discord,
'commands': commands
}
env.update(globals())
body = self.cleanup_code(body)
stdout = io.StringIO()
to_compile = f'async def func():\n _ctx = ctx\n{textwrap.indent(body, " ")}'
try:
exec(to_compile, env)
except Exception as e:
return await ctx.send(f'```py\n{e.__class__.__name__}: {e}\n```')
func = env['func']
try:
with contextlib.redirect_stdout(stdout):
ret = await func()
except Exception as e:
value = stdout.getvalue()
await ctx.send(f'```py\n{value}{traceback.format_exc()}\n```')
else:
value = stdout.getvalue()
if ret is None:
if value:
await ctx.send(f'```py\n{value}\n```')
else:
await ctx.send(f'```py\n{value}{ret}\n```')
async def is_owner_check(interaction: discord.Interaction) -> bool:
"""Checks if the interacting user is the bot owner."""
return interaction.user.id == interaction.client.owner_id
@app_commands.command(name="eval", description="Evaluate Python code using a modal form.")
@app_commands.check(is_owner_check)
async def eval_slash(self, interaction: discord.Interaction):
"""Opens a modal to evaluate Python code."""
await interaction.response.send_modal(EvalModal())
async def setup(bot: commands.Bot):
await bot.add_cog(EvalCog(bot))
# After adding the cog, sync the commands
# If you have guild-specific commands, you might need to sync to specific guilds:
# await bot.tree.sync(guild=discord.Object(id=YOUR_GUILD_ID))