Heimdallr/heimdallr/commands/self_roles.py

831 lines
27 KiB
Python

# pylint: disable=not-an-iterable
from typing import Dict, List, Optional, Tuple
import logging
from naff import (
Extension,
Client,
slash_command,
slash_option,
OptionTypes,
Permissions,
InteractionContext,
Role,
Embed,
AutocompleteContext,
StringSelectMenu,
SelectOption,
listen,
context_menu,
CommandTypes,
Message,
)
from naff.api import events
from database import (
SelfRoles as SelfRolesModel,
SelfRoleGroups as SelfRoleGroupsModel,
SelfRoleGroupRelations as SelfRoleGroupRelationsModel,
)
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-group",
description="Manage self-role groups",
sub_cmd_name="add-group",
sub_cmd_description="Add a self-role group",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
@slash_option(
name="group",
description="Group to add",
required=True,
opt_type=OptionTypes.STRING,
)
@slash_option(
name="description",
description="Description of the group",
required=False,
opt_type=OptionTypes.STRING,
)
@slash_option(
name="exclusive",
description="Whether each role in this group is exclusive",
required=False,
opt_type=OptionTypes.BOOLEAN,
)
async def role_group_add_group(
self,
ctx: InteractionContext,
group: str,
description: str = None,
exclusive: bool = False,
):
await ctx.defer(ephemeral=True)
SelfRoleGroupsModel.insert(
guild_id=int(ctx.guild_id),
group_name=group,
group_description=description,
exclusive=exclusive,
).on_conflict_replace().execute()
await ctx.send(
f"Added self-role group {group} to the database."
f"\n{'Description: ' + description if description is not None else ''}",
ephemeral=True,
)
@slash_command(
name="role-group",
description="Manage self-role groups",
sub_cmd_name="remove-group",
sub_cmd_description="Remove a self-role group",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
@slash_option(
name="group",
description="Group to remove",
required=True,
opt_type=OptionTypes.STRING,
)
async def role_group_remove_group(self, ctx: InteractionContext, group: str):
await ctx.defer(ephemeral=True)
SelfRoleGroupRelationsModel.delete().where(
SelfRoleGroupRelationsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == group,
).execute()
SelfRoleGroupsModel.delete().where(
SelfRoleGroupsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupsModel.group_name == group,
).execute()
await ctx.send(
f"Removed self-role group {group} from the database.", ephemeral=True
)
@slash_command(
name="role-group",
description="Manage self-role groups",
sub_cmd_name="add-role",
sub_cmd_description="Add a role to a self-role group",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
@slash_option(
name="group",
description="Group to add the role to",
required=True,
opt_type=OptionTypes.STRING,
autocomplete=True,
)
@slash_option(
name="role",
description="Role to add to the group",
required=True,
opt_type=OptionTypes.STRING,
autocomplete=True,
)
async def role_group_add_role(self, ctx: InteractionContext, group: str, role: str):
await ctx.defer(ephemeral=True)
r: Role = await ctx.guild.fetch_role(role)
if r is None:
await ctx.send(f"Role {role} not found in this server.", ephemeral=True)
return
SelfRoleGroupRelationsModel.insert(
guild_id=int(ctx.guild_id),
group_name=group,
role_id=int(role),
).on_conflict_replace().execute()
await ctx.send(f"Added role {r.mention} to the group {group}.", ephemeral=True)
@role_group_add_role.autocomplete("group")
async def role_group_add_role_group_autocomplete(
self, ctx: AutocompleteContext, group: str
):
groups_q: List[SelfRoleGroupsModel] = (
SelfRoleGroupsModel.select()
.where(
SelfRoleGroupsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupsModel.group_name.startswith(group),
)
.execute()
)
await ctx.send(choices=[g.group_name for g in groups_q])
@role_group_add_role.autocomplete("role")
async def role_group_add_role_role_autocomplete(
self, ctx: AutocompleteContext, role: str, **kwargs
):
roles_q: List[SelfRolesModel] = (
SelfRolesModel.select()
.where(
SelfRolesModel.guild_id == int(ctx.guild_id),
SelfRolesModel.role_name.startswith(role),
)
.execute()
)
await ctx.send(
choices=[{"name": r.role_name, "value": str(r.role_id)} for r in roles_q]
)
@slash_command(
name="role-group",
description="Manage self-role groups",
sub_cmd_name="remove-role",
sub_cmd_description="Remove a role from a self-role group",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
@slash_option(
name="group",
description="Group to remove the role from",
required=True,
opt_type=OptionTypes.STRING,
autocomplete=True,
)
@slash_option(
name="role",
description="Role to remove from the group",
required=True,
opt_type=OptionTypes.STRING,
autocomplete=True,
)
async def role_group_remove_role(
self, ctx: InteractionContext, group: str, role: str
):
await ctx.defer(ephemeral=True)
r: Role = await ctx.guild.fetch_role(role)
if r is None:
await ctx.send(f"Role {role} not found in this server.", ephemeral=True)
return
SelfRoleGroupRelationsModel.delete().where(
SelfRoleGroupRelationsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == group,
SelfRoleGroupRelationsModel.role_id == int(role),
).execute()
await ctx.send(
f"Removed role {r.mention} from the group {group}.", ephemeral=True
)
@role_group_remove_role.autocomplete("group")
async def role_group_remove_role_group_autocomplete(
self, ctx: AutocompleteContext, group: str
):
groups_q: List[SelfRoleGroupsModel] = (
SelfRoleGroupsModel.select()
.where(
SelfRoleGroupsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupsModel.group_name.startswith(group),
)
.execute()
)
await ctx.send(choices=[g.group_name for g in groups_q])
@role_group_remove_role.autocomplete("role")
async def role_group_remove_role_role_autocomplete(
self, ctx: AutocompleteContext, role: str, **kwargs
):
roles_q = (
SelfRolesModel.select()
.join(
SelfRoleGroupRelationsModel,
on=(SelfRolesModel.role_id == SelfRoleGroupRelationsModel.role_id),
)
.where(
SelfRolesModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == kwargs["group"],
SelfRolesModel.role_name.startswith(role),
)
.execute()
)
await ctx.send(
choices=[{"name": r.role_name, "value": str(r.role_id)} for r in roles_q]
)
@slash_command(
name="role-group",
description="Manage self-role groups",
sub_cmd_name="list-groups",
sub_cmd_description="List all self-role groups",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
async def role_group_list_groups(self, ctx: InteractionContext):
await ctx.defer(ephemeral=True)
embeds: List[Embed] = []
groups_q = (
SelfRolesModel.select(
SelfRolesModel.role_name,
SelfRolesModel.role_description,
SelfRoleGroupsModel.group_name,
SelfRoleGroupsModel.group_description,
)
.join(
SelfRoleGroupRelationsModel,
on=(SelfRoleGroupRelationsModel.role_id == SelfRolesModel.role_id),
)
.join(
SelfRoleGroupsModel,
on=(
SelfRoleGroupRelationsModel.group_name
== SelfRoleGroupsModel.group_name
),
)
.where(SelfRoleGroupsModel.guild_id == int(ctx.guild_id))
.objects()
)
groups: Dict[str, List[str, List[Tuple[str, str]]]] = {}
for row in groups_q:
if row.group_name not in groups:
groups[row.group_name] = [
row.group_description,
[(row.role_name, row.role_description)],
]
else:
groups[row.group_name][1].append((row.role_name, row.role_description))
for group, data in groups.items():
embed = Embed(title=group, description=data[0])
for role in data[1]:
embed.add_field(name=role[0], value=role[1], inline=False)
embeds.append(embed)
await ctx.send(
embeds=embeds,
ephemeral=True,
)
@slash_command(
name="role-group",
description="Manage self-role groups",
sub_cmd_name="generate",
sub_cmd_description="Generate a message with a selection of roles from a group",
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
@slash_option(
name="group",
description="Group to generate the message from",
required=True,
opt_type=OptionTypes.STRING,
autocomplete=True,
)
async def role_group_generate(self, ctx: InteractionContext, group: str):
await ctx.defer(ephemeral=True)
roles_q: List[SelfRolesModel] = (
SelfRolesModel.select(
SelfRolesModel.role_name,
SelfRolesModel.role_id,
SelfRolesModel.role_description,
)
.join(
SelfRoleGroupRelationsModel,
on=(SelfRolesModel.role_id == SelfRoleGroupRelationsModel.role_id),
)
.where(
SelfRoleGroupRelationsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == group,
)
.objects()
)
group_q: SelfRoleGroupsModel = SelfRoleGroupsModel.get(
SelfRoleGroupsModel.group_name == group
)
options = []
for role in roles_q:
opt = SelectOption(
label=role.role_name,
value=str(role.role_id),
description=role.role_description,
)
options.append(opt)
max_vals = 1 if group_q.exclusive else len(options)
select = StringSelectMenu(
options=options,
placeholder=f"Select roles from the group {group}",
custom_id=f"role-group-assign:{group}",
min_values=0,
max_values=max_vals,
)
msg = await ctx.channel.send(
embed=Embed(
title=f"{group}",
description=group_q.group_description,
),
components=select,
)
if msg is not None:
await ctx.send("Created!", ephemeral=True)
else:
await ctx.send("Failed to create message!", ephemeral=True)
@role_group_generate.autocomplete("group")
async def role_group_generate_group_autocomplete(
self, ctx: AutocompleteContext, group: str
):
groups_q: List[SelfRoleGroupsModel] = (
SelfRoleGroupsModel.select()
.where(
SelfRoleGroupsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupsModel.group_name.startswith(group),
)
.execute()
)
await ctx.send(choices=[g.group_name for g in groups_q])
@listen(events.Select)
async def on_role_selected(self, event: events.Select):
ctx = event.context
await ctx.defer(ephemeral=True)
if not ctx.custom_id.startswith("role-group-assign"):
return
group = ctx.custom_id.split(":")[1]
role_ids = [int(r) for r in ctx.values]
roles_q: List[SelfRolesModel] = (
SelfRolesModel.select(
SelfRolesModel.role_name,
SelfRolesModel.role_id,
SelfRolesModel.requires,
SelfRoleGroupRelationsModel.group_name,
)
.join(
SelfRoleGroupRelationsModel,
on=(SelfRolesModel.role_id == SelfRoleGroupRelationsModel.role_id),
)
.where(
SelfRoleGroupRelationsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == group,
)
.objects()
)
actions = []
for role in roles_q:
if role.role_id in role_ids:
if not ctx.author.has_role(role.role_id):
if not role.requires is None:
if not ctx.author.has_role(role.requires):
actions.append(
f"{role.role_name} requires {role.requires} and was not added."
)
continue
await ctx.author.add_role(role.role_id)
actions.append(f"Added role {role.role_name}")
else:
if ctx.author.has_role(role.role_id):
if not role.requires is None:
if not ctx.author.has_role(role.requires):
actions.append(
f"{role.role_name} requires {role.requires} to be managed, and was therefore not removed."
)
continue
await ctx.author.remove_role(role.role_id)
actions.append(f"Removed role {role.role_name}")
if len(actions) == 0:
await ctx.send("No roles were added or removed.")
else:
await ctx.send(
content="\n".join(actions),
ephemeral=True,
)
@context_menu(
name="Regenerate role group message",
context_type=CommandTypes.MESSAGE,
dm_permission=False,
default_member_permissions=Permissions.MANAGE_ROLES,
)
async def role_group_regenerate(self, ctx: InteractionContext):
await ctx.defer(ephemeral=True)
msg: Message = ctx.target
if msg.author.id != self.client.user.id:
await ctx.send("This message is not from the bot!", ephemeral=True)
return
if (
len(msg.embeds) < 1
or msg.embeds[0].title is None
or len(msg.components) < 1
or len(msg.components[0].components) < 1
or not isinstance(msg.components[0].components[0], StringSelectMenu)
or not msg.components[0].components[0].custom_id.startswith("role-group-assign")
):
await ctx.send("Could not identify this as a role-group message!", ephemeral=True)
return
embed = msg.embeds[0]
group = msg.components[0].components[0].custom_id.split(":")[1]
if (
group is None
or group == ""
or embed.title != group
):
await ctx.send("Could not identify this as a role-group message!", ephemeral=True)
return
roles_q: List[SelfRolesModel] = (
SelfRolesModel.select(
SelfRolesModel.role_name,
SelfRolesModel.role_id,
SelfRolesModel.role_description,
)
.join(
SelfRoleGroupRelationsModel,
on=(SelfRolesModel.role_id == SelfRoleGroupRelationsModel.role_id),
)
.where(
SelfRoleGroupRelationsModel.guild_id == int(ctx.guild_id),
SelfRoleGroupRelationsModel.group_name == group,
)
.objects()
)
group_q: SelfRoleGroupsModel = SelfRoleGroupsModel.get(
SelfRoleGroupsModel.group_name == group
)
options = []
for role in roles_q:
opt = SelectOption(
label=role.role_name,
value=str(role.role_id),
description=role.role_description,
)
options.append(opt)
max_vals = 1 if group_q.exclusive else len(options)
select = StringSelectMenu(
options=options,
placeholder=f"Select roles from the group {group}",
custom_id=f"role-group-assign:{group}",
min_values=0,
max_values=max_vals,
)
msg = await msg.edit(
embed=Embed(
title=f"{group}",
description=group_q.group_description,
),
components=select,
)
await ctx.send("Role group message regenerated!", 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("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, **kwargs
):
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("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, **kwargs
):
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()
SelfRoleGroupsModel.create_table()
SelfRoleGroupRelationsModel.create_table()
logging.info("SelfRoles extension loaded.")