Merge pull request #259 from SylteA/dev

custom roles

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"

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))

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))

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)

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

bot/models/custom_roles.py

@@ -3,22 +3,23 @@
 
 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,
                   color = $4
              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)

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;

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;

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",

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