From 71f84895692ff814bf95f3b22dd29ee6e3a6556d Mon Sep 17 00:00:00 2001
From: nova <novanai.dev@gmail.com>
Date: Wed, 31 Jan 2024 18:07:17 +0000
Subject: [PATCH] Add  command & disable commands in DMs

---
 src/bot.py                  |  2 +-
 src/extensions/userroles.py | 68 +++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 src/extensions/userroles.py

diff --git a/src/bot.py b/src/bot.py
index 74ed884..e0eefef 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -19,5 +19,5 @@
 
 logging.info(f"Debug mode is {DEBUG}; You can safely ignore this.")
 
-arc_client = arc.GatewayClient(bot)
+arc_client = arc.GatewayClient(bot, is_dm_enabled=False)
 arc_client.load_extensions_from("./src/extensions/")
diff --git a/src/extensions/userroles.py b/src/extensions/userroles.py
new file mode 100644
index 0000000..db9047c
--- /dev/null
+++ b/src/extensions/userroles.py
@@ -0,0 +1,68 @@
+import arc
+import hikari
+
+plugin = arc.GatewayPlugin("User Roles")
+
+role = plugin.include_slash_group("role", "Get/remove assignable roles.")
+
+role_choices = [
+    hikari.CommandChoice(name="Webgroup", value="1166751688598761583"),
+    hikari.CommandChoice(name="gamz", value="1089204642241581139"),
+    hikari.CommandChoice(name="Croomer", value="1172696659097047050"),
+]
+
+
+@role.include
+@arc.slash_subcommand("add", "Add an assignable role.")
+async def add_role(
+    ctx: arc.GatewayContext, role: arc.Option[str, arc.StrParams("The role to add.", choices=role_choices)]
+) -> None:
+    assert ctx.guild_id
+    assert ctx.member
+
+    role_id = int(role)
+    if role_id not in ctx.member.role_ids:
+        await ctx.client.rest.add_role_to_member(
+            ctx.guild_id, ctx.author, int(role), reason=f"{ctx.author} added role."
+        )
+        await ctx.respond(f"Done! Added <@&{role}> to your roles.", flags=hikari.MessageFlag.EPHEMERAL)
+        return
+
+    await ctx.respond("You already have this role!", flags=hikari.MessageFlag.EPHEMERAL)
+
+
+@role.include
+@arc.slash_subcommand("remove", "Remove an assignable role.")
+async def remove_role(
+    ctx: arc.GatewayContext, role: arc.Option[str, arc.StrParams("The role to remove.", choices=role_choices)]
+) -> None:
+    assert ctx.guild_id
+    assert ctx.member
+
+    role_id = int(role)
+    if role_id in ctx.member.role_ids:
+        await ctx.client.rest.remove_role_from_member(
+            ctx.guild_id, ctx.author, int(role), reason=f"{ctx.author} removed role."
+        )
+        await ctx.respond(f"Done! Removed <@&{role}> from your roles.", flags=hikari.MessageFlag.EPHEMERAL)
+        return
+
+    await ctx.respond("You don't have this role!", flags=hikari.MessageFlag.EPHEMERAL)
+
+
+@role.set_error_handler
+async def role_error_handler(ctx: arc.GatewayContext, exc: Exception) -> None:
+    role = ctx.get_option("role", arc.OptionType.STRING)
+    assert role is not None
+    role_id = int(role)
+
+    if isinstance(exc, hikari.ForbiddenError):
+        await ctx.respond(f"You don't have permission to obtain <@&{role_id}>.", flags=hikari.MessageFlag.EPHEMERAL)
+        return
+
+    raise exc
+
+
+@arc.loader
+def loader(client: arc.GatewayClient) -> None:
+    client.add_plugin(plugin)