From 98d533910ce6a3049a35e9d986e513e92f35c348 Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 1 Feb 2024 23:54:54 +0100 Subject: [PATCH] [CustomRoles] guild setup for privileged role, fix deletion loop --- cogs/CustomRoles.py | 88 +++++++++++++++++++++++++++++++------------ db.py | 10 +++++ models/__init__.py | 3 +- models/custom_role.py | 8 ++++ views/__init__.py | 4 +- views/role_creator.py | 83 +++++++++++++++++++++++++++++++++++----- 6 files changed, 159 insertions(+), 37 deletions(-) diff --git a/cogs/CustomRoles.py b/cogs/CustomRoles.py index db8f9e5..1a1e011 100644 --- a/cogs/CustomRoles.py +++ b/cogs/CustomRoles.py @@ -7,27 +7,30 @@ import db from cogs import AinitMixin, CustomCog - from models import CustomRole -from views import RoleCreatorView, RoleCreatorResult +from views import RoleCreatorView, RoleCreatorResult, CustomRoleSetup logger = logging.getLogger(__name__) -class NotBoosting(commands.CheckFailure): +class MissingRole(commands.CheckFailure): def __init__(self): - super().__init__("You are not a server booster.") - + super().__init__("You are missing a role to do this.") -def is_booster(): - def predicate(ctx: commands.Context): - if ( - ctx.author.premium_since is None - and not ctx.author.guild_permissions.administrator - ): - raise NotBoosting - return True +def has_guild_role(): + async def predicate(ctx: commands.Context): + async with ctx.bot.Session() as session: + custom_role_settings = await db.get_custom_role_settings( + session, ctx.guild.id + ) + if ctx.author.guild_permissions.administrator or ( + custom_role_settings is not None + and custom_role_settings._role in [role.id for role in ctx.author.roles] + ): + return True + else: + raise MissingRole return commands.check(predicate) @@ -42,13 +45,32 @@ def __init__(self, bot): CustomRole.inject_bot(bot) - @commands.group(aliases=["customrole", "cr"], brief="Custom roles") + self._role_removal_loop.start() + + def cog_unload(self): + self._role_removal_loop.stop() + + @commands.group( + name="customroles", aliases=["customrole", "cr"], brief="Custom roles" + ) async def custom_roles(self, ctx: commands.Context): if not ctx.invoked_subcommand: await ctx.send_help(self.custom_roles) + @custom_roles.command(brief="Set up custom roles") + @commands.has_permissions(administrator=True) + async def setup(self, ctx: commands.Context): + embed = Embed( + title="Set up custom roles", + description="Choose the role below that members need to have in order to be allowed to create" + " custom roles in the server, then hit confirm.", + ) + embed.set_author(name=ctx.guild.name, icon_url=ctx.guild.icon.url) + + await ctx.reply(view=CustomRoleSetup(), embed=embed, ephemeral=True) + @custom_roles.command(brief="Removes your custom role") - @is_booster() + @has_guild_role() async def delete(self, ctx: commands.Context): async with self.bot.Session() as session: custom_role = await db.get_user_custom_role_in_guild( @@ -68,7 +90,7 @@ async def delete(self, ctx: commands.Context): await ctx.reply(f"Your role '{role_name}' has been deleted.") @custom_roles.command(brief="Creates a custom role for you") - @is_booster() + @has_guild_role() async def create(self, ctx: commands.Context): embed = Embed( title="Create custom role", @@ -110,13 +132,14 @@ async def _delete_custom_role(self, session, custom_role: CustomRole) -> None: try: logger.info( "removing custom role from %s (%d) in %s (%d)", + str(custom_role.member), + custom_role._user, str(custom_role.guild), custom_role._guild, - str(custom_role.member), - custom_role._member, ) await custom_role.role.delete( - reason=f"Member {custom_role.member} ({custom_role.member.id}) removed their boost" + reason=f"Member {custom_role.member} ({custom_role._user}) no longer has required role," + f" removing custom role" ) await db.delete_custom_role(session, custom_role._guild, custom_role._user) except discord.Forbidden: @@ -125,10 +148,10 @@ async def _delete_custom_role(self, session, custom_role: CustomRole) -> None: str(custom_role.guild), custom_role._guild, str(custom_role.member), - custom_role._member, + custom_role._user, ) - @tasks.loop(hours=24.0) + @tasks.loop(minutes=1) async def _role_removal_loop(self) -> None: logger.info("running role removal task") async with self.bot.Session() as session: @@ -137,10 +160,27 @@ async def _role_removal_loop(self) -> None: role_deletion_tasks = [] for custom_role in custom_roles: - if custom_role.member.premium_since is None: - # member is no longer boosting, add to deletion list + # TODO fix this N+1 issue + custom_role_settings = await db.get_custom_role_settings( + session, custom_role._guild + ) + + has_guild_role = custom_role_settings._role in ( + member_role.id for member_role in custom_role.member.roles + ) + + if ( + not custom_role.member.guild_permissions.administrator + and not has_guild_role + ): + # member no longer has the guild's required role, add to deletion list role_deletion_tasks.append( self._delete_custom_role(session, custom_role) ) - await asyncio.gather(*custom_roles) + await asyncio.gather(*role_deletion_tasks) + await session.commit() + + @_role_removal_loop.before_loop + async def loop_before(self): + await self.bot.wait_until_ready() diff --git a/db.py b/db.py index ec702a0..2a43af3 100644 --- a/db.py +++ b/db.py @@ -27,6 +27,7 @@ RoleSettings, GuildCog, CustomRole, + CustomRoleSettings, ) @@ -381,6 +382,15 @@ async def get_custom_roles(session: AsyncSession) -> list[CustomRole]: return [r for (r,) in result] +async def get_custom_role_settings( + session: AsyncSession, guild_id: int +) -> typing.Optional[CustomRoleSettings]: + statement = select(CustomRoleSettings).where(CustomRoleSettings._guild == guild_id) + result = (await session.execute(statement)).first() + + return result[0] if result else None + + async def get_greeter(session, guild_id, greeter_type): statement = select(Greeter).where( (Greeter._guild == guild_id) & (Greeter.type == greeter_type) diff --git a/models/__init__.py b/models/__init__.py index 9e91a1e..1d8f8e8 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1,6 +1,6 @@ from .botw import BotwState, BotwWinner, Nomination, Idol, BotwSettings from .channel_mirror import ChannelMirror -from .custom_role import CustomRole +from .custom_role import CustomRole, CustomRoleSettings from .greeter import Greeter, GreeterType from .guild_settings import GuildSettings, EmojiSettings, GuildCog from .log import CommandLog @@ -15,6 +15,7 @@ "BotwWinner", "ChannelMirror", "CustomRole", + "CustomRoleSettings", "GuildSettings", "GuildCog", "EmojiSettings", diff --git a/models/custom_role.py b/models/custom_role.py index d03345c..08eaa48 100644 --- a/models/custom_role.py +++ b/models/custom_role.py @@ -2,6 +2,7 @@ from sqlalchemy.ext.hybrid import hybrid_property from models.base import Base +from models.guild_settings import GuildSettingsMixin from models.role import RoleMixin @@ -18,3 +19,10 @@ class CustomRole(RoleMixin, Base): @hybrid_property def member(self) -> Member: return self.guild.get_member(self._user) + + +class CustomRoleSettings(GuildSettingsMixin, Base): + __tablename__ = "custom_role_settings" + + """The role that gets to create custom roles/is the insertion point""" + _role = Column(BigInteger, nullable=False) diff --git a/views/__init__.py b/views/__init__.py index 0506b07..2b2f655 100644 --- a/views/__init__.py +++ b/views/__init__.py @@ -1,4 +1,4 @@ -from .role_creator import RoleCreatorView, RoleCreatorResult +from .role_creator import RoleCreatorView, RoleCreatorResult, CustomRoleSetup -__all__ = ("RoleCreatorView", "RoleCreatorResult") +__all__ = ("RoleCreatorView", "RoleCreatorResult", "CustomRoleSetup") diff --git a/views/role_creator.py b/views/role_creator.py index 54ae322..fc34cf7 100644 --- a/views/role_creator.py +++ b/views/role_creator.py @@ -2,10 +2,12 @@ import discord from discord import Interaction, Button -from discord.ui import View, TextInput, Modal +from discord.ui import View, TextInput, Modal, RoleSelect import db +from models import CustomRoleSettings + class RoleCreatorResult: name: str | None @@ -28,6 +30,64 @@ def __init__(self, callback, result): self.result = result +class RolePicker(RoleSelect): + def __init__(self): + super().__init__(placeholder="Choose one role...") + + async def callback(self, interaction: Interaction) -> typing.Any: + role = self.values[0] + client_member = interaction.guild.get_member(interaction.client.user.id) + if role.position >= client_member.top_role.position: + await interaction.response.send_message( + "Please choose a role the bot can manage, i.e., one that is below its own highest role.", + ephemeral=True, + ) + + await interaction.response.defer() + + +class CustomRoleSetup(View): + def __init__(self): + super().__init__() + + self.add_item(RolePicker()) + + async def interaction_check(self, interaction: Interaction, /) -> bool: + if not interaction.user.guild_permissions.administrator: + await interaction.response.send_message( + "You are not allowed to set up custom roles.", ephemeral=True + ) + return False + + return True + + @discord.ui.button(label="Confirm", style=discord.ButtonStyle.blurple, row=2) + async def confirm(self, interaction: Interaction, button: Button): + self.stop() + + values = self.children[1].values + + if len(values) != 1: + await interaction.response.send_message( + "Please try again and choose exactly one role!", + ephemeral=True, + ) + return + + role = values[0] + + async with interaction.client.Session() as session: + custom_role_settings = CustomRoleSettings( + _role=role.id, _guild=interaction.guild_id + ) + await session.merge(custom_role_settings) + await session.commit() + + await interaction.response.send_message( + f"Members with the role {role.mention} will now be able to create custom roles!", + ) + + class RoleCreatorView(CallbackView, View): @discord.ui.button(label="Click me to start", style=discord.ButtonStyle.blurple) async def create_role(self, interaction: Interaction, button: Button): @@ -48,20 +108,23 @@ async def interaction_check(self, interaction: Interaction, /) -> bool: ) return False - # FIXME remove me! - return True + custom_role_settings = await db.get_custom_role_settings( + session, interaction.guild_id + ) + member = interaction.guild.get_member(interaction.user.id) + + if member.guild_permissions.administrator or ( + custom_role_settings is not None + and member is not None + and custom_role_settings._role in [role.id for role in member.roles] + ): + return True - if ( - interaction.user.premium_since is None - and not interaction.user.guild_permissions.administrator - ): await interaction.response.send_message( - "You are not a server booster.", ephemeral=True + "You are missing a role to do this.", ephemeral=True ) return False - return True - class RoleCreatorNameConfirmationView(CallbackView, View): @discord.ui.button(label="I confirm", style=discord.ButtonStyle.blurple)