From f1987d2788111e33e330c32d423eb363b42ffa68 Mon Sep 17 00:00:00 2001 From: Vegard Berg Date: Sat, 6 Aug 2022 02:01:29 +0200 Subject: [PATCH] Added bot messages. Updated polls and self-roles. Added bot messages, intended for moderators to be able to post as the bot. These messages are also editable through a context menu. Intended for rules/info messages etc. Polls: Added some error handling, as the bot only supports Unicode emojis. Self-roles: Edited the name of the context menu for updating a role-group message. --- Heimdallr.py | 1 + commands/bot_messages.py | 190 +++++++++++++++++++++++++++++++++++++++ commands/polls.py | 59 +++++++----- commands/self_roles.py | 2 +- 4 files changed, 229 insertions(+), 23 deletions(-) create mode 100644 commands/bot_messages.py diff --git a/Heimdallr.py b/Heimdallr.py index 4dce5c5..7e54dbe 100644 --- a/Heimdallr.py +++ b/Heimdallr.py @@ -134,4 +134,5 @@ if __name__ == "__main__": bot.load_extension("commands.infractions") bot.load_extension("commands.self_roles") bot.load_extension("commands.polls") + bot.load_extension("commands.bot_messages") bot.start(getenv("DISCORD_TOKEN")) diff --git a/commands/bot_messages.py b/commands/bot_messages.py new file mode 100644 index 0000000..d036df8 --- /dev/null +++ b/commands/bot_messages.py @@ -0,0 +1,190 @@ +from copy import deepcopy +import json +from json.decoder import JSONDecodeError +import logging +from naff import ( + Client, + Extension, + slash_command, + slash_option, + InteractionContext, + Modal, + ParagraphText, + ModalContext, + context_menu, + Permissions, + CommandTypes, + Message, +) + +from database import BotMessages as BotMessagesModel + + +message_creation_modal = Modal( + custom_id="bot-message-create", + title=f"Create a message as the bot", + components=[ + ParagraphText( + custom_id="embeds", + label="Embeds", + placeholder="Embeds as JSON", + required=False, + ), + ParagraphText( + custom_id="content", + label="Text", + placeholder="Lorem ipsum dolor sit amet and so on", + required=False, + ), + ] +) + + +class BotMessages(Extension): + def __init__(self, client: Client) -> None: + self.client = client + + @slash_command( + name="bot-message-create", + description="Create a message as the bot.", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + async def bot_message_create_command(self, ctx: InteractionContext): + + + await ctx.send_modal(message_creation_modal) + + modal_ctx: ModalContext = await self.client.wait_for_modal(message_creation_modal, author=ctx.author) + if modal_ctx.custom_id != "bot-message-create": + return + + embeds_string: str = modal_ctx.responses["embeds"] + content_string: str = modal_ctx.responses["content"] + if ( + (embeds_string is None or embeds_string == "") + and (content_string is None or content_string == "") + ): + await modal_ctx.send( + "You must provide either an embed or text.", + ephemeral=True, + ) + return + + embed: dict | None = None + embeds: list | None = None + + try: + if embeds_string: + embeds_temp = json.loads(embeds_string) + if isinstance(embeds_temp, list): + embeds = embeds_temp + elif isinstance(embeds_temp, dict): + embed = embeds_temp + except JSONDecodeError: + await modal_ctx.send( + "The embeds were not valid JSON.", + ephemeral=True, + ) + return + + msg = await ctx.channel.send( + content=content_string if content_string else None, + embed=embed, + embeds=embeds, + ) + BotMessagesModel.create( + guild_id=msg.guild.id, + channel_id=msg.channel.id, + message_id=msg.id, + ) + await modal_ctx.send( + "Message created!", + ephemeral=True, + ) + + @context_menu( + name="Edit bot message", + context_type=CommandTypes.MESSAGE, + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + async def edit_bot_message_context_menu(self, ctx: InteractionContext): + message: Message = ctx.target + if message.author.id != self.client.user.id: + await ctx.send( + "This is not a bot message.", + ephemeral=True, + ) + return + + bot_message: BotMessagesModel | None = BotMessagesModel.get( + BotMessagesModel.guild_id == message.channel.guild.id, + BotMessagesModel.channel_id == message.channel.id, + BotMessagesModel.message_id == message.id, + ) + + if bot_message is None: + await ctx.send( + "This is not an editable bot message.", + ephemeral=True, + ) + return + modal = deepcopy(message_creation_modal) + modal.title = "Edit bot message" + modal.components[0].value = json.dumps( + [e.to_dict() for e in message.embeds] if message.embeds else "", + indent=4, + ) + modal.components[1].value = message.content + + await ctx.send_modal(modal) + + modal_ctx: ModalContext = await self.client.wait_for_modal(modal, author=ctx.author) + if modal_ctx.custom_id != "bot-message-create": + return + + embeds_string: str = modal_ctx.responses["embeds"] + content_string: str = modal_ctx.responses["content"] + if ( + (embeds_string is None or embeds_string == "") + and (content_string is None or content_string == "") + ): + await modal_ctx.send( + "You must provide either an embed or text.", + ephemeral=True, + ) + return + + embed: dict | None = None + embeds: list | None = None + + try: + if embeds_string: + embeds_temp = json.loads(embeds_string) + if isinstance(embeds_temp, list): + embeds = embeds_temp + elif isinstance(embeds_temp, dict): + embed = embeds_temp + except JSONDecodeError: + await modal_ctx.send( + "The embeds were not valid JSON.", + ephemeral=True, + ) + return + + await message.edit( + content=content_string if content_string else None, + embed=embed, + embeds=embeds, + ) + await modal_ctx.send( + "Message edited!", + ephemeral=True, + ) + + +def setup(client: Client): + BotMessagesModel.create_table() + BotMessages(client) + logging.info("BotMessages extension loaded") \ No newline at end of file diff --git a/commands/polls.py b/commands/polls.py index 39634d8..73f47b0 100644 --- a/commands/polls.py +++ b/commands/polls.py @@ -23,8 +23,10 @@ from naff import ( Task, IntervalTrigger, GuildText, + PartialEmoji, ) from naff.api import events +from naff.client.errors import HTTPException from database import Polls as PollsModel, PollVotes as PollVotesModel from peewee import fn @@ -70,7 +72,7 @@ def generate_poll_embed( embed = Embed( title=title, description=("\n".join(data) if data else None), - ) + ) sum_votes = sum(votes) for i, (emoji, option) in enumerate(options): @@ -217,9 +219,13 @@ class Polls(Extension): buttons: List[Button] = [] for i, option in enumerate(options): + try: + emoji = PartialEmoji.from_str((option[0] or num_to_emoji(i + 1))) + except ValueError: + emoji = num_to_emoji(i + 1) buttons.append( Button( - emoji=(option[0] or num_to_emoji(i + 1)), + emoji=emoji, style=ButtonStyles.PRIMARY, custom_id=f"poll-vote:{poll_entry.id}:{i}", ) @@ -227,9 +233,10 @@ class Polls(Extension): buttons.append( Button( - label="Delete", + emoji="🔒", + label="Lock", style=ButtonStyles.DANGER, - custom_id=f"poll-delete:{poll_entry.id}", + custom_id=f"poll-lock:{poll_entry.id}", ) ) embed = generate_poll_embed( @@ -239,10 +246,18 @@ class Polls(Extension): multiple_choice=multiple_choice, expires=duration, ) - poll_message: Message = await modal_ctx.send( - embed=embed, - components=spread_to_rows(*buttons), - ) + try: + poll_message: Message = await modal_ctx.send( + embed=embed, + components=spread_to_rows(*buttons), + ) + except HTTPException as e: + logging.error(f"Error sending poll: {e}") + await modal_ctx.send( + "Error during poll creation. NB: You can not use server-specific emojis", + ephemeral=True, + ) + return poll_entry.message_id = poll_message.id poll_entry.channel_id = poll_message.channel.id @@ -346,8 +361,8 @@ class Polls(Extension): ) await ctx.message.edit(embed=embed) - - elif ctx.custom_id.startswith("poll-delete:"): + + elif ctx.custom_id.startswith("poll-lock:"): poll_id = ctx.custom_id.split(":", 1)[1] poll_entry: PollsModel | None = PollsModel.get_or_none( @@ -357,17 +372,17 @@ class Polls(Extension): await ctx.send("That poll doesn't exist.") return - if ( - not ctx.author.id == int(poll_entry.author_id) - or not ctx.author.has_permission(Permissions.MANAGE_MESSAGES) - ): - await ctx.send("You don't have permission to delete that poll.") + if not ctx.author.id == int( + poll_entry.author_id + ) or not ctx.author.has_permission(Permissions.MANAGE_MESSAGES): + await ctx.send("You don't have permission to lock that poll.") return - poll_entry.delete_instance() - PollVotesModel.delete().where(PollVotesModel.poll_id == poll_id).execute() - await ctx.message.delete() - await ctx.send("Poll deleted.") + poll_entry.expires = datetime.now() - timedelta(minutes=1) + poll_entry.save() + + await self.poll_expiry_check() + await ctx.send("Poll locked.") @Task.create(IntervalTrigger(minutes=1)) async def poll_expiry_check(self): @@ -376,7 +391,9 @@ class Polls(Extension): polls_q: List[PollsModel] = PollsModel.select().where(PollsModel.expires < now) for poll_entry in polls_q: - channel: GuildText = await self.client.fetch_channel(int(poll_entry.channel_id)) + channel: GuildText = await self.client.fetch_channel( + int(poll_entry.channel_id) + ) if not channel: continue @@ -386,8 +403,6 @@ class Polls(Extension): await message.edit(components=[]) - - def setup(client: Client): PollsModel.create_table() diff --git a/commands/self_roles.py b/commands/self_roles.py index acaea6e..05537fe 100644 --- a/commands/self_roles.py +++ b/commands/self_roles.py @@ -536,7 +536,7 @@ class SelfRoles(Extension): ) @context_menu( - name="Regenerate group", + name="Regenerate role group message", context_type=CommandTypes.MESSAGE, dm_permission=False, default_member_permissions=Permissions.MANAGE_ROLES,