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, LOG class Archive(commands.Cog): def __init__(self, bot): self.bot = bot async def download_attachments(attachments) -> list: """Download attachments and return a list of their paths.""" paths = [] for attachment in attachments: async with aiohttp.ClientSession().get(attachment.url) as response: content_type = response.headers.get("Content-Type") if not content_type: LOG.warn( f"Failed to get content type for: {attachment.url}" ) continue # Create a randomized filename file_extension = content_type.split("/")[-1] filename = ( "".join( random.choice(string.ascii_letters) for i in range(10) ) + f".{file_extension}" ) # Save the image to the filesystem with open(f"images/{filename}]", "wb") as file: file.write(await response.read()) # Add the path to the attachments list paths.append(f"images/{filename}") return paths @app_commands.command() async def archive( self, interaction: discord.Interaction, channel: discord.TextChannel, amount: int, ): """Archive a channel's messages.""" # Ensure valid channel permissions 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, ) # Ensure valid amount 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 attachments = [] if not os.path.exists("images"): os.makedirs("images") attachments = await self.download_attachments(message.attachments) db_message = Message( author_id=message.author.id, channel_id=channel.id, stickers=[sticker.name for sticker in message.stickers], role_mentions=[role.id for role in message.role_mentions], mention_everyone=message.mention_everyone, mentions=[mention.id for mention in message.mentions], attachments=attachments, content=message.content, ) db.add(db_message) db.commit() async def setup(bot): await bot.add_cog(Archive(bot))