diff --git a/bot/config.py b/bot/config.py index ac390a14..989652cb 100644 --- a/bot/config.py +++ b/bot/config.py @@ -95,6 +95,11 @@ class ErrorHandling(BaseModel): webhook_url: str +class CustomRoles(BaseModel): + log_channel_id: int + divider_role_id: int + + class Settings(BaseSettings): aoc: AoC bot: Bot @@ -109,6 +114,7 @@ class Settings(BaseSettings): timathon: Timathon hastebin: Hastebin errors: ErrorHandling + custom_roles: CustomRoles class Config: env_file = ".env" diff --git a/bot/extensions/custom_roles/__init__.py b/bot/extensions/custom_roles/__init__.py new file mode 100644 index 00000000..1239ad49 --- /dev/null +++ b/bot/extensions/custom_roles/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .commands import CustomRoles +from .events import CustomRoleEvents + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(CustomRoles(bot=bot)) + await bot.add_cog(CustomRoleEvents(bot=bot)) diff --git a/bot/extensions/custom_roles/commands.py b/bot/extensions/custom_roles/commands.py new file mode 100644 index 00000000..e6b6abab --- /dev/null +++ b/bot/extensions/custom_roles/commands.py @@ -0,0 +1,111 @@ +import datetime + +import discord +from discord import app_commands, utils +from discord.ext import commands + +from bot import core +from bot.config import settings +from bot.models import CustomRole + + +class CustomRoles(commands.Cog): + def __init__(self, bot): + self.bot = bot + + self.color_converter = commands.ColorConverter() + + @staticmethod + def role_embed(heading: str, user: discord.Member, role: discord.Role): + embed = discord.Embed( + title=heading, + color=role.color, + timestamp=datetime.datetime.utcnow(), + ) + embed.add_field(name="Name", value=utils.escape_markdown(role.name)) + embed.add_field(name="Color", value=str(role.color)) + embed.add_field(name="Created at", value=utils.format_dt(role.created_at)) + embed.set_thumbnail(url=user.avatar) + return embed + + @app_commands.command() + @app_commands.default_permissions(administrator=True) + @app_commands.describe(name="New name", color="New color") + async def myrole( + self, interaction: core.InteractionType, name: app_commands.Range[str, 2, 100] | None, color: str = None + ): + """Manage your custom role""" + if color is not None: + try: + color = await self.color_converter.convert(None, color) # noqa + except commands.BadColourArgument as e: + return await interaction.response.send_message(str(e), ephemeral=True) + + query = "SELECT * FROM custom_roles WHERE guild_id = $1 AND user_id = $2" + before = await CustomRole.fetchrow(query, interaction.guild.id, interaction.user.id) + + if before is None: + if name is None: + return await interaction.response.send_message("You don't have a custom role yet!", ephemeral=True) + + # Create and assign the role to user + role = await interaction.guild.create_role(name=name, colour=color or discord.Color.random()) + + divider_role = interaction.guild.get_role(settings.custom_roles.divider_role_id) + await role.edit(position=divider_role.position + 1) + + record = await CustomRole.ensure_exists( + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_id=role.id, + name=role.name, + color=role.color.value, + ) + + self.bot.dispatch("custom_role_create", custom_role=record) + self.bot.dispatch( + "persist_roles", + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_ids=[role.id], + ) + + return await interaction.response.send_message( + embed=self.role_embed("**Custom Role has been assigned**", interaction.user, role), + ephemeral=True, + ) + + role = interaction.guild.get_role(before.role_id) + + # Return role information if no parameter is passed + if (name == before.name or name is None) and (color.value == before.color or color is None): + return await interaction.response.send_message( + embed=self.role_embed("Custom Role for", interaction.user, interaction.guild.get_role(before.role_id)), + ephemeral=True, + ) + + await role.edit( + name=name or before.name, + colour=color or discord.Color(int(before.color)), + ) + + after = await CustomRole.ensure_exists( + guild_id=interaction.guild.id, + user_id=interaction.user.id, + role_id=role.id, + name=role.name, + color=role.color.value, + ) + + self.bot.dispatch("custom_role_update", before, after) + + return await interaction.response.send_message( + embed=self.role_embed( + "**Custom Role has been updated**", interaction.user, interaction.guild.get_role(before.role_id) + ), + ephemeral=True, + ) + + +async def setup(bot: commands.Bot): + await bot.add_cog(CustomRoles(bot=bot)) diff --git a/bot/extensions/custom_roles/events.py b/bot/extensions/custom_roles/events.py new file mode 100644 index 00000000..aaad6302 --- /dev/null +++ b/bot/extensions/custom_roles/events.py @@ -0,0 +1,80 @@ +import datetime +import time + +import discord +from discord.ext import commands + +from bot import core +from bot.config import settings +from bot.models.custom_roles import CustomRole + + +class CustomRoleEvents(commands.Cog): + """Events for Levelling in discord.""" + + def __init__(self, bot: core.DiscordBot): + self.bot = bot + self.updated_at: dict[int, datetime] = {} + + @property + def custom_roles_logs_channel(self) -> discord.TextChannel | None: + return self.bot.guild.get_channel(settings.custom_roles.log_channel_id) + + @commands.Cog.listener() + async def on_guild_role_update(self, before: discord.Role, after: discord.Role): + last_update = self.updated_at.get(before.id) + + # Ignore events less than 10 seconds since a user updated their role + if last_update is not None: + if time.time() - last_update < 10: + return + + query = """ + UPDATE custom_roles + SET name = $2, color = $3 + WHERE role_id = $1 + """ + await CustomRole.execute(query, after.id, after.name, after.color.value) + + @commands.Cog.listener() + async def on_custom_role_create(self, custom_role: CustomRole): + """Logs the creation of new role""" + self.updated_at[custom_role.role_id] = time.time() + + user = await self.bot.fetch_user(custom_role.user_id) + + embed = discord.Embed( + title="Custom Role Created", + color=discord.Color.brand_green(), + timestamp=datetime.datetime.utcnow(), + ) + embed.set_author(name=user.display_name, icon_url=user.display_avatar) + embed.add_field(name="Name", value=custom_role.name) + embed.add_field(name="Color", value="#" + hex(custom_role.color)[2:]) + embed.set_thumbnail(url=user.avatar) + embed.set_footer(text=f"user_id: {custom_role.user_id}") + + return await self.custom_roles_logs_channel.send(embed=embed) + + @commands.Cog.listener() + async def on_custom_role_update(self, before: CustomRole, after: CustomRole): + """Logs the update of custom role.""" + self.updated_at[after.role_id] = time.time() + + user = await self.bot.fetch_user(after.user_id) + + embed = discord.Embed( + title="Custom Role Updated", + color=discord.Color.brand_green(), + timestamp=datetime.datetime.utcnow(), + ) + + embed.add_field(name="Old Name", value=before.name) + embed.add_field(name="New Name", value=after.name) + embed.add_field(name="\u200B", value="\u200B") + embed.add_field(name="Old Color", value="#" + hex(before.color)[2:]) + embed.add_field(name="New Color", value="#" + hex(after.color)[2:]) + embed.set_thumbnail(url=user.avatar) + embed.set_footer(text=f"user_id: {after.user_id}") + + return await self.custom_roles_logs_channel.send(embed=embed) diff --git a/bot/models/__init__.py b/bot/models/__init__.py index 3fb6422e..2bb857fe 100644 --- a/bot/models/__init__.py +++ b/bot/models/__init__.py @@ -1,3 +1,4 @@ +from .custom_roles import CustomRole from .gconfig import FilterConfig from .levelling_ignored_channels import IgnoredChannel from .levelling_roles import LevellingRole @@ -20,4 +21,5 @@ PersistedRole, IgnoredChannel, LevellingRole, + CustomRole, ) # Fixes F401 diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py index 73c90fba..35736fae 100644 --- a/bot/models/custom_roles.py +++ b/bot/models/custom_roles.py @@ -3,17 +3,18 @@ class CustomRole(Model): id: int + user_id: int | None guild_id: int role_id: int name: str - color: str + color: int @classmethod - async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: str): + async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: int, user_id: int = None): """Inserts or updates the custom role.""" query = """ - INSERT INTO custom_roles (guild_id, role_id, name, color) - VALUES ($1, $2, $3, $4) + INSERT INTO custom_roles (guild_id, role_id, name, color, user_id) + VALUES ($1, $2, $3, $4, $5) ON CONFLICT (guild_id, role_id) DO UPDATE SET name = $3, @@ -21,4 +22,4 @@ async def ensure_exists(cls, guild_id: int, role_id: int, name: str, color: str) RETURNING * """ - return await cls.fetchrow(query, guild_id, role_id, name, color) + return await cls.fetchrow(query, guild_id, role_id, name, color, user_id) diff --git a/bot/models/migrations/005_down__custom_role.sql b/bot/models/migrations/005_down__custom_role.sql new file mode 100644 index 00000000..e03e6cd7 --- /dev/null +++ b/bot/models/migrations/005_down__custom_role.sql @@ -0,0 +1,8 @@ +ALTER TABLE custom_roles DROP COLUMN user_id; + +BEGIN; +ALTER TABLE custom_roles ADD COLUMN old_color VARCHAR NOT NULL; +UPDATE custom_roles SET old_color = CAST(color AS VARCHAR); +ALTER TABLE custom_roles DROP COLUMN color; +ALTER TABLE custom_roles RENAME COLUMN old_color TO color; +COMMIT; diff --git a/bot/models/migrations/005_up__custom_roles.sql b/bot/models/migrations/005_up__custom_roles.sql new file mode 100644 index 00000000..80d33a36 --- /dev/null +++ b/bot/models/migrations/005_up__custom_roles.sql @@ -0,0 +1,11 @@ +ALTER TABLE custom_roles ADD COLUMN user_id BIGINT; + + +-- Create a temporary column `new_color` and update the table with the values from this table. +-- Drop the old color column, then rename the new one and set it to non-nullable. +BEGIN; +ALTER TABLE custom_roles ADD COLUMN new_color INTEGER NOT NULL; +UPDATE custom_roles SET new_color = CAST(color AS INTEGER); +ALTER TABLE custom_roles DROP COLUMN color; +ALTER TABLE custom_roles RENAME COLUMN new_color TO color; +COMMIT; diff --git a/cli.py b/cli.py index 571676d0..b151094b 100644 --- a/cli.py +++ b/cli.py @@ -137,6 +137,7 @@ async def main(ctx): "bot.extensions.tags", "bot.extensions.levelling", "bot.extensions.persistent_roles", + "bot.extensions.custom_roles", "bot.extensions.polls", "bot.extensions.youtube", "bot.cogs._help", diff --git a/example.env b/example.env index bd049c39..efc11cf8 100644 --- a/example.env +++ b/example.env @@ -63,3 +63,7 @@ TAGS__REQUIRED_ROLE_ID=0 # --- Timathon TIMATHON__CHANNEL_ID=0 TIMATHON__PARTICIPANT_ROLE_ID=0 + +# --- Custom Roles +CUSTOM_ROLES__LOG_CHANNEL_ID = 0 +CUSTOM_ROLES__DIVIDER_ROLE_ID = 0