diff --git a/Heimdallr.py b/Heimdallr.py index 930f291..144f5c1 100644 --- a/Heimdallr.py +++ b/Heimdallr.py @@ -69,4 +69,5 @@ if __name__ == "__main__": load_dotenv() bot.load_extension("commands.quote") bot.load_extension("commands.infractions") + bot.load_extension("commands.self_roles") bot.start(getenv("DISCORD_TOKEN")) diff --git a/commands/infractions.py b/commands/infractions.py index ebf4f85..6aba85f 100644 --- a/commands/infractions.py +++ b/commands/infractions.py @@ -1,5 +1,5 @@ import logging -from typing import List +from typing import List, Optional from naff import ( Extension, slash_command, @@ -16,6 +16,7 @@ from naff import ( ButtonStyles, ) from database import ( + GuildSettings, Infractions as InfractionsModel, ) @@ -181,6 +182,22 @@ class Infractions(Extension): f'Warned {user.display_name} ({user.username}#{user.discriminator}, {user.id}) for "{reason}" with severity {weight}', # pylint: disable=line-too-long ephemeral=True, ) + + guild_settings: Optional[GuildSettings] = GuildSettings.get_or_none(GuildSettings.guild_id == int(ctx.guild_id)) + if guild_settings is not None: + if guild_settings.admin_channel is not None: + admin_channel = self.client.fetch_channel(int(guild_settings.admin_channel)) + if admin_channel is not None: + await admin_channel.send(embed=Embed( + title=f"Warned {user.display_name} ({user.username}#{user.discriminator}, {user.id})", + description=f"{reason}", + color=infraction_colour(0x0000FF), + fields=[ + EmbedField(name="**Severity**", value=f"{weight}"), + EmbedField("**Issued by**", f"{ctx.author.display_name}"), + ], + )) + if not silent and warning_msg is None: await ctx.send( f"{user.mention} has been warned, but I couldn't DM them.", diff --git a/commands/self_roles.py b/commands/self_roles.py new file mode 100644 index 0000000..6361a64 --- /dev/null +++ b/commands/self_roles.py @@ -0,0 +1,303 @@ +from re import S +from typing import List, Optional +import logging +from naff import ( + Extension, + Client, + slash_command, + slash_option, + OptionTypes, + Permissions, + InteractionContext, + Role, + Embed, + EmbedField, + AutocompleteContext, +) +from database import SelfRoles as SelfRolesModel + + +class SelfRoles(Extension): + def __init__(self, client: Client): + self.client: Client = client + + @slash_command( + name="role-admin", + description="Manage self-roles", + sub_cmd_name="add", + sub_cmd_description="Add a self-role", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_ROLES, + ) + @slash_option( + name="role", description="Role to add", required=True, opt_type=OptionTypes.ROLE + ) + @slash_option( + name="description", + description="Description of the role", + required=False, + opt_type=OptionTypes.STRING, + ) + @slash_option( + name="requires", + description="Role required to add this role", + required=False, + opt_type=OptionTypes.ROLE, + ) + async def role_admin_add( + self, + ctx: InteractionContext, + role: Role, + description: str = None, + requires: Role = None, + ): + await ctx.defer(ephemeral=True) + + SelfRolesModel.insert( + guild_id=int(ctx.guild_id), + role_id=int(role.id), + role_name=role.name, + role_description=description, + requires=int(requires.id) if requires is not None else None, + ).on_conflict_replace().execute() + + await ctx.send( + f"Added self-role {role.mention} to the database." + f"\n{'Description: ' + description if description is not None else ''}" + f"\n{'Requires: ' + requires.mention if requires is not None else ''}", + ephemeral=True, + ) + + @slash_command( + name="role-admin", + description="Manage self-roles", + sub_cmd_name="remove", + sub_cmd_description="Remove a self-role", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_ROLES, + ) + @slash_option( + name="role", + description="Role to remove", + required=True, + opt_type=OptionTypes.ROLE, + ) + async def role_admin_remove(self, ctx: InteractionContext, role: Role): + await ctx.defer(ephemeral=True) + + SelfRolesModel.delete().where( + SelfRolesModel.guild_id == int(ctx.guild_id), + SelfRolesModel.role_id == int(role.id), + ).execute() + + await ctx.send( + f"Removed self-role {role.mention} from the database.", ephemeral=True + ) + + @slash_command( + name="role", + description="Manage your roles", + dm_permission=False, + sub_cmd_name="add", + sub_cmd_description="Add a role", + ) + @slash_option( + name="role", + description="Role to add", + required=True, + opt_type=OptionTypes.STRING, + autocomplete=True, + ) + async def role_add(self, ctx: InteractionContext, role: str): + await ctx.defer(ephemeral=True) + + role_q: Optional[SelfRolesModel] = SelfRolesModel.get_or_none( + SelfRolesModel.guild_id == int(ctx.guild_id), + SelfRolesModel.role_id == int(role), + ) + + if role_q is None: + await ctx.send(f"Failed to retrieve self-role.", ephemeral=True) + return + + role = await ctx.guild.fetch_role(role) + + if role_q.requires is not None: + required_role = await ctx.guild.fetch_role(role_q.requires) + if required_role is None: + await ctx.send( + "Getting role failed. This role requires another role but the bot could not figure out which.", + ephemeral=True, + ) + return + + if required_role not in ctx.author.roles: + await ctx.send( + f"You do not have the required role {required_role.mention} to add this role.", + ephemeral=True, + ) + return + + await ctx.author.add_role(role, reason="Self-role added") + + await ctx.send( + f"Added role {role.mention} to {ctx.author.mention}.", ephemeral=True + ) + + @role_add.autocomplete("role") + async def role_add_autocomplete_role( + self, ctx: AutocompleteContext, role: str = None + ): + choices = [] + + if role is None: + roles: List[SelfRolesModel] = ( + SelfRolesModel.select() + .where( + SelfRolesModel.guild_id == int(ctx.guild_id), + ~(SelfRolesModel.guild_id.in_(ctx.author.roles)), + ) + .execute() + ) + choices = [ + {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"} + for r in roles + ] + await ctx.send(choices=choices) + return + + roles: List[SelfRolesModel] = ( + SelfRolesModel.select() + .where( + SelfRolesModel.guild_id == int(ctx.guild_id), + SelfRolesModel.role_name.startswith(role), + ~(SelfRolesModel.role_id.in_(ctx.author.roles)), + ) + .execute() + ) + + for r in roles: + choices.append( + {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"} + ) + + await ctx.send(choices=choices) + + @slash_command( + name="role", + description="Manage your roles", + dm_permission=False, + sub_cmd_name="remove", + sub_cmd_description="Remove a role", + ) + @slash_option( + name="role", + description="Role to remove", + required=True, + opt_type=OptionTypes.STRING, + autocomplete=True, + ) + async def role_remove(self, ctx: InteractionContext, role: str): + await ctx.defer(ephemeral=True) + role = await ctx.guild.fetch_role(role) + if role is None: + await ctx.send(f"Failed to retrieve role.", ephemeral=True) + return + + if role not in ctx.author.roles: + await ctx.send( + f"{ctx.author.mention} does not have the role {role.mention}.", + ephemeral=True, + ) + return + + role_q: Optional[SelfRolesModel] = SelfRolesModel.get_or_none( + SelfRolesModel.guild_id == int(ctx.guild_id), + SelfRolesModel.role_id == int(role.id), + ) + if role_q is None: + await ctx.send(f"{role.mention} is not a self-role.", ephemeral=True) + return + + await ctx.author.remove_role(role, reason="Self-role removed") + await ctx.send( + f"Removed role {role.mention} from {ctx.author.mention}.", ephemeral=True + ) + + @role_remove.autocomplete("role") + async def role_remove_autocomplete_role( + self, ctx: AutocompleteContext, role: str = None + ): + choices = [] + + if role is None: + roles: List[SelfRolesModel] = ( + SelfRolesModel.select() + .where( + SelfRolesModel.guild_id == int(ctx.guild_id), + (SelfRolesModel.guild_id.in_(ctx.author.roles)), + ) + .execute() + ) + choices = [ + {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"} + for r in roles + ] + await ctx.send(choices=choices) + return + + roles: List[SelfRolesModel] = ( + SelfRolesModel.select() + .where( + SelfRolesModel.guild_id == int(ctx.guild_id), + SelfRolesModel.role_name.startswith(role), + (SelfRolesModel.role_id.in_(ctx.author.roles)), + ) + .execute() + ) + + for r in roles: + choices.append( + {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"} + ) + + await ctx.send(choices=choices) + + @slash_command( + name="role", + description="Manage your roles", + dm_permission=False, + sub_cmd_name="list", + sub_cmd_description="List all self-roles", + ) + async def role_list(self, ctx: InteractionContext): + await ctx.defer(ephemeral=True) + + roles: List[SelfRolesModel] = SelfRolesModel.select().where( + SelfRolesModel.guild_id == int(ctx.guild_id) + ) + if roles.count() == 0: + await ctx.send("No self-roles found.", ephemeral=True) + return + + embeds: List[Embed] = [] + + for role in roles: + r = await ctx.guild.fetch_role(role.role_id) + embed = Embed( + title=f"Self-role {r.name}", + description=role.role_description, + color=r.color, + ) + if role.requires is not None: + req_role = await ctx.guild.fetch_role(role.requires) + embed.add_field(name="Requires", value=req_role.mention, inline=False) + embeds.append(embed) + + await ctx.send(embeds=embeds) + + +def setup(client: Client): + SelfRoles(client) + SelfRolesModel.create_table() + logging.info("SelfRoles extension loaded.") diff --git a/database.py b/database.py index 00df1e1..e71bf88 100644 --- a/database.py +++ b/database.py @@ -62,8 +62,8 @@ class SelfRoles(Model): guild_id = BigIntegerField() role_id = BigIntegerField() role_name = TextField() - role_description = TextField() - requires = BigIntegerField() + role_description = TextField(null=True) + requires = BigIntegerField(null=True) class Meta: primary_key = CompositeKey("guild_id", "role_id")