diff --git a/techsupport_bot/commands/archive.py b/techsupport_bot/commands/archive.py new file mode 100644 index 00000000..834ddd98 --- /dev/null +++ b/techsupport_bot/commands/archive.py @@ -0,0 +1,154 @@ +""" +TO HANDLE: +- Reactions +- Threads +- Embeds +- Channel Info +- Discriminator (for bot accounts) +""" + +from __future__ import annotations + +import json +import zipfile +from datetime import timezone +from pathlib import Path +from typing import Self + +import bot +import discord +from core import cogs +from discord import app_commands + +ARCHIVE_DIR = Path("/archive-output") + + +async def setup(bot: bot.TechSupportBot) -> None: + await bot.add_cog(ChannelTextArchiver(bot=bot)) + + +class ChannelTextArchiver(cogs.BaseCog): + + async def write_channel_archive_zip( + self: Self, + channel: discord.TextChannel | discord.Thread, + filename: str, + ) -> Path: + ARCHIVE_DIR.mkdir(parents=True, exist_ok=True) + zip_path = ARCHIVE_DIR / filename + + authors_seen: set[int] = set() + + with zipfile.ZipFile( + zip_path, + mode="w", + compression=zipfile.ZIP_DEFLATED, + compresslevel=9, + ) as zip_file: + async for message in channel.history( + limit=None, + oldest_first=True, + ): + message_id = str(message.id) + message_folder = f"Messages/{message_id}/" + + author = message.author + username = author.name + author_id = author.id + + if author_id not in authors_seen: + authors_seen.add(author_id) + + author_data = { + "id": str(author.id), + "name": author.name, + "display_name": author.display_name, + "global_name": author.global_name, + "display_color": f"#{author.color.value:06x}", + } + + zip_file.writestr( + f"Authors/{author_id}/author.json", + json.dumps(author_data, indent=2, ensure_ascii=False), + ) + + if author.display_avatar: + avatar_bytes = await author.display_avatar.read() + zip_file.writestr( + f"Authors/{author_id}/profile-pic.png", + avatar_bytes, + ) + + reply_data = None + + if message.reference and isinstance( + message.reference.resolved, + discord.Message, + ): + ref = message.reference.resolved + reply_data = { + "reply_id": str(ref.id), + "reply_deleted": False, + "reply_content": ref.clean_content, + "reply_author": ref.author.name, + } + elif message.reference: + reply_data = { + "reply_id": str(message.reference.message_id), + "reply_deleted": True, + "reply_content": None, + "reply_author": None, + } + + message_data = { + "id": message_id, + "content": message.clean_content, + "timestamp": message.created_at.astimezone( + timezone.utc + ).isoformat(), + "author_username": username, + "attachments": [ + f"{idx}-{a.filename}" + for idx, a in enumerate(message.attachments, start=1) + ], + } + + if reply_data is not None: + message_data["reply"] = reply_data + + zip_file.writestr( + f"{message_folder}message.json", + json.dumps(message_data, indent=2, ensure_ascii=False), + ) + + attachments_folder = f"{message_folder}attachments/" + + for idx, attachment in enumerate(message.attachments, start=1): + data = await attachment.read() + safe_name = f"{idx}-{attachment.filename}" + + zip_file.writestr( + f"{attachments_folder}{safe_name}", + data, + ) + + return zip_path + + @app_commands.command( + name="archive", + description="Archive messages, authors, attachments, and replies", + ) + async def archive_text( + self: Self, + interaction: discord.Interaction, + ) -> None: + await interaction.response.defer(thinking=True) + + zip_path = await self.write_channel_archive_zip( + channel=interaction.channel, + filename=f"{interaction.channel.id}_archive.zip", + ) + + await interaction.followup.send( + f"Archive completed successfully.\nSaved to `{zip_path}`" + )