Way cooler images (#59)

* Add healthcheck

* Move myvote (aaron) to new cog

* At a pausing point

* Kredcool?

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))

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):
         """

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)

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

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()

admin_interface/Dockerfile

@@ -43,5 +43,7 @@ USER ${USER_ID}
 
 EXPOSE 8000
 
+HEALTHCHECK NONE
+
 # Run the entrypoint bin
 ENTRYPOINT ["entrypoint"]