diff --git a/admin_bot/admin_bot/cogs/image_cog.py b/admin_bot/admin_bot/cogs/image_cog.py new file mode 100644 index 0000000..ad29971 --- /dev/null +++ b/admin_bot/admin_bot/cogs/image_cog.py @@ -0,0 +1,108 @@ +import io +import discord +import logging + +from discord import app_commands +from discord.ext import commands +from utilities.image_utils import ( + apply_image_task, + load_image_from_bytes, + save_image_to_bytes, +) + +from typing import Optional, Callable + +# Just makes sure we decorate where the image_tasks go, cool python! +import utilities.image_tasks + + +class ImageCog(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.register_image_tasks() + + def register_image_tasks(self): + """ + Dynamically register context menu commands for image tasks. + """ + from utilities.image_utils import IMAGE_TASKS # Import the registered tasks + + for task_name in IMAGE_TASKS: + # Create a context menu for each task + context_menu = app_commands.ContextMenu( + name=task_name.title(), + callback=self.process_image_context, + ) + # Attach the task name to the context menu + context_menu.task_name = task_name + # Add the command to the bot + self.bot.tree.add_command(context_menu) + + async def process_image_context( + self, interaction: discord.Interaction, message: discord.Message + ): + """ + Handle image processing via context menu. + """ + task_name = interaction.command.task_name + await self.process_image_task(interaction, task_name, message) + + async def process_image_task( + self, + interaction: discord.Interaction, + task_name: str, + message: Optional[discord.Message] = None, + ): + """ + Process the image using the selected task. + """ + if not message: + await interaction.response.send_message( + "No message to process!", ephemeral=True + ) + return + + if not message.attachments: + await interaction.response.send_message( + "No image attached to this message!", ephemeral=True + ) + return + + attachment = message.attachments[0] + + if not attachment.content_type.startswith("image/"): + await interaction.response.send_message( + "The attachment is not an image!", ephemeral=True + ) + return + + await interaction.response.defer() # Acknowledge the interaction + image_bytes = await attachment.read() + + try: + # Process the image + input_image = load_image_from_bytes(image_bytes) + output_image = apply_image_task(input_image, task_name) + output_bytes = save_image_to_bytes(output_image) + + # Send the result + await interaction.followup.send( + file=discord.File(io.BytesIO(output_bytes), f"{task_name}.png") + ) + except Exception as e: + logging.error(f"Error processing image task '{task_name}': {e}") + await interaction.followup.send( + "Failed to process the image.", ephemeral=True + ) + + async def cog_unload(self): + """ + Unload the context menus when the cog is unloaded. + """ + for command in self.bot.tree.get_commands(type=discord.AppCommandType.message): + if hasattr(command, "task_name"): + self.bot.tree.remove_command(command.name) + + +async def setup(bot): + await bot.add_cog(ImageCog(bot)) diff --git a/admin_bot/admin_bot/cogs/misc.py b/admin_bot/admin_bot/cogs/misc.py index f544b6c..dbd3ee9 100644 --- a/admin_bot/admin_bot/cogs/misc.py +++ b/admin_bot/admin_bot/cogs/misc.py @@ -3,17 +3,17 @@ # MIT License -import discord -import logging -import cv2 as cv import io import os +import yaml import random -import discord.utils +import discord +import logging import asyncio -import random -import yaml import subprocess +import discord.utils + +import cv2 as cv from PIL import Image, ImageDraw, ImageFilter, UnidentifiedImageError @@ -53,29 +53,6 @@ async def bee(self, ctx: discord.Interaction): await ctx.response.send_message(str(line)) - @app_commands.command(name="myvote") - async def myvote(self, ctx: discord.Interaction, img: str): - """ - Automates that silly thing Aaron does. - """ - - # This is literally begging to be used for remote code injection - subprocess.run(["wget", img, "-O", "/tmp/to_vote.png"]) - - try: - pip = Image.open("/tmp/to_vote.png") - frame = Image.open("admin_bot/resources/vote.png") - - pip = pip.resize((375, 250)) - - frame.paste(pip, (40, 240)) - frame.save("/tmp/myvoteout.png", quality=95) - - await ctx.response.send_message(file=discord.File("/tmp/myvoteout.png")) - except UnidentifiedImageError: - logging.error("Could not download image when requested.") - await ctx.response.send_message("Error opening that image!") - @app_commands.command(name="aaron") async def aaron(self, ctx: discord.Interaction): """ @@ -88,36 +65,6 @@ async def aaron(self, ctx: discord.Interaction): file=discord.File(srcdir + random.choice(os.listdir(srcdir))) ) - @app_commands.command(name="whiteboard") - async def whiteboard(self, ctx: discord.Interaction, img: str): - """ - Puts a thing on the whiteboard - """ - - await ctx.response.defer() # We can expect that this command will take a while - - subprocess.run(["wget", img, "-O", "/tmp/to_vote.png"]) - - try: - background = Image.open("admin_bot/resources/look_at_this/background.png") - pip = Image.open("/tmp/to_vote.png") - foreground = Image.open("admin_bot/resources/look_at_this/foreground.png") - - pip = pip.resize((1000, 2000)) - - background.paste(pip, (1500, 75)) - background.paste( - foreground, - (0, 0), - foreground, # Transparency layer - ) - background.save("/tmp/myvoteout.png", quality=95) - - await ctx.followup.send(file=discord.File("/tmp/myvoteout.png")) - except UnidentifiedImageError: - logging.error("Could not download image when requested.") - await ctx.followup.send("Error opening that image!") - @app_commands.command(name="snap") async def snap(self, ctx: discord.Interaction): """ diff --git a/admin_bot/admin_bot/main.py b/admin_bot/admin_bot/main.py index ed7ee5e..c43f850 100644 --- a/admin_bot/admin_bot/main.py +++ b/admin_bot/admin_bot/main.py @@ -61,6 +61,9 @@ async def on_ready(self): except Exception as e: logging.fatal(f"Error loading {filename} as a cog, error: {e}") + # Sync tree after all changes + await self.bot.tree.sync() + async def on_message(self, ctx): # hehe, sneaky every time await self.bot.rick(ctx) diff --git a/admin_bot/admin_bot/resources/what_is_keegan_looking_at/fore.png b/admin_bot/admin_bot/resources/what_is_keegan_looking_at/fore.png new file mode 100644 index 0000000..0c3ca0a Binary files /dev/null and b/admin_bot/admin_bot/resources/what_is_keegan_looking_at/fore.png differ diff --git a/admin_bot/admin_bot/utilities/image_tasks.py b/admin_bot/admin_bot/utilities/image_tasks.py new file mode 100644 index 0000000..8b0b1a6 --- /dev/null +++ b/admin_bot/admin_bot/utilities/image_tasks.py @@ -0,0 +1,48 @@ +from PIL import Image +from utilities.image_utils import register_image_task + + +@register_image_task("myvote") +def myvote_task(input_image: Image.Image) -> Image.Image: + frame = Image.open("admin_bot/resources/vote.png") + input_image = input_image.resize((375, 250)) + frame.paste(input_image, (40, 240)) + return frame + + +@register_image_task("whiteboard") +def whiteboard_task(input_image: Image.Image) -> Image.Image: + background = Image.open("admin_bot/resources/look_at_this/background.png") + foreground = Image.open("admin_bot/resources/look_at_this/foreground.png") + input_image = input_image.resize((1000, 2000)) + background.paste(input_image, (1500, 75)) + background.paste(foreground, (0, 0), foreground) + return background + + +@register_image_task("keegan") +def keegan_task(input_image: Image.Image) -> Image.Image: + """ + Keegan transformations lol! + """ + # Load the foreground overlay (ensure it has an alpha channel) + foreground = Image.open( + "admin_bot/resources/what_is_keegan_looking_at/fore.png" + ).convert("RGBA") + + # Create a new blank RGBA image with the same size as the foreground + result_image = Image.new("RGBA", foreground.size, (0, 0, 0, 0)) # Fully transparent + + # Resize image to fit screen + resized_image = input_image.resize((1550, 900)) + + # Rotate image counterclockwise by 38.2 degrees + rotated_image = resized_image.rotate(38.2, expand=True) + + # Paste the rotated image onto the result image + result_image.paste(rotated_image, (240, 780)) + + # Place the foreground on top + result_image.paste(foreground, (0, 0), foreground) + + return result_image diff --git a/admin_bot/admin_bot/utilities/image_utils.py b/admin_bot/admin_bot/utilities/image_utils.py new file mode 100644 index 0000000..3f88d1e --- /dev/null +++ b/admin_bot/admin_bot/utilities/image_utils.py @@ -0,0 +1,62 @@ +import io +import logging + +from PIL import Image, UnidentifiedImageError +from typing import Callable, Dict + +# Registry for image manipulation tasks +IMAGE_TASKS: Dict[str, Callable[[Image.Image], Image.Image]] = {} + + +def register_image_task(name: str): + """ + Decorator to register an image task. + + :param name: The name of the task. + """ + + def decorator(func: Callable[[Image.Image], Image.Image]): + IMAGE_TASKS[name] = func + return func + + return decorator + + +def apply_image_task(image: Image.Image, task_name: str) -> Image.Image: + """ + Apply a registered image task to the given image. + + :param image: The input image. + :param task_name: The task to apply. + :return: The resulting image. + """ + if task_name not in IMAGE_TASKS: + raise ValueError(f"Task '{task_name}' is not registered.") + + return IMAGE_TASKS[task_name](image) + + +def load_image_from_bytes(image_bytes: bytes) -> Image.Image: + """ + Load an PIL image from bytes. + + :param image_bytes: The image data in bytes. + :return: A PIL Image. + """ + try: + return Image.open(io.BytesIO(image_bytes)) + except UnidentifiedImageError as e: + logging.error(f"Failed to load image: {e}") + raise + + +def save_image_to_bytes(image: Image.Image) -> bytes: + """ + Save a PIL Image to bytes. + + :param image: The PIL Image. + :return: Image data as bytes. + """ + with io.BytesIO() as output: + image.save(output, format="PNG") + return output.getvalue() diff --git a/admin_interface/Dockerfile b/admin_interface/Dockerfile index 3f69790..0f255ba 100644 --- a/admin_interface/Dockerfile +++ b/admin_interface/Dockerfile @@ -43,5 +43,7 @@ USER ${USER_ID} EXPOSE 8000 +HEALTHCHECK NONE + # Run the entrypoint bin ENTRYPOINT ["entrypoint"]