Implement new attachment handlers

.dockerignore

@@ -5,7 +5,6 @@
 !/cogs
 !/core
 !/plugins
-!/src
 !*.py
 !LICENSE
 !pdm.lock

.pre-commit-config.yaml

@@ -8,13 +8,13 @@ repos:
     hooks:
       - id: ruff
   - repo: https://github.com/pdm-project/pdm
-    rev: 2.10.0 # a PDM release exposing the hook
+    rev: 2.11.2 # a PDM release exposing the hook
     hooks:
       - id: pdm-export
         # command arguments, e.g.:
         args: [ '-o', 'requirements.txt', '--without-hashes' ]
         files: ^pdm.lock$
   - repo: https://github.com/pdm-project/pdm
-    rev: 2.10.0 # a PDM release exposing the hook
+    rev: 2.11.2 # a PDM release exposing the hook
     hooks:
       - id: pdm-lock-check

bot.py

@@ -24,6 +24,10 @@
 from emoji import UNICODE_EMOJI
 from pkg_resources import parse_version
 
+from core.attachments.attachment_handler import IAttachmentHandler
+from core.attachments.errors import AttachmentSizeException
+from core.attachments.mongo_attachment_client import MongoAttachmentHandler
+from core.attachments.s3_attachment_client import S3AttachmentHandler
 from core.blocklist import Blocklist, BlockReason
 
 try:
@@ -92,6 +96,27 @@ def __init__(self):
 
         self.blocklist = Blocklist(bot=self)
 
+        if self.config["attachment_datastore"] == "internal":
+            logger.info("Using internal attachment handler.")
+            self.attachment_handler: IAttachmentHandler = MongoAttachmentHandler(self.api.db)
+        elif self.config["attachment_datastore"] == "s3":
+            logger.info("Using S3 attachment handler.")
+            endpoint = self.config["s3_endpoint"]
+            if endpoint is None:
+                logger.critical("S3 endpoint must be set when using the S3 attachment datastore.")
+                raise InvalidConfigError("s3_endpoint must be set.")
+            self.attachment_handler: IAttachmentHandler = S3AttachmentHandler(
+                endpoint=endpoint,
+                access_key=self.config["s3_access_key"] or None,
+                secret_key=self.config["s3_secret_key"] or None,
+                region=self.config["s3_region"] or None,
+                bucket=self.config["s3_bucket"] or None,
+            )
+        else:
+            raise InvalidConfigError("Invalid image_store option set.")
+        if self.config["max_attachment_size"] is not None:
+            self.attachment_handler.max_size = self.config["max_attachment_size"]
+
         self.startup()
 
     def get_guild_icon(
@@ -961,11 +986,22 @@ async def process_dm_modmail(self, message: discord.Message) -> None:
                 )
                 logger.info("A message was blocked from %s due to disabled Modmail.", message.author)
                 await self.add_reaction(message, blocked_emoji)
-                return await message.channel.send(embed=embed)
+                await message.channel.send(embed=embed)
+                return
 
         if not thread.cancelled:
             try:
                 await thread.send(message)
+            except AttachmentSizeException as e:
+                await self.add_reaction(message, blocked_emoji)
+                await message.channel.send(
+                    embed=discord.Embed(
+                        title="Attachment too large",
+                        description=str(e),
+                        color=self.error_color,
+                    )
+                )
+                return
             except Exception:
                 logger.error("Failed to send message:", exc_info=True)
                 await self.add_reaction(message, blocked_emoji)

core/attachments/attachment_handler.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+
+from discord.message import Message
+
+
+class IAttachmentHandler(ABC):
+    _attachment_max_size: int = 1024 * 1024 * 500
+
+    @abstractmethod
+    async def upload_attachments(self, message: Message) -> list[dict]:
+        """
+        Uploads all attachments from a message to the database
+        Parameters
+        ----------
+        message
+        Returns
+        -------
+        A dict containing what should be appended to the thread documents attachments field
+        """
+        pass
+
+    @property
+    def max_size(self) -> int:
+        """
+        The maximum size of an attachment in bytes
+        Returns
+        -------
+        int
+        """
+        return self._attachment_max_size
+
+    @max_size.setter
+    def max_size(self, value: int):
+        """
+        Set the maximum size of an attachment in bytes
+        """
+        self._attachment_max_size = value

