diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/cogs/archive.py | 101 | ||||
| -rw-r--r-- | src/cogs/owner/sync.py | 87 | ||||
| -rw-r--r-- | src/schemas/message_schema.py | 12 | ||||
| -rw-r--r-- | src/utils/command_tree.py | 19 | ||||
| -rw-r--r-- | src/utils/db.py | 9 |
5 files changed, 228 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)) diff --git a/src/cogs/owner/sync.py b/src/cogs/owner/sync.py new file mode 100644 index 0000000..d6647b5 --- /dev/null +++ b/src/cogs/owner/sync.py @@ -0,0 +1,87 @@ +import discord +from discord.ext import commands + + +class TreeSync(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @commands.group(invoke_without_command=True) + @commands.dm_only() + @commands.is_owner() + async def tree(self, ctx): + await ctx.author.send( + "This is a group command. Use either" + f" `{self.bot.command_prefix}tree sync` or" + f" `{self.bot.command_prefix}tree clear` followed by an optional" + " guild ID." + ) + + @commands.dm_only() + @commands.is_owner() + @tree.command() + async def sync( + self, ctx: commands.Context, *, guild: discord.Object = None + ): + """Sync the command tree to a guild or globally.""" + if guild: + self.bot.tree.copy_global_to(guild=guild) + await self.bot.tree.sync(guild=guild) + return await ctx.author.send( + "Synced the command tree to" + f" `{self.bot.get_guild(guild.id).name}`" + ) + else: + await self.bot.tree.sync() + return await ctx.author.send("Synced the command tree globally.") + + @sync.error + async def tree_sync_error(self, ctx, error): + if isinstance(error, commands.ObjectNotFound): + return await ctx.author.send( + "The guild you provided does not exist." + ) + if isinstance(error, commands.CommandInvokeError): + return await ctx.author.send( + "Guild ID provided is not a guild that the bot is in." + ) + else: + return await ctx.author.send( + "An unknown error occurred. Perhaps you've been rate limited." + ) + + @commands.dm_only() + @commands.is_owner() + @tree.command() + async def clear(self, ctx: commands.Context, *, guild: discord.Object): + """Clear the command tree from a guild.""" + self.bot.tree.clear_commands(guild=guild) + await self.bot.tree.sync(guild=guild) + return await ctx.author.send( + "Cleared the command tree from" + f" `{self.bot.get_guild(guild.id).name}`" + ) + + @clear.error + async def tree_sync_error(self, ctx, error): + if isinstance(error, commands.MissingRequiredArgument): + return await ctx.author.send( + "You need to provide a guild ID to clear the command tree" + " from." + ) + if isinstance(error, commands.ObjectNotFound): + return await ctx.author.send( + "The guild you provided does not exist." + ) + if isinstance(error, commands.CommandInvokeError): + return await ctx.author.send( + "Guild ID provided is not a guild that the bot is in." + ) + else: + return await ctx.author.send( + "An unknown error occurred. Perhaps you've been rate limited." + ) + + +async def setup(bot): + await bot.add_cog(TreeSync(bot)) diff --git a/src/schemas/message_schema.py b/src/schemas/message_schema.py new file mode 100644 index 0000000..79c9659 --- /dev/null +++ b/src/schemas/message_schema.py @@ -0,0 +1,12 @@ +from pydantic import BaseModel + + +class Message(BaseModel): + author_id: int + channel_id: int + stickers: list[str] + role_mentions: list[int] + mention_everyone: bool + mentions: list[int] + attachments: list[str] + content: str diff --git a/src/utils/command_tree.py b/src/utils/command_tree.py new file mode 100644 index 0000000..31ec2c1 --- /dev/null +++ b/src/utils/command_tree.py @@ -0,0 +1,19 @@ +import discord +from discord import app_commands +from discord.ext.commands.errors import * + +from config import create_embed + + +class Tree(app_commands.CommandTree): + async def on_error( + self, + interaction: discord.Interaction, + error: app_commands.AppCommandError, + ): + + if isinstance(error, CommandNotFound): + return + + else: + raise error diff --git a/src/utils/db.py b/src/utils/db.py new file mode 100644 index 0000000..a6734ea --- /dev/null +++ b/src/utils/db.py @@ -0,0 +1,9 @@ +from database import SessionLocal + + +def get_db(): + db = SessionLocal() + try: + yield db + finally: + db.close() |