diff --git a/bot.py b/bot.py index edc2e4b319..a3ea05cd58 100644 --- a/bot.py +++ b/bot.py @@ -4,6 +4,7 @@ import asyncio import copy import hashlib +import io import os import re import string @@ -59,6 +60,24 @@ logger = getLogger(__name__) + +class SnippetAttachment: + """In-memory attachment loaded from the snippet GridFS bucket.""" + + def __init__(self, file_data, metadata, is_image): + self.file_data = file_data + self.id = 0 + self.url = f"attachment://{metadata['filename']}" + self.filename = metadata["filename"] + self.size = metadata["length"] + self.width = None + self.is_snippet_attachment = True + self.is_snippet_image = is_image + + async def to_file(self): + return discord.File(io.BytesIO(self.file_data), filename=self.filename) + + temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp") if not os.path.exists(temp_dir): os.mkdir(temp_dir) @@ -374,7 +393,7 @@ async def wait_for_connected(self) -> None: await self.config.wait_until_ready() @property - def snippets(self) -> typing.Dict[str, str]: + def snippets(self) -> typing.Dict[str, typing.Union[str, typing.Dict[str, str]]]: return self.config["snippets"] @property @@ -1291,6 +1310,20 @@ def _get_snippet_command(self) -> commands.Command: return self.get_command(f"{modifiers}reply") + async def _download_snippet_attachment(self, snippet_data): + """Download and wrap a snippet attachment, returning None on failure.""" + if not isinstance(snippet_data, dict) or not snippet_data.get("file_id"): + return None + + try: + file_data, metadata = await self.api.download_snippet_attachment(snippet_data["file_id"]) + except Exception as e: + logger.warning("Failed to download snippet attachment: %s", e) + return None + + content_type = metadata.get("content_type", "") + return SnippetAttachment(file_data, metadata, content_type.startswith("image/")) + async def get_contexts(self, message, *, cls=commands.Context): """ Returns all invocation contexts from the message. @@ -1317,7 +1350,14 @@ async def get_contexts(self, message, *, cls=commands.Context): # snippets can have multiple words. try: # Use removeprefix once PY3.9+ - snippet_text = self.snippets[message.content[len(invoked_prefix) :]] + snippet_data = self.snippets[message.content[len(invoked_prefix) :]] + # Extract text from snippet (handle both old string format and new dict format) + if isinstance(snippet_data, str): + snippet_text = snippet_data + elif isinstance(snippet_data, dict): + snippet_text = snippet_data.get("text", "") + else: + snippet_text = None except KeyError: snippet_text = None @@ -1332,15 +1372,27 @@ async def get_contexts(self, message, *, cls=commands.Context): for alias in aliases: command = None + context_message = copy.copy(message) try: - snippet_text = self.snippets[alias] + snippet_data = self.snippets[alias] + # Extract text from snippet (handle both old string format and new dict format) + if isinstance(snippet_data, str): + snippet_text = snippet_data + elif isinstance(snippet_data, dict): + snippet_text = snippet_data.get("text", "") + attachment = await self._download_snippet_attachment(snippet_data) + if attachment is not None: + context_message.attachments = [*message.attachments, attachment] + else: + snippet_text = None except KeyError: command_invocation_text = alias + snippet_text = None else: command = self._get_snippet_command() command_invocation_text = f"{invoked_prefix}{command} {snippet_text}" view = StringView(invoked_prefix + command_invocation_text) - ctx_ = cls(prefix=self.prefix, view=view, bot=self, message=message) + ctx_ = cls(prefix=self.prefix, view=view, bot=self, message=context_message) ctx_.thread = thread discord.utils.find(view.skip_string, prefixes) ctx_.invoked_with = view.get_word().lower() @@ -1352,6 +1404,13 @@ async def get_contexts(self, message, *, cls=commands.Context): if snippet_text is not None: # Process snippets + snippet_name = message.content[len(invoked_prefix) :] + snippet_data = self.snippets.get(snippet_name) + attachment = await self._download_snippet_attachment(snippet_data) + if attachment is not None: + snippet_message = copy.copy(message) + snippet_message.attachments = [*message.attachments, attachment] + ctx.message = snippet_message ctx.command = self._get_snippet_command() reply_view = StringView(f"{invoked_prefix}{ctx.command} {snippet_text}") discord.utils.find(reply_view.skip_string, prefixes) diff --git a/cogs/modmail.py b/cogs/modmail.py index ca1498c2bc..758f1a9540 100644 --- a/cogs/modmail.py +++ b/cogs/modmail.py @@ -28,6 +28,9 @@ # Arg names reserved by formatreply commands (channel, recipient, author). RESERVED_ARG_NAMES = {"channel", "recipient", "author"} +SNIPPET_ATTACHMENT_BYTES_PER_MIB = 1024**2 +SNIPPET_ATTACHMENT_CONFIRMATION_SIZE = 2 * SNIPPET_ATTACHMENT_BYTES_PER_MIB + class Modmail(commands.Cog): """Commands directly related to Modmail functionality.""" @@ -116,6 +119,157 @@ def _resolve_user(self, user_str): return int(match.group(1)) return None + def _get_snippet_text(self, snippet_data) -> str: + """ + Extract text from a snippet, handling both old string format and new dict format. + + Parameters + ---------- + snippet_data : str or dict + The snippet data (either old string format or new dict format). + + Returns + ------- + str + The text content of the snippet. + """ + if isinstance(snippet_data, str): + return snippet_data + elif isinstance(snippet_data, dict): + return snippet_data.get("text", "") + return "" + + def _has_snippet_attachment(self, snippet_data) -> bool: + """Check if a snippet has an attachment.""" + return isinstance(snippet_data, dict) and bool(snippet_data.get("file_id")) + + async def _get_snippet_attachment_file(self, snippet_data): + """Return a stored snippet attachment as a Discord file, if available.""" + attachment = await self.bot._download_snippet_attachment(snippet_data) + if attachment is None: + return None + + try: + return await attachment.to_file() + except Exception: + logger.warning("Failed to prepare snippet attachment for display.", exc_info=True) + return None + + async def _delete_snippet_attachment(self, file_id, snippet_name): + """Delete a stored attachment without allowing cleanup errors to mask command errors.""" + try: + deleted = await self.bot.api.delete_snippet_attachment(file_id) + except Exception: + logger.warning("Failed to delete snippet attachment for %s.", snippet_name, exc_info=True) + return False + + if not deleted: + logger.warning("Failed to delete snippet attachment for %s.", snippet_name) + return deleted + + async def _send_snippet_preview(self, ctx, embed, snippet_data): + """Send a snippet preview and include its stored attachment.""" + if not self._has_snippet_attachment(snippet_data): + return await ctx.send(embed=embed) + + file = await self._get_snippet_attachment_file(snippet_data) + if file is None: + embed.set_footer(text="The attachment could not be retrieved.") + return await ctx.send(embed=embed) + try: + return await ctx.send(embed=embed, file=file) + except discord.HTTPException: + logger.warning("Failed to send snippet attachment preview.", exc_info=True) + embed.set_footer(text="The attachment could not be sent in this channel.") + return await ctx.send(embed=embed) + + @staticmethod + async def _send_snippet_result(ctx, embed, confirmation_message=None): + """Send a result, replacing a large-attachment confirmation when present.""" + if confirmation_message is not None: + return await confirmation_message.edit(content=None, embed=embed, view=None) + return await ctx.send(embed=embed) + + async def _validate_snippet_attachment(self, ctx, attachment, name, action): + """Validate an attachment and ask for confirmation when it is at least 2 MiB.""" + max_size_mib = self.bot.config.get("snippet_attachment_max_size") + # Convert the configured MiB limit to bytes for comparison with Discord's byte size. + max_size_bytes = max_size_mib * SNIPPET_ATTACHMENT_BYTES_PER_MIB + if attachment.size > max_size_bytes: + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=f"Attachment exceeds the maximum file size of {max_size_mib} MiB. " + f"Your file is {attachment.size / SNIPPET_ATTACHMENT_BYTES_PER_MIB:.2f} MiB.", + ) + await ctx.send(embed=embed) + return False, None + + if attachment.size < SNIPPET_ATTACHMENT_CONFIRMATION_SIZE: + return True, None + + view = discord.ui.View(timeout=30) + confirmed = None + past_tense = {"create": "created", "update": "updated"}[action] + + def error_embed(description): + return discord.Embed( + title="Error", + description=description, + color=self.bot.error_color, + ) + + async def confirm_callback(interaction: discord.Interaction): + nonlocal confirmed + if interaction.user.id != ctx.author.id: + return await interaction.response.send_message( + embed=error_embed("Only the command author can confirm."), ephemeral=True + ) + confirmed = True + await interaction.response.edit_message(view=None) + view.stop() + + async def cancel_callback(interaction: discord.Interaction): + nonlocal confirmed + if interaction.user.id != ctx.author.id: + return await interaction.response.send_message( + embed=error_embed("Only the command author can cancel."), ephemeral=True + ) + confirmed = False + await interaction.response.edit_message( + content=None, + embed=error_embed(f"Cancelled. Snippet not {past_tense}."), + view=None, + ) + view.stop() + + confirm_button = discord.ui.Button(label="Confirm", style=discord.ButtonStyle.green) + cancel_button = discord.ui.Button(label="Cancel", style=discord.ButtonStyle.red) + confirm_button.callback = confirm_callback + cancel_button.callback = cancel_callback + view.add_item(confirm_button) + view.add_item(cancel_button) + + embed = discord.Embed( + title="Confirm large attachment", + description=( + f"The attachment is " + f"{attachment.size / SNIPPET_ATTACHMENT_BYTES_PER_MIB:.2f} MiB " + f"(at least 2 MiB).\nDo you want to {action} the snippet `{name}` with this attachment?" + ), + color=self.bot.main_color, + ) + confirm_msg = await ctx.send(embed=embed, view=view) + await view.wait() + + if confirmed is None: + await confirm_msg.edit( + content=None, + embed=error_embed(f"Timed out. Snippet not {past_tense}."), + view=None, + ) + return confirmed is True, confirm_msg + @commands.command() @trigger_typing @checks.has_permissions(PermissionLevel.OWNER) @@ -250,13 +404,25 @@ async def snippet(self, ctx, *, name: str.lower = None): if snippet_name is None: embed = create_not_found_embed(name, self.bot.snippets.keys(), "Snippet") else: - val = self.bot.snippets[snippet_name] + snippet_data = self.bot.snippets[snippet_name] + snippet_text = self._get_snippet_text(snippet_data) + has_attachment = self._has_snippet_attachment(snippet_data) + + if snippet_text: + description = snippet_text + elif has_attachment: + description = "This is an attachment-only snippet." + else: + description = "(No text content)" + embed = discord.Embed( title=f'Snippet - "{snippet_name}":', - description=val, + description=description, color=self.bot.main_color, ) - return await ctx.send(embed=embed) + if snippet_name is None: + return await ctx.send(embed=embed) + return await self._send_snippet_preview(ctx, embed, snippet_data) if not self.bot.snippets: embed = discord.Embed( @@ -274,10 +440,24 @@ async def snippet(self, ctx, *, name: str.lower = None): for embed in embeds: embed.set_author(name="Snippets", icon_url=self.bot.get_guild_icon(guild=ctx.guild, size=128)) - for i, snippet in enumerate(sorted(self.bot.snippets.items())): - embeds[i // 10].add_field( - name=snippet[0], value=return_or_truncate(snippet[1], 350), inline=False - ) + for i, (snippet_name, snippet_data) in enumerate(sorted(self.bot.snippets.items())): + snippet_text = self._get_snippet_text(snippet_data) + has_attachment = self._has_snippet_attachment(snippet_data) + + if snippet_text: + display_value = return_or_truncate(snippet_text, 350) + elif has_attachment: + display_value = "This is an attachment-only snippet." + else: + display_value = "(No text)" + if has_attachment: + display_value = "📎 " + display_value + + embeds[i // 10].add_field(name=snippet_name, value=display_value, inline=False) + + if any(self._has_snippet_attachment(data) for data in self.bot.snippets.values()): + for embed in embeds: + embed.set_footer(text="📎 indicates that a snippet has an attachment.") session = EmbedPaginatorSession(ctx, *embeds) await session.run() @@ -292,20 +472,33 @@ async def snippet_raw(self, ctx, *, name: str.lower): if snippet_name is None: embed = create_not_found_embed(name, self.bot.snippets.keys(), "Snippet") else: - val = truncate(escape_code_block(self.bot.snippets[snippet_name]), 2048 - 7) + snippet_data = self.bot.snippets[snippet_name] + snippet_text = self._get_snippet_text(snippet_data) + has_attachment = self._has_snippet_attachment(snippet_data) + + if snippet_text: + val = truncate(escape_code_block(snippet_text), 2048 - 7) + description = f"```\n{val}```" + elif has_attachment: + description = "This is an attachment-only snippet." + else: + description = "(No text content)" + embed = discord.Embed( title=f'Raw snippet - "{snippet_name}":', - description=f"```\n{val}```", + description=description, color=self.bot.main_color, ) - return await ctx.send(embed=embed) + if snippet_name is None: + return await ctx.send(embed=embed) + return await self._send_snippet_preview(ctx, embed, snippet_data) @snippet.command(name="add", aliases=["create", "make"]) @checks.has_permissions(PermissionLevel.SUPPORTER) - async def snippet_add(self, ctx, name: str.lower, *, value: commands.clean_content): + async def snippet_add(self, ctx, name: str.lower, *, value: commands.clean_content = None): """ - Add a snippet. + Add a snippet with an optional attachment. Simply to add a snippet, do: ``` {prefix}snippet add hey hello there :) @@ -315,7 +508,21 @@ async def snippet_add(self, ctx, name: str.lower, *, value: commands.clean_conte To add a multi-word snippet name, use quotes: ``` {prefix}snippet add "two word" this is a two word snippet. ``` + + You can also attach one file (10 MiB by default) to include with the snippet. + Attachments are stored in the database and frequent use can consume significant storage. """ + if not value and len(ctx.message.attachments) == 0: + return await ctx.send_help(ctx.command) + + if len(ctx.message.attachments) > 1: + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description="You can only attach one file to a snippet.", + ) + return await ctx.send(embed=embed) + if self.bot.get_command(name): embed = discord.Embed( title="Error", @@ -347,15 +554,91 @@ async def snippet_add(self, ctx, name: str.lower, *, value: commands.clean_conte ) return await ctx.send(embed=embed) - self.bot.snippets[name] = value - await self.bot.config.update() + # Handle optional attachment + file_id = None + attachment_info = None + confirm_msg = None + if ctx.message.attachments: + attachment = ctx.message.attachments[0] + is_valid, confirm_msg = await self._validate_snippet_attachment(ctx, attachment, name, "create") + if not is_valid: + return + + # Download and upload to GridFS + try: + file_data = await attachment.read() + file_id = await self.bot.api.upload_snippet_attachment( + file_data, + attachment.filename, + attachment.content_type or "application/octet-stream", + ) + attachment_info = attachment.filename + except Exception as e: + logger.error("Failed to upload snippet attachment: %s", e) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description="Failed to upload attachment. Please try again.", + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + if file_id and (self.bot.get_command(name) or name in self.bot.snippets or name in self.bot.aliases): + cleanup_failed = not await self._delete_snippet_attachment(file_id, name) + description = ( + f"The name `{name}` became unavailable while the attachment was being uploaded. " + "Please try again with another name." + ) + if cleanup_failed: + description += ( + " The uploaded file could not be cleaned up. If this error happens again, " + "contact the bot owner and ask them to check the logs." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + # Store snippet as dict with text and optional file_id + snippet_data = {"text": value or ""} + if file_id: + snippet_data["file_id"] = file_id + + self.bot.snippets[name] = snippet_data + try: + await self.bot.config.update() + except Exception: + logger.error("Failed to save snippet %s.", name, exc_info=True) + self.bot.snippets.pop(name, None) + cleanup_failed = False + if file_id: + cleanup_failed = not await self._delete_snippet_attachment(file_id, name) + + description = "Failed to save the snippet. Please try again." + if cleanup_failed: + description += ( + " The uploaded file could not be cleaned up. If this error happens again, " + "contact the bot owner and ask them to check the logs." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + description = "Successfully created snippet." + if attachment_info: + description += f"\nAttachment: `{attachment_info}`" embed = discord.Embed( title="Added snippet", color=self.bot.main_color, - description="Successfully created snippet.", + description=description, ) - return await ctx.send(embed=embed) + + return await self._send_snippet_result(ctx, embed, confirm_msg) def _fix_aliases(self, snippet_being_deleted: str) -> Tuple[List[str]]: """ @@ -412,6 +695,9 @@ def _fix_aliases(self, snippet_being_deleted: str) -> Tuple[List[str]]: async def snippet_remove(self, ctx, *, name: str.lower): """Remove a snippet.""" if name in self.bot.snippets: + snippet_data = self.bot.snippets[name] + file_id = snippet_data.get("file_id") if isinstance(snippet_data, dict) else None + deleted_aliases, edited_aliases = self._fix_aliases(name) deleted_aliases_string = ",".join(f"`{alias}`" for alias in deleted_aliases) @@ -425,10 +711,10 @@ async def snippet_remove(self, ctx, *, name: str.lower): deleted_aliases_output = None if len(edited_aliases) == 1: - alias, val = edited_aliases.popitem() + alias, val = next(iter(edited_aliases.items())) edited_aliases_output = ( f"Steps pointing to this snippet have been removed from the `{alias}` alias" - f" (previous value: `{val}`).`" + f" (previous value: `{val}`)." ) elif edited_aliases: alias_list = "\n".join( @@ -456,33 +742,232 @@ async def snippet_remove(self, ctx, *, name: str.lower): description=description, ) self.bot.snippets.pop(name) - await self.bot.config.update() + try: + await self.bot.config.update() + except Exception: + logger.error("Failed to save removal of snippet %s.", name, exc_info=True) + self.bot.snippets[name] = snippet_data + self.bot.aliases.update(deleted_aliases) + self.bot.aliases.update(edited_aliases) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=f"Failed to remove snippet `{name}`. Please try again.", + ) + return await ctx.send(embed=embed) + + if file_id: + attachment_deleted = await self._delete_snippet_attachment(file_id, name) + + if not attachment_deleted: + self.bot.snippets[name] = snippet_data + self.bot.aliases.update(deleted_aliases) + self.bot.aliases.update(edited_aliases) + + rollback_failed = False + try: + await self.bot.config.update() + except Exception: + logger.error("Failed to roll back removal of snippet %s.", name, exc_info=True) + rollback_failed = True + + description = ( + f"Snippet `{name}` was not removed because its attachment could not be " + "deleted. Please try again." + ) + if rollback_failed: + description = ( + f"The attachment for snippet `{name}` could not be deleted, and the " + "database rollback also failed. Contact the bot owner and check the logs " + "before retrying." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) else: embed = create_not_found_embed(name, self.bot.snippets.keys(), "Snippet") await ctx.send(embed=embed) @snippet.command(name="edit") @checks.has_permissions(PermissionLevel.SUPPORTER) - async def snippet_edit(self, ctx, name: str.lower, *, value): + async def snippet_edit(self, ctx, name: str.lower, *, value: commands.clean_content = None): """ - Edit a snippet. + Edit a snippet's text and/or attachment. To edit a multi-word snippet name, use quotes: ``` {prefix}snippet edit "two word" this is a new two word snippet. ``` - """ - if name in self.bot.snippets: - self.bot.snippets[name] = value - await self.bot.config.update() + Attach one new file to replace the existing attachment. + Editing without a file removes the existing attachment. + Omit the text to keep the existing text. + """ + if len(ctx.message.attachments) > 1: embed = discord.Embed( - title="Edited snippet", - color=self.bot.main_color, - description=f'`{name}` will now send "{value}".', + title="Error", + color=self.bot.error_color, + description="You can only attach one file to a snippet.", ) - else: + return await ctx.send(embed=embed) + + if name not in self.bot.snippets: embed = create_not_found_embed(name, self.bot.snippets.keys(), "Snippet") - await ctx.send(embed=embed) + return await ctx.send(embed=embed) + + snippet_data = self.bot.snippets[name] + + # Handle the legacy string format as well as attachment-aware snippets. + if isinstance(snippet_data, str): + old_text = snippet_data + old_file_id = None + else: + old_text = snippet_data.get("text", "") + old_file_id = snippet_data.get("file_id") + + has_new_attachment = bool(ctx.message.attachments) + if value is None and not has_new_attachment and not old_file_id: + return await ctx.send_help(ctx.command) + + new_text = value if value is not None else old_text + if not has_new_attachment and not new_text: + if old_file_id: + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=( + "Removing this attachment would leave the snippet empty. " + "Provide replacement text or attach a new file." + ), + ) + return await ctx.send(embed=embed) + return await ctx.send_help(ctx.command) + + new_file_id = None + attachment_info = None + confirm_msg = None + if has_new_attachment: + attachment = ctx.message.attachments[0] + is_valid, confirm_msg = await self._validate_snippet_attachment(ctx, attachment, name, "update") + if not is_valid: + return + + try: + file_data = await attachment.read() + new_file_id = await self.bot.api.upload_snippet_attachment( + file_data, + attachment.filename, + attachment.content_type or "application/octet-stream", + ) + attachment_info = attachment.filename + except Exception as e: + logger.error("Failed to upload snippet attachment: %s", e) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description="Failed to upload attachment. Please try again.", + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + if new_file_id and self.bot.snippets.get(name) is not snippet_data: + cleanup_failed = not await self._delete_snippet_attachment(new_file_id, name) + description = ( + f"Snippet `{name}` changed while the attachment was being uploaded. Please try again." + ) + if cleanup_failed: + description += ( + " The uploaded file could not be cleaned up. If this error happens again, " + "contact the bot owner and ask them to check the logs." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + updated_snippet = {"text": new_text or ""} + if new_file_id: + updated_snippet["file_id"] = new_file_id + + self.bot.snippets[name] = updated_snippet + try: + await self.bot.config.update() + except Exception: + logger.error("Failed to save edits to snippet %s.", name, exc_info=True) + self.bot.snippets[name] = snippet_data + cleanup_failed = False + if new_file_id: + cleanup_failed = not await self._delete_snippet_attachment(new_file_id, name) + + description = "Failed to save the snippet changes. Please try again." + if cleanup_failed: + description += ( + " The newly uploaded file could not be cleaned up. If this error happens again, " + "contact the bot owner and ask them to check the logs." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + if old_file_id: + old_attachment_deleted = await self._delete_snippet_attachment(old_file_id, name) + + if not old_attachment_deleted: + self.bot.snippets[name] = snippet_data + + rollback_failed = False + try: + await self.bot.config.update() + except Exception: + logger.error("Failed to roll back edits to snippet %s.", name, exc_info=True) + rollback_failed = True + + cleanup_failed = False + if new_file_id and not rollback_failed: + cleanup_failed = not await self._delete_snippet_attachment(new_file_id, name) + + if rollback_failed: + description = ( + f"The previous attachment for snippet `{name}` could not be deleted, and " + "the database rollback also failed. Contact the bot owner and check the logs " + "before retrying." + ) + else: + description = ( + f"Snippet `{name}` was not updated because its previous attachment could not " + "be deleted. Please try again." + ) + if cleanup_failed: + description += ( + " The newly uploaded file could not be cleaned up. If this error happens " + "again, contact the bot owner and ask them to check the logs." + ) + embed = discord.Embed( + title="Error", + color=self.bot.error_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) + + description = f"`{name}` has been updated." + if value is not None: + description += f'\nText: "{truncate(value, 100)}"' + if attachment_info: + description += f"\nNew attachment: `{attachment_info}`" + elif old_file_id: + description += "\nAttachment removed." + + embed = discord.Embed( + title="Edited snippet", + color=self.bot.main_color, + description=description, + ) + return await self._send_snippet_result(ctx, embed, confirm_msg) @snippet.command(name="rename") @checks.has_permissions(PermissionLevel.SUPPORTER) diff --git a/cogs/utility.py b/cogs/utility.py index 283823d39c..173690f0d2 100644 --- a/cogs/utility.py +++ b/cogs/utility.py @@ -195,6 +195,8 @@ async def send_error_message(self, error): command = self.context.kwargs.get("command") val = self.context.bot.snippets.get(command) if val is not None: + if isinstance(val, dict): + val = val.get("text") or "This is an attachment-only snippet." embed = discord.Embed(title=f"{command} is a snippet.", color=self.context.bot.main_color) embed.add_field(name=f"`{command}` will send:", value=val, inline=False) diff --git a/core/clients.py b/core/clients.py index ff0d3ff8cd..8884baa8bc 100644 --- a/core/clients.py +++ b/core/clients.py @@ -1,14 +1,17 @@ import secrets import sys from json import JSONDecodeError -from typing import Any, Dict, Union, Optional +from typing import Any, Dict, Union, Optional, Tuple import discord from discord import Member, DMChannel, TextChannel, Message from discord.ext import commands from aiohttp import ClientResponseError, ClientResponse -from motor.motor_asyncio import AsyncIOMotorClient +from bson import ObjectId +from bson.errors import InvalidId +from gridfs.errors import NoFile +from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket from pymongo.errors import ConfigurationError from core.models import InvalidConfigError, getLogger @@ -399,6 +402,7 @@ async def append_log( message_id: str = "", channel_id: str = "", type_: str = "thread_message", + attachments=None, ) -> dict: return NotImplemented @@ -461,6 +465,7 @@ def __init__(self, bot): sys.exit(0) super().__init__(bot, db) + self.fs = AsyncIOMotorGridFSBucket(db, bucket_name="snippet_attachments") async def setup_indexes(self): """Setup text indexes so we can use the $search operator""" @@ -661,11 +666,13 @@ async def append_log( message_id: str = "", channel_id: str = "", type_: str = "thread_message", + attachments=None, ) -> dict: channel_id = str(channel_id) or (str(message.channel.id) if message else "") message_id = str(message_id) or (str(message.id) if message else "") if message: + log_attachments = message.attachments if attachments is None else attachments content = message.content or "" if forwarded := extract_forwarded_content(message): if content: @@ -695,7 +702,7 @@ async def append_log( "size": a.size, "url": a.url, } - for a in message.attachments + for a in log_attachments ], } else: @@ -715,7 +722,6 @@ async def append_log( "type": type_, "attachments": [], } - return await self.logs.find_one_and_update( {"channel_id": channel_id}, {"$push": {"messages": data}}, @@ -807,6 +813,89 @@ async def get_user_info(self) -> Optional[dict]: } } + # ==================== GridFS Methods for Snippet Attachments ==================== + + async def upload_snippet_attachment( + self, file_data: bytes, filename: str, content_type: str = "application/octet-stream" + ) -> str: + """ + Upload a file to GridFS for snippet attachments. + + Parameters + ---------- + file_data : bytes + The raw file data to upload. + filename : str + The original filename. + content_type : str + The MIME type of the file. + + Returns + ------- + str + The string representation of the GridFS file ID. + """ + file_id = await self.fs.upload_from_stream( + filename, + file_data, + metadata={"content_type": content_type, "filename": filename}, + ) + logger.debug("Uploaded snippet attachment %s with file_id %s.", filename, file_id) + return str(file_id) + + async def download_snippet_attachment(self, file_id: str) -> Tuple[bytes, Dict[str, Any]]: + """ + Download a file from GridFS. + + Parameters + ---------- + file_id : str + The string representation of the GridFS file ID. + + Returns + ------- + Tuple[bytes, Dict[str, Any]] + A tuple of (file_data, metadata) where metadata includes filename and content_type. + """ + grid_out = await self.fs.open_download_stream(ObjectId(file_id)) + file_data = await grid_out.read() + metadata = { + "filename": grid_out.filename, + "content_type": ( + grid_out.metadata.get("content_type", "application/octet-stream") + if grid_out.metadata + else "application/octet-stream" + ), + "length": grid_out.length, + } + logger.debug("Downloaded snippet attachment with file_id %s.", file_id) + return file_data, metadata + + async def delete_snippet_attachment(self, file_id: str) -> bool: + """ + Delete a file from GridFS. + + Parameters + ---------- + file_id : str + The string representation of the GridFS file ID. + + Returns + ------- + bool + True if deletion was successful. + """ + try: + await self.fs.delete(ObjectId(file_id)) + logger.debug("Deleted snippet attachment with file_id %s.", file_id) + return True + except (InvalidId, NoFile): + logger.info("Snippet attachment %s was already absent.", file_id) + return True + except Exception as e: + logger.warning("Failed to delete snippet attachment %s: %s", file_id, e) + return False + class PluginDatabaseClient: def __init__(self, bot): diff --git a/core/config.py b/core/config.py index b3c79bcd69..51d3233d21 100644 --- a/core/config.py +++ b/core/config.py @@ -164,6 +164,8 @@ class ConfigManager: "thread_creation_menu_embed_large_image": False, "thread_creation_menu_embed_footer_icon_url": None, "thread_creation_menu_embed_color": str(discord.Color.green()), + # snippet attachments + "snippet_attachment_max_size": 10, # in MiB } private_keys = { @@ -243,6 +245,8 @@ class ConfigManager: duration_seconds = {"snooze_default_duration", "thread_creation_menu_timeout"} + mebibytes = {"snippet_attachment_max_size"} + booleans = { "use_user_id_channel_name", "use_timestamp_channel_name", @@ -314,6 +318,23 @@ def __init__(self, bot): def __repr__(self): return repr(self._cache) + @staticmethod + def _convert_mebibytes(value: typing.Any) -> int: + """Convert a positive whole-number MiB value without rounding it.""" + if isinstance(value, bool): + raise InvalidConfigError("Must be a positive whole number of MiB.") + + if isinstance(value, int): + converted = value + elif isinstance(value, str) and re.fullmatch(r"[1-9]\d*", value.strip()): + converted = int(value) + else: + raise InvalidConfigError("Must be a positive whole number of MiB.") + + if converted <= 0: + raise InvalidConfigError("Must be a positive whole number of MiB.") + return converted + def populate_cache(self) -> dict: data = deepcopy(self.defaults) @@ -422,6 +443,13 @@ def get(self, key: str, *, convert: bool = True) -> typing.Any: logger.warning("Invalid %s %s.", key, value) value = self.remove(key) + elif key in self.mebibytes: + try: + value = self._convert_mebibytes(value) + except InvalidConfigError: + logger.warning("Invalid %s %s.", key, value) + value = self.remove(key) + elif key in self.force_str: # Temporary: as we saved in int previously, leading to int32 overflow, # this is transitioning IDs to strings @@ -525,6 +553,9 @@ async def set(self, key: str, item: typing.Any, convert=True) -> None: duration_seconds = int((time.dt - now).total_seconds()) return self.__setitem__(key, duration_seconds) + elif key in self.mebibytes: + return self.__setitem__(key, self._convert_mebibytes(item)) + elif key in self.enums: if isinstance(item, self.enums[key]): # value is an enum type diff --git a/core/config_help.json b/core/config_help.json index fedf9279ed..0679dfb73c 100644 --- a/core/config_help.json +++ b/core/config_help.json @@ -856,6 +856,19 @@ "See also: `anonymous_snippets`." ] }, + "snippet_attachment_max_size": { + "default": "10 MiB", + "description": "Maximum file size in mebibytes (MiB) for attachments when creating or editing snippets.", + "examples": [ + "`{prefix}config set snippet_attachment_max_size 5` (5 MiB)", + "`{prefix}config set snippet_attachment_max_size 20` (20 MiB)" + ], + "notes": [ + "Attachments larger than this size will be rejected when adding or editing snippets.", + "The value must be a positive whole number of MiB; fractional and negative values are rejected.", + "Increasing this limit can substantially increase long-term database storage usage, especially when attachments are used frequently." + ] + }, "require_close_reason": { "default" : "No", "description": "Require a reason to close threads.", @@ -1594,4 +1607,4 @@ "Color names map to the built-in palette (e.g., 'red', 'green', 'blurple')." ] } -} \ No newline at end of file +} diff --git a/core/thread.py b/core/thread.py index 8bc83f7324..7c34355ec7 100644 --- a/core/thread.py +++ b/core/thread.py @@ -1750,12 +1750,23 @@ async def reply( msg = None if msg is not None: + log_attachments = None + if any( + getattr(attachment, "is_snippet_attachment", False) for attachment in message.attachments + ): + original_attachments = [ + attachment + for attachment in message.attachments + if not getattr(attachment, "is_snippet_attachment", False) + ] + log_attachments = [*original_attachments, *msg.attachments] tasks.append( self.bot.api.append_log( message, message_id=msg.id, channel_id=self.channel.id, type_="anonymous" if anonymous else "thread_message", + attachments=log_attachments, ) ) else: @@ -2004,7 +2015,27 @@ async def send( images = [] attachments = [] - for attachment in ext: + files_to_upload = [] + + # List to track snippet images that should be uploaded but not listed as file attachments + snippet_images_to_upload = [] + + for i, a in enumerate(message.attachments): + attachment = ext[i] + if getattr(a, "is_snippet_attachment", False): + if getattr(a, "is_snippet_image", False): + # Embed snippet images using attachment:// syntax. + snippet_images_to_upload.append(a) + else: + files_to_upload.append(a) + elif is_image_url(attachment[0]): + images.append(attachment) + else: + attachments.append(attachment) + + # Forwarded attachments are represented only in ``ext`` rather than + # ``message.attachments``, so classify the remaining entries separately. + for attachment in ext[len(message.attachments) :]: if is_image_url(attachment[0]): images.append(attachment) else: @@ -2078,6 +2109,26 @@ def lottie_to_png(data): images.append((None, i.name, True)) embedded_image = False + discord_files = [] + + # Handle snippet images first (embedded directly) + for a in snippet_images_to_upload: + try: + file = await a.to_file() + except Exception: + logger.warning( + "Failed to convert snippet image %s to file.", + getattr(a, "filename", "unknown"), + exc_info=True, + ) + continue + + if not embedded_image: + embed.set_image(url=f"attachment://{a.filename}") + embed.add_field(name="Image", value=a.filename) + embedded_image = True + # Only reference images that were successfully converted and will be sent. + discord_files.append(file) prioritize_uploads = any(i[1] is not None for i in images) @@ -2171,7 +2222,9 @@ def lottie_to_png(data): embed.colour = self.bot.recipient_color if (from_mod or note) and not thread_creation: - delete_message = not bool(message.attachments) + delete_message = not any( + not getattr(attachment, "is_snippet_attachment", False) for attachment in message.attachments + ) # Only delete the source command message when it's in a guild text # channel; attempting to delete a DM message can raise 50003. if ( @@ -2223,6 +2276,12 @@ def lottie_to_png(data): else: mentions = None + for att in files_to_upload: + try: + discord_files.append(await att.to_file()) + except Exception: + logger.warning("Failed to convert snippet attachment to file.", exc_info=True) + if plain: if from_mod and not isinstance(destination, discord.TextChannel): # Plain to user (DM) @@ -2234,8 +2293,10 @@ def lottie_to_png(data): body = embed.description or "" plain_message = f"{prefix}{embed.author.name}:** {body}" - files = [] + files = discord_files[:] for att in message.attachments: + if getattr(att, "is_snippet_attachment", False): + continue try: files.append(await att.to_file()) except Exception: @@ -2246,11 +2307,11 @@ def lottie_to_png(data): # Plain to mods footer_text = embed.footer.text if embed.footer else "" embed.set_footer(text=f"[PLAIN] {footer_text}".strip()) - msg = await destination.send(mentions, embed=embed) + msg = await destination.send(mentions, embed=embed, files=discord_files or None) else: try: - msg = await destination.send(mentions, embed=embed) + msg = await destination.send(mentions, embed=embed, files=discord_files or None) except discord.NotFound: if ( isinstance(destination, discord.TextChannel) @@ -2260,7 +2321,7 @@ def lottie_to_png(data): logger.info("Thread channel missing while sending; attempting restore and resend.") await self.restore_from_snooze() destination = self.channel or destination - msg = await destination.send(mentions, embed=embed) + msg = await destination.send(mentions, embed=embed, files=discord_files or None) else: logger.warning("Channel not found during send.") raise