jarvis-bot/jarvis/cogs/rolegiver.py

392 lines
14 KiB
Python

"""JARVIS Role Giver Cog."""
import asyncio
import logging
from jarvis_core.db import q
from jarvis_core.db.models import Rolegiver
from naff import Client, Extension, InteractionContext, Permissions
from naff.client.utils.misc_utils import get
from naff.models.discord.components import ActionRow, Select, SelectOption
from naff.models.discord.embed import EmbedField
from naff.models.discord.role import Role
from naff.models.naff.application_commands import (
OptionTypes,
SlashCommand,
slash_option,
)
from naff.models.naff.command import check, cooldown
from naff.models.naff.cooldowns import Buckets
from jarvis.utils import build_embed
from jarvis.utils.permissions import admin_or_permissions
class RolegiverCog(Extension):
"""JARVIS Role Giver Cog."""
def __init__(self, bot: Client):
self.bot = bot
self.logger = logging.getLogger(__name__)
rolegiver = SlashCommand(name="rolegiver", description="Allow users to choose their own roles")
@rolegiver.subcommand(
sub_cmd_name="add",
sub_cmd_description="Add a role to rolegiver",
)
@slash_option(name="role", description="Role to add", opt_type=OptionTypes.ROLE, required=True)
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _rolegiver_add(self, ctx: InteractionContext, role: Role) -> None:
if role.id == ctx.guild.id:
await ctx.send("Cannot add `@everyone` to rolegiver", ephemeral=True)
return
if role.bot_managed or not role.is_assignable:
await ctx.send(
"Cannot assign this role, try lowering it below my role or using a different role",
ephemeral=True,
)
return
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if setting and setting.roles and role.id in setting.roles:
await ctx.send("Role already in rolegiver", ephemeral=True)
return
if not setting:
setting = Rolegiver(guild=ctx.guild.id, roles=[])
setting.roles = setting.roles or []
if len(setting.roles) >= 20:
await ctx.send("You can only have 20 roles in the rolegiver", ephemeral=True)
return
setting.roles.append(role.id)
await setting.commit()
roles = []
for role_id in setting.roles:
if role_id == role.id:
continue
e_role = await ctx.guild.fetch_role(role_id)
if not e_role:
continue
roles.append(e_role)
if roles:
roles.sort(key=lambda x: -x.position)
value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [
EmbedField(name="New Role", value=f"{role.mention}"),
EmbedField(name="Existing Role(s)", value=value),
]
embed = build_embed(
title="Rolegiver Updated",
description="New role added to rolegiver",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon.url)
embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.display_avatar.url)
embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embeds=embed)
@rolegiver.subcommand(sub_cmd_name="remove", sub_cmd_description="Remove a role from rolegiver")
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _rolegiver_remove(self, ctx: InteractionContext) -> None:
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", ephemeral=True)
return
options = []
for role in setting.roles:
role: Role = await ctx.guild.fetch_role(role)
option = SelectOption(label=role.name, value=str(role.id))
options.append(option)
select = Select(
options=options,
custom_id="to_delete",
placeholder="Select roles to remove",
min_values=1,
max_values=len(options),
)
components = [ActionRow(select)]
message = await ctx.send(content="\u200b", components=components)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.context.author.id,
messages=message,
timeout=60 * 1,
)
removed_roles = []
for to_delete in context.context.values:
role = await ctx.guild.fetch_role(to_delete)
if role:
removed_roles.append(role)
setting.roles.remove(int(to_delete))
await setting.commit()
for row in components:
for component in row.components:
component.disabled = True
roles = []
for role_id in setting.roles:
e_role = await ctx.guild.fetch_role(role_id)
if not e_role:
continue
roles.append(e_role)
if roles:
roles.sort(key=lambda x: -x.position)
value = "\n".join([r.mention for r in roles]) if roles else "None"
rvalue = "\n".join([r.mention for r in removed_roles]) if removed_roles else "None"
fields = [
EmbedField(name="Removed Role(s)", value=rvalue),
EmbedField(name="Remaining Role(s)", value=value),
]
embed = build_embed(
title="Rolegiver Updated",
description="Role removed from rolegiver",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon.url)
embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.display_avatar.url)
embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
)
await context.context.edit_origin(
content=f"Removed {len(context.context.values)} role(s)",
embeds=embed,
components=components,
)
except asyncio.TimeoutError:
for row in components:
for component in row.components:
component.disabled = True
await message.edit(components=components)
@rolegiver.subcommand(sub_cmd_name="list", sub_cmd_description="List rolegiver roles")
async def _rolegiver_list(self, ctx: InteractionContext) -> None:
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", ephemeral=True)
return
roles = []
for role_id in setting.roles:
e_role = await ctx.guild.fetch_role(role_id)
if not e_role:
continue
roles.append(e_role)
if roles:
roles.sort(key=lambda x: -x.position)
value = "\n".join([r.mention for r in roles]) if roles else "None"
embed = build_embed(
title="Rolegiver",
description=f"Available roles:\n{value}",
fields=[],
)
embed.set_thumbnail(url=ctx.guild.icon.url)
embed.set_author(
name=ctx.author.display_name,
icon_url=ctx.author.display_avatar.url,
)
embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embeds=embed)
role = SlashCommand(name="role", description="Get/Remove Rolegiver roles")
@role.subcommand(sub_cmd_name="get", sub_cmd_description="Get a role")
@cooldown(bucket=Buckets.USER, rate=1, interval=10)
async def _role_get(self, ctx: InteractionContext) -> None:
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", ephemeral=True)
return
options = []
for role in setting.roles:
role: Role = await ctx.guild.fetch_role(role)
option = SelectOption(label=role.name, value=str(role.id))
options.append(option)
select = Select(
options=options,
placeholder="Select roles to add",
min_values=1,
max_values=len(options),
)
components = [ActionRow(select)]
message = await ctx.send(content="\u200b", components=components)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.context.author.id,
messages=message,
timeout=60 * 5,
)
added_roles = []
for role in context.context.values:
role = await ctx.guild.fetch_role(int(role))
added_roles.append(role)
await ctx.author.add_role(role, reason="Rolegiver")
roles = ctx.author.roles
if roles:
roles.sort(key=lambda x: -x.position)
_ = roles.pop(-1)
avalue = "\n".join([r.mention for r in added_roles]) if added_roles else "None"
value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [
EmbedField(name="Added Role(s)", value=avalue),
EmbedField(name="Prior Role(s)", value=value),
]
embed = build_embed(
title="User Given Role",
description=f"{len(added_roles)} role(s) given to {ctx.author.mention}",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon.url)
embed.set_author(
name=ctx.author.display_name,
icon_url=ctx.author.display_avatar.url,
)
embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
)
for row in components:
for component in row.components:
component.disabled = True
await context.context.edit_origin(embeds=embed, content="\u200b", components=components)
except asyncio.TimeoutError:
for row in components:
for component in row.components:
component.disabled = True
await message.edit(components=components)
@role.subcommand(sub_cmd_name="remove", sub_cmd_description="Remove a role")
@cooldown(bucket=Buckets.USER, rate=1, interval=10)
async def _role_remove(self, ctx: InteractionContext) -> None:
user_roles = ctx.author.roles
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", ephemeral=True)
return
elif not any(x.id in setting.roles for x in user_roles):
await ctx.send("You have no rolegiver roles", ephemeral=True)
return
valid = list(filter(lambda x: x.id in setting.roles, user_roles))
options = []
for role in valid:
option = SelectOption(label=role.name, value=str(role.id))
options.append(option)
select = Select(
options=options,
custom_id="to_remove",
placeholder="Select roles to remove",
min_values=1,
max_values=len(options),
)
components = [ActionRow(select)]
message = await ctx.send(content="\u200b", components=components)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.context.author.id,
messages=message,
timeout=60 * 5,
)
removed_roles = []
for to_remove in context.context.values:
role = get(user_roles, id=int(to_remove))
await ctx.author.remove_role(role, reason="Rolegiver")
user_roles.remove(role)
removed_roles.append(role)
user_roles.sort(key=lambda x: -x.position)
_ = user_roles.pop(-1)
value = "\n".join([r.mention for r in user_roles]) if user_roles else "None"
rvalue = "\n".join([r.mention for r in removed_roles]) if removed_roles else "None"
fields = [
EmbedField(name="Removed Role(s)", value=rvalue),
EmbedField(name="Remaining Role(s)", value=value),
]
embed = build_embed(
title="User Forfeited Role",
description=f"{len(removed_roles)} role(s) removed from {ctx.author.mention}",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon.url)
embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.display_avatar.url)
embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
)
for row in components:
for component in row.components:
component.disabled = True
await context.context.edit_origin(embeds=embed, components=components, content="\u200b")
except asyncio.TimeoutError:
for row in components:
for component in row.components:
component.disabled = True
await message.edit(components=components)
@rolegiver.subcommand(
sub_cmd_name="cleanup", sub_cmd_description="Removed deleted roles from rolegiver"
)
@check(admin_or_permissions(Permissions.MANAGE_GUILD))
async def _rolegiver_cleanup(self, ctx: InteractionContext) -> None:
setting = await Rolegiver.find_one(q(guild=ctx.guild.id))
if not setting or not setting.roles:
await ctx.send("Rolegiver has no roles", ephemeral=True)
guild_role_ids = [r.id for r in ctx.guild.roles]
for role_id in setting.roles:
if role_id not in guild_role_ids:
setting.roles.remove(role_id)
await setting.commit()
await ctx.send("Rolegiver cleanup finished")
def setup(bot: Client) -> None:
"""Add RolegiverCog to JARVIS"""
RolegiverCog(bot)