jarvis-bot/jarvis/cogs/rolegiver.py

363 lines
12 KiB
Python

"""J.A.R.V.I.S. Role Giver Cog."""
import asyncio
from discord import Role
from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from discord_slash.utils.manage_components import (
create_actionrow,
create_select,
create_select_option,
wait_for_component,
)
from jarvis.db.models import Rolegiver
from jarvis.utils import build_embed
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions
class RolegiverCog(commands.Cog):
"""J.A.R.V.I.S. Role Giver Cog."""
def __init__(self, bot: commands.Bot):
self.bot = bot
@cog_ext.cog_subcommand(
base="rolegiver",
name="add",
description="Add a role to rolegiver",
options=[
create_option(
name="role",
description="Role to add",
option_type=8,
required=True,
)
],
)
@admin_or_permissions(manage_guild=True)
async def _rolegiver_add(self, ctx: SlashContext, role: Role) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if setting and role.id in setting.roles:
await ctx.send("Role already in rolegiver", hidden=True)
return
if not setting:
setting = Rolegiver(guild=ctx.guild.id, roles=[])
if len(setting.roles) >= 20:
await ctx.send("You can only have 20 roles in the rolegiver", hidden=True)
return
setting.roles.append(role.id)
setting.save()
roles = []
for role_id in setting.roles:
if role_id == role.id:
continue
e_role = ctx.guild.get_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 = [
Field(name="New Role", value=f"{role.mention}"),
Field(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.nick if ctx.author.nick else ctx.author.name,
icon_url=ctx.author.avatar_url,
)
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed)
@cog_ext.cog_subcommand(
base="rolegiver",
name="remove",
description="Remove a role from rolegiver",
)
@admin_or_permissions(manage_guild=True)
async def _rolegiver_remove(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True)
return
options = []
for role in setting.roles:
role: Role = ctx.guild.get_role(role)
option = create_select_option(label=role.name, value=str(role.id))
options.append(option)
select = create_select(
options=options,
custom_id="to_delete",
placeholder="Select roles to remove",
min_values=1,
max_values=len(options),
)
components = [create_actionrow(select)]
message = await ctx.send(content="\u200b", components=components)
try:
context = await wait_for_component(
self.bot,
check=lambda x: ctx.author.id == x.author.id,
message=message,
timeout=60 * 1,
)
for to_delete in context.selected_options:
setting.roles.remove(int(to_delete))
setting.save()
for row in components:
for component in row["components"]:
component["disabled"] = True
roles = []
for role_id in setting.roles:
e_role = ctx.guild.get_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 = [
Field(name="Removed Role", value=f"{role.mention}"),
Field(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.nick if ctx.author.nick else ctx.author.name,
icon_url=ctx.author.avatar_url,
)
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}")
await context.edit_origin(
content=f"Removed {len(context.selected_options)} role(s)",
embed=embed,
components=components,
)
except asyncio.TimeoutError:
for row in components:
for component in row["components"]:
component["disabled"] = True
await message.edit(components=components)
@cog_ext.cog_subcommand(
base="rolegiver",
name="list",
description="List roles rolegiver",
)
async def _rolegiver_list(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True)
return
roles = []
for role_id in setting.roles:
e_role = ctx.guild.get_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.nick if ctx.author.nick else ctx.author.name,
icon_url=ctx.author.avatar_url,
)
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed)
@cog_ext.cog_subcommand(
base="role",
name="get",
description="Get a role from rolegiver",
)
@commands.cooldown(1, 10, commands.BucketType.user)
async def _role_get(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True)
return
options = []
for role in setting.roles:
role: Role = ctx.guild.get_role(role)
option = create_select_option(label=role.name, value=str(role.id))
options.append(option)
select = create_select(
options=options,
custom_id="to_delete",
placeholder="Select roles to add",
min_values=1,
max_values=len(options),
)
components = [create_actionrow(select)]
message = await ctx.send(content="\u200b", components=components)
try:
context = await wait_for_component(
self.bot,
check=lambda x: ctx.author.id == x.author.id,
messages=message,
timeout=60 * 5,
)
for role in context.selected_options:
role = ctx.guild.get_role(int(role))
await ctx.author.add_roles(role, reason="Rolegiver")
roles = ctx.author.roles
if roles:
roles.sort(key=lambda x: -x.position)
_ = roles.pop(-1)
value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [
Field(name="Added Role", value=f"{role.mention}"),
Field(name="Prior Role(s)", value=value),
]
embed = build_embed(
title="User Given Role",
description=f"{role.mention} given to {ctx.author.mention}",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name,
icon_url=ctx.author.avatar_url,
)
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed)
except asyncio.TimeoutError:
for row in components:
for component in row["components"]:
component["disabled"] = True
await message.edit(components=components)
@cog_ext.cog_subcommand(
base="role",
name="forfeit",
description="Have rolegiver take away role",
options=[
create_option(
name="role",
description="Role to remove",
option_type=8,
required=True,
)
],
)
@commands.cooldown(1, 10, commands.BucketType.user)
async def _role_forfeit(self, ctx: SlashContext, role: Role) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True)
return
elif role.id not in setting.roles:
await ctx.send("Role not in rolegiver", hidden=True)
return
elif role not in ctx.author.roles:
await ctx.send("You do not have that role", hidden=True)
return
await ctx.author.remove_roles(role, reason="Rolegiver")
roles = ctx.author.roles
if roles:
roles.sort(key=lambda x: -x.position)
_ = roles.pop(-1)
value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [
Field(name="Taken Role", value=f"{role.mention}"),
Field(name="Remaining Role(s)", value=value),
]
embed = build_embed(
title="User Forfeited Role",
description=f"{role.mention} taken from {ctx.author.mention}",
fields=fields,
)
embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name,
icon_url=ctx.author.avatar_url,
)
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed)
@cog_ext.cog_subcommand(
base="rolegiver",
name="cleanup",
description="Cleanup rolegiver roles",
)
@admin_or_permissions(manage_guild=True)
async def _rolegiver_cleanup(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or not setting.roles:
await ctx.send("Rolegiver has no roles", hidden=True)
guild_roles = await ctx.guild.fetch_roles()
guild_role_ids = [x.id for x in guild_roles]
for role_id in setting.roles:
if role_id not in guild_role_ids:
setting.roles.remove(role_id)
setting.save()
await ctx.send("Rolegiver cleanup finished")
def setup(bot: commands.Bot) -> None:
"""Add RolegiverCog to J.A.R.V.I.S."""
bot.add_cog(RolegiverCog(bot))