core/attachments/errors.py

@@ -0,0 +1,13 @@
+class AttachmentSizeException(Exception):
+    """Raised when an attachment is too large.
+    Attributes:
+            size -- size of the attachment
+            max -- maximum allowed size of the attachment
+            message -- explanation of why the specific transition is invalid
+    """
+
+    def __init__(self, size: int, max_size: int, message: str):
+        self.size = size
+        self.max = max_size
+        self.message = message
+        super().__init__(self.message)

core/attachments/mongo_attachment_client.py

@@ -0,0 +1,113 @@
+import asyncio
+import datetime
+from typing import Any, Final
+
+from discord.message import Attachment, Message
+from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
+from pymongo.results import InsertOneResult
+
+from core.attachments.attachment_handler import IAttachmentHandler
+from core.attachments.errors import AttachmentSizeException
+from core.models import getLogger
+
+
+def _mongo_attachment_dict(attachment: Attachment, data: bytes) -> dict:
+    """
+    Convert a discord attachment to a dict that can be stored in the database
+    Parameters
+    ----------
+    attachment
+    data
+
+    Returns
+    -------
+    dict
+    """
+    return {
+        # same as discord id
+        "_id": attachment.id,
+        "filename": attachment.filename,
+        "content_type": attachment.content_type,
+        "width": attachment.width,
+        "height": attachment.height,
+        "description": attachment.description,
+        "size": attachment.size,
+        "data": data,
+        "uploaded_at": datetime.datetime.now(datetime.timezone.utc),
+    }
+
+
+class MongoAttachmentHandler(IAttachmentHandler):
+    logger = getLogger(__name__)
+
+    # 8 MB to bytes
+    MONGODB_MAX_SIZE: Final[int] = 1024 * 1024 * 15
+
+    def __init__(self, database: AsyncIOMotorDatabase) -> None:
+        self.client = database
+        self.attachment_collection: AsyncIOMotorCollection = database["attachments"]
+        self._attachment_max_size = self.MONGODB_MAX_SIZE
+        # self.log_collection: AsyncIOMotorCollection = database["logs"]
+
+    @IAttachmentHandler.max_size.setter
+    def max_size(self, size: int) -> None:
+        if size > self.MONGODB_MAX_SIZE:
+            self.logger.warning(
+                f"MongoDB attachment storage has a maximum attachment size of {self.MONGODB_MAX_SIZE} bytes. "
+                f"The max attachment size will be set to this value.."
+            )
+            return
+        self._attachment_max_size = size
+
+    async def _store_attachments_bulk(self, attachments: list[Attachment]) -> Any:
+        attachment_data = await asyncio.gather(*[attachment.read() for attachment in attachments])
+
+        results = await self.attachment_collection.insert_many(
+            [
+                _mongo_attachment_dict(attachment, attachment_data[index])
+                for index, attachment in enumerate(attachments)
+            ]
+        )
+
+        return results.inserted_ids
+
+    async def _store_attachment(self, attachment: Attachment) -> Any:
+        result: InsertOneResult = await self.attachment_collection.insert_one(
+            _mongo_attachment_dict(attachment, await attachment.read())
+        )
+        return result.inserted_id
+
+    async def upload_attachments(
+        self,
+        message: Message,
+    ) -> list[dict]:
+        attachments = []
+
+        for attachment in message.attachments:
+            if attachment.size > self.max_size:
+                self.logger.warning(
+                    "Attachment %s is too large to be stored in the database. It will not be uploaded...",
+                    attachment.filename,
+                )
+                raise AttachmentSizeException("Attachment too large")
+
+        if len(message.attachments) > 1:
+            await self._store_attachments_bulk(message.attachments)
+        else:
+            await self._store_attachment(message.attachments[0])
+
+        for attachment in message.attachments:
+            attachments.append(
+                {
+                    "id": attachment.id,
+                    "filename": attachment.filename,
+                    "type": "internal",
+                    # URL points to the original discord URL
+                    "url": attachment.url,
+                    "content_type": attachment.content_type,
+                    "width": attachment.width,
+                    "height": attachment.height,
+                }
+            )
+
+        return attachments

