1 files changed, 101 insertions, 0 deletions
diff --git a/src/cogs/archive.py b/src/cogs/archive.py
new file mode 100644
index 0000000..36c7044
--- /dev/null
+++ b/src/cogs/archive.py
@@ -0,0 +1,101 @@
+import discord
+from discord import app_commands
+from discord.ext import commands
+import aiohttp
+import random
+import string
+import os
+
+from src.utils.db import get_db
+from models import Message
+from config import BOT_COLOR
+
+
+class Archive(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @app_commands.command()
+    async def archive(
+        self,
+        interaction: discord.Interaction,
+        channel: discord.TextChannel,
+        amount: int,
+    ):
+        """Archive a channel's messages."""
+        if not channel.permissions_for(
+            interaction.guild.me
+        ).read_message_history:
+            return await interaction.response.send_message(
+                "I do not have permission to read message history in that"
+                " channel.",
+                ephemeral=True,
+            )
+
+        if amount < 1:
+            return await interaction.response.send_message(
+                "You must provide a number greater than 0.",
+                ephemeral=True,
+            )
+
+        await interaction.response.send_message("Archiving messages now.")
+
+        db = next(get_db())
+        count = 0
+        messages = channel.history(limit=amount)
+        async for message in messages:
+            count += 1
+
+            author = message.author
+            stickers = [sticker.name for sticker in message.stickers]
+            role_mentions = [
+                role_mention.id for role_mention in message.role_mentions
+            ]
+            mention_everyone = message.mention_everyone
+            mentions = [mention.id for mention in message.mentions]
+            attachments = [
+                attachment.url for attachment in message.attachments
+            ]
+
+            if not os.path.exists("images"):
+                os.makedirs("images")
+
+            # Download all images before saving everything to database
+            for url in attachments:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(url) as response:
+                        if response.status == 200:
+                            # Create a randomized filename
+                            filename = "".join(
+                                random.choice(string.ascii_letters)
+                                for i in range(10)
+                            )
+
+                            # Save the image to the filesystem
+                            with open(f"images/{filename}.jpg", "wb") as file:
+                                file.write(await response.read())
+
+                            # Update the attachment URL to the new filename
+                            attachments[attachments.index(url)] = (
+                                f"images/{filename}.jpg"
+                            )
+
+            content = message.content
+
+            db_message = Message(
+                author_id=author.id,
+                channel_id=channel.id,
+                stickers=stickers,
+                role_mentions=role_mentions,
+                mention_everyone=mention_everyone,
+                mentions=mentions,
+                attachments=attachments,
+                content=content,
+            )
+
+            db.add(db_message)
+            db.commit()
+
+
+async def setup(bot):
+    await bot.add_cog(Archive(bot))