From 71c08bab7ac08321a0f87c59f7e2bde641da6d8e Mon Sep 17 00:00:00 2001
From: eip <36315290+eip618@users.noreply.github.com>
Date: Sun, 10 Mar 2024 21:31:55 +1100
Subject: [PATCH] possibly fix wii restriction commands?

---
 cogs/mod.py     |  8 +++-----
 utils/checks.py | 15 ++++++++++++++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/cogs/mod.py b/cogs/mod.py
index b45db0d0a..2c51f1411 100644
--- a/cogs/mod.py
+++ b/cogs/mod.py
@@ -12,7 +12,7 @@
 from typing import Optional, TYPE_CHECKING
 from utils import Restriction
 from utils.converters import DateOrTimeToSecondsConverter, TimeTransformer
-from utils.checks import is_staff, check_staff, check_bot_or_staff, is_staff_app
+from utils.checks import is_staff, check_staff, check_bot_or_staff, is_staff_app, wii_check
 from utils.utils import paginate_message, send_dm_message, parse_time, text_to_discord_file, gen_color, \
     create_error_embed, create_userinfo_embed
 
@@ -937,23 +937,21 @@ async def giverestriction(self,
         await interaction.response.send_message(f"{member.mention} now has the {restriction} restriction role{f' until {format_dt(end_time)}.' if end_time else '.'}")
         await self.logs.post_action_log(interaction.user, member, restriction_action[restriction], reason=reason, until=end_time)
 
+    @wii_check()
     @commands.guild_only()
     @commands.command()
     async def takewii(self, ctx: GuildContext, member: discord.Member):
         """Remove access to the Wii channel."""
-        if not check_staff(ctx.bot, 'Helper', member.id) or ctx.bot.roles['Wii-Assistance'] not in ctx.author.roles:
-            return await ctx.send("You can't use this command!")
         if await check_bot_or_staff(ctx, member, "no-wii"):
             return
         await member.add_roles(ctx.bot.roles['No-Wii'])
         await ctx.send("Wii access taken.")
 
+    @wii_check()
     @commands.guild_only()
     @commands.command()
     async def givewii(self, ctx: GuildContext, member: discord.Member):
         """Restore access to the Wii channel."""
-        if not check_staff(ctx.bot, 'Helper', member.id) or ctx.bot.roles['Wii-Assistance'] not in ctx.author.roles:
-            return await ctx.send("You can't use this command!")
         await member.remove_roles(ctx.bot.roles['No-Wii'])
         await ctx.send("Wii access restored.")
 
diff --git a/utils/checks.py b/utils/checks.py
index bc1a6980b..554656ce4 100644
--- a/utils/checks.py
+++ b/utils/checks.py
@@ -68,8 +68,21 @@ async def predicate(ctx):
 def soap_check():
     async def predicate(ctx):
         author = ctx.author
-        if not check_staff(ctx.bot, 'Helper', author.id) and check_staff(ctx.bot, 'Staff', author.id) and (
+        if not check_staff(ctx.bot, 'Helper', author.id) and not check_staff(ctx.bot, 'Staff', author.id) and (
                 ctx.bot.roles['crc'] not in author.roles) and (ctx.bot.roles['Small Help'] not in author.roles):
+            raise InsufficientStaffRank("You can't use this command.")
+            return False
+        return True
+
+    return commands.check(predicate)
+
+
+def wii_check():
+    async def predicate(ctx):
+        author = ctx.author
+        if not check_staff(ctx.bot, 'Helper', author.id) and not check_staff(ctx.bot, 'Staff', author.id) and (
+                ctx.bot.roles['Wii-Assistance'] not in author.roles):
+            raise InsufficientStaffRank("You can't use this command.")
             return False
         return True