core/attachments/s3_attachment_client.py

@@ -0,0 +1,110 @@
+import datetime
+import io
+
+from discord import Message
+from minio import Minio
+from minio.commonconfig import Tags
+
+from core.attachments.attachment_handler import IAttachmentHandler
+from core.models import getLogger
+
+
+class S3AttachmentHandler(IAttachmentHandler):
+    logger = getLogger(__name__)
+
+    def __init__(
+        self,
+        endpoint: str,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        region: str | None = None,
+        bucket: str | None = "modmail-attachments",
+    ) -> None:
+        """
+        Initialize the S3AttachmentHandler with the given access key, secret key, and endpoint.
+
+        Parameters
+        ----------
+        access_key : str | None
+            The access key for the S3 bucket.
+        secret_key : str | None
+            The secret key for the S3 bucket.
+        endpoint : str
+            The endpoint for the S3 bucket.
+        region : str | None
+            The region for the S3 bucket.
+        bucket : str | None
+            The name of the S3 bucket.
+        """
+        self.bucket = bucket
+        self.region = region
+
+        self.client = Minio(
+            endpoint, access_key=access_key, secret_key=secret_key, secure=False, region=region
+        )
+
+        try:
+            self.client.bucket_exists(self.bucket)
+        except Exception as e:
+            self.logger.error(f"An error occurred while checking if the s3 bucket exists: {e}", exc_info=True)
+            raise
+
+        if not self.client.bucket_exists(self.bucket):
+            self.client.make_bucket(self.bucket)
+
+    async def upload_attachments(self, message: Message) -> list[dict]:
+        """
+        Upload attachments from a given message to the S3 bucket.
+
+        Parameters
+        ----------
+        message : Message
+            The message containing the attachments to upload.
+
+        Returns
+        -------
+        list[dict]
+            A list of dictionaries containing information about the uploaded attachments.
+        """
+        attachments = []
+
+        # Setup metadata for S3 object
+        tags = Tags.new_object_tags()
+        tags["message_id"] = str(message.id)
+        tags["channel_id"] = str(message.channel.id)
+        # tags["bot_id"] = str(bot.bot_id)
+        for attachment in message.attachments:
+            if attachment.size > self.max_size:
+                raise ValueError(
+                    f"Attachment {attachment.filename} is too large. Max size is {self.max_size} bytes."
+                )
+
+        for attachment in message.attachments:
+            self.logger.debug("Uploading attachment %s to S3", attachment.filename)
+            result = self.client.put_object(
+                self.bucket,
+                str(attachment.id),
+                io.BytesIO(await attachment.read()),
+                length=attachment.size,
+                content_type=attachment.content_type,
+                tags=tags,
+            )
+            attachments.append(
+                {
+                    "id": attachment.id,
+                    "filename": attachment.filename,
+                    "type": "openmodmail_s3",
+                    "s3": {
+                        "object": result.object_name,
+                        "bucket": result.bucket_name,
+                    }
+                    "content_type": attachment.content_type,
+                    "width": attachment.width,
+                    "height": attachment.height,
+                    "description": attachment.description,
+                    "size": attachment.size,
+                    "uploaded_at": datetime.datetime.now(datetime.timezone.utc),
+                }
+            )
+
+        return attachments