diff --git a/requirements.txt b/requirements.txt index b64820a..950813e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -51,3 +51,5 @@ typing_extensions==4.3.0 websockets==10.4 yarl==1.7.2 zonis==1.2.5 +types-aiobotocore==2.11.2 +aiobotocore==2.11.2 diff --git a/suggestions/bot.py b/suggestions/bot.py index fad449b..000fcd8 100644 --- a/suggestions/bot.py +++ b/suggestions/bot.py @@ -36,6 +36,7 @@ QueueImbalance, BlocklistedUser, PartialResponse, + InvalidFileType, ) from suggestions.http_error_parser import try_parse_http_error from suggestions.objects import Error, GuildConfig, UserConfig @@ -447,6 +448,18 @@ async def on_slash_command_error( ephemeral=True, ) + elif isinstance(exception, InvalidFileType): + return await interaction.send( + embed=self.error_embed( + "Invalid file type", + "The file you attempted to upload is not an accepted type.\n\n" + "If you believe this is an error please reach out to us via our support discord.", + error_code=ErrorCode.INVALID_FILE_TYPE, + error=error, + ), + ephemeral=True, + ) + elif isinstance(exception, ConfiguredChannelNoLongerExists): return await interaction.send( embed=self.error_embed( diff --git a/suggestions/codes.py b/suggestions/codes.py index ef7e0d4..12356b2 100644 --- a/suggestions/codes.py +++ b/suggestions/codes.py @@ -26,6 +26,7 @@ class ErrorCode(IntEnum): QUEUE_IMBALANCE = 20 MISSING_QUEUE_CHANNEL = 21 BLOCKLISTED_USER = 22 + INVALID_FILE_TYPE = 23 @classmethod def from_value(cls, value: int) -> ErrorCode: diff --git a/suggestions/cogs/suggestion_cog.py b/suggestions/cogs/suggestion_cog.py index 190185d..b98a0bd 100644 --- a/suggestions/cogs/suggestion_cog.py +++ b/suggestions/cogs/suggestion_cog.py @@ -1,21 +1,22 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Optional, cast +from typing import TYPE_CHECKING, Optional import cooldowns import disnake from commons.caching import NonExistentEntry from bot_base.wraps import WrappedChannel -from disnake import Guild, Localized +from disnake import Guild from disnake.ext import commands, components -from suggestions import checks, Stats, ErrorCode +from suggestions import checks, Stats from suggestions.clunk2 import update_suggestion_message from suggestions.cooldown_bucket import InteractionBucket from suggestions.exceptions import SuggestionTooLong, ErrorHandled -from suggestions.objects import Suggestion, GuildConfig, UserConfig, QueuedSuggestion +from suggestions.objects import Suggestion, GuildConfig, QueuedSuggestion from suggestions.objects.suggestion import SuggestionState +from suggestions.utility import r2 if TYPE_CHECKING: from alaric import Document @@ -188,15 +189,23 @@ async def suggest( ) raise ErrorHandled - image_url = image.url if isinstance(image, disnake.Attachment) else None - if image_url and not guild_config.can_have_images_in_suggestions: - await interaction.send( - self.bot.get_locale( - "SUGGEST_INNER_NO_IMAGES_IN_SUGGESTIONS", interaction.locale - ), - ephemeral=True, + image_url = None + if image is not None: + if not guild_config.can_have_images_in_suggestions: + await interaction.send( + self.bot.get_locale( + "SUGGEST_INNER_NO_IMAGES_IN_SUGGESTIONS", interaction.locale + ), + ephemeral=True, + ) + raise ErrorHandled + + image_url = await r2.upload_file_to_r2( + file_name=image.filename, + file_data=await image.read(use_cached=True), + guild_id=interaction.guild_id, + user_id=interaction.author.id, ) - raise ErrorHandled if guild_config.uses_suggestion_queue: await QueuedSuggestion.new( diff --git a/suggestions/exceptions.py b/suggestions/exceptions.py index 6669a63..ba3c1d4 100644 --- a/suggestions/exceptions.py +++ b/suggestions/exceptions.py @@ -53,3 +53,7 @@ class BlocklistedUser(CheckFailure): class PartialResponse(Exception): """A garven route returned a partial response when we require a full response""" + + +class InvalidFileType(disnake.DiscordException): + """The file you attempted to upload is not allowed.""" diff --git a/suggestions/utility/__init__.py b/suggestions/utility/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/suggestions/utility/r2.py b/suggestions/utility/r2.py new file mode 100644 index 0000000..1a886f6 --- /dev/null +++ b/suggestions/utility/r2.py @@ -0,0 +1,57 @@ +import hashlib +import logging +import mimetypes +import os +import secrets + +from aiobotocore.session import get_session + +from suggestions.exceptions import InvalidFileType + +log = logging.getLogger(__name__) + + +async def upload_file_to_r2( + *, + file_name: str, + file_data: bytes, + guild_id: int, + user_id: int, +) -> str: + """Upload a file to R2 and get the cdn url back""" + + session = get_session() + async with session.create_client( + "s3", + endpoint_url=os.environ["ENDPOINT"], + aws_access_key_id=os.environ["ACCESS_KEY"], + aws_secret_access_key=os.environ["SECRET_ACCESS_KEY"], + ) as client: + mimetype_guessed, _ = mimetypes.guess_type(file_name) + accepted_mimetypes: dict[str, set[str]] = { + "image/jpeg": {".jpeg", ".jpg"}, + "image/png": {".png"}, + "image/gif": {".gif"}, + "video/mp3": {".mp3"}, + "video/mp4": {".mp4"}, + "video/mpeg": {".mpeg"}, + "video/webm": {".webm"}, + "image/webp": {".webp"}, + "audio/webp": {".weba"}, + } + file_names = accepted_mimetypes.get(mimetype_guessed) + if file_names is None: + raise InvalidFileType + + for ext in file_names: + if file_name.endswith(ext): + break + else: + raise InvalidFileType + + file_key = hashlib.sha256(file_data + secrets.token_bytes(16)).hexdigest() + key = "{}/{}.{}".format(guild_id, file_key, ext) + await client.put_object(Bucket=os.environ["BUCKET"], Key=key, Body=file_data) + log.debug("User %s in guild %s uploaded an image", user_id, guild_id) + + return f"https://cdn.suggestions.bot/{key}"