Make checks in Bank, Economy and Trivia cogs Permissions-friendly (#3672)

* Create converters.py

* Update trivia.py

* Create checks.py

* Update checks.py

* Update checks.py

* Update trivia.py

* Update checks.py

* Update checks.py

* Update trivia.py

* Update bank.py

* Update economy.py

* Update trivia.py

* Update checks.py

* Update checks.py

* Update __init__.py

redbot/cogs/bank/__init__.py

@@ -1,4 +1,4 @@
-from .bank import Bank, check_global_setting_guildowner, check_global_setting_admin
+from .bank import Bank, is_owner_if_bank_global
 
 
 def setup(bot):

redbot/cogs/bank/bank.py

@@ -9,10 +9,13 @@ from redbot.core.bot import Red  # Only used for type hints
 _ = Translator("Bank", __file__)
 
 
-def check_global_setting_guildowner():
+def is_owner_if_bank_global():
     """
-    Command decorator. If the bank is not global, it checks if the author is
-     either the guildowner or has the administrator permission.
+    Command decorator. If the bank is global, it checks if the author is
+    bot owner, otherwise it does nothing.
+     
+    When used on the command, this should be combined
+    with permissions check like `guildowner_or_permissions()`.
     """
 
     async def pred(ctx: commands.Context):
@@ -30,33 +33,6 @@ def check_global_setting_guildowner():
     return commands.check(pred)
 
 
-def check_global_setting_admin():
-    """
-    Command decorator. If the bank is not global, it checks if the author is
-     either a bot admin or has the manage_guild permission.
-    """
-
-    async def pred(ctx: commands.Context):
-        author = ctx.author
-        if not await bank.is_global():
-            if not isinstance(ctx.channel, discord.abc.GuildChannel):
-                return False
-            if await ctx.bot.is_owner(author):
-                return True
-            if author == ctx.guild.owner:
-                return True
-            if ctx.channel.permissions_for(author).manage_guild:
-                return True
-            admin_role_ids = await ctx.bot.get_admin_role_ids(ctx.guild.id)
-            for role in author.roles:
-                if role.id in admin_role_ids:
-                    return True
-        else:
-            return await ctx.bot.is_owner(author)
-
-    return commands.check(pred)
-
-
 @cog_i18n(_)
 class Bank(commands.Cog):
     """Bank"""
@@ -67,7 +43,7 @@ class Bank(commands.Cog):
 
     # SECTION commands
 
-    @check_global_setting_guildowner()
+    @is_owner_if_bank_global()
     @checks.guildowner_or_permissions(administrator=True)
     @commands.group(autohelp=True)
     async def bankset(self, ctx: commands.Context):
@@ -117,22 +93,25 @@ class Bank(commands.Cog):
             await bank.set_global(not cur_setting)
             await ctx.send(_("The bank is now {banktype}.").format(banktype=word))
 
+    @is_owner_if_bank_global()
+    @checks.guildowner_or_permissions(administrator=True)
     @bankset.command(name="bankname")
-    @check_global_setting_guildowner()
     async def bankset_bankname(self, ctx: commands.Context, *, name: str):
         """Set the bank's name."""
         await bank.set_bank_name(name, ctx.guild)
         await ctx.send(_("Bank name has been set to: {name}").format(name=name))
 
+    @is_owner_if_bank_global()
+    @checks.guildowner_or_permissions(administrator=True)
     @bankset.command(name="creditsname")
-    @check_global_setting_guildowner()
     async def bankset_creditsname(self, ctx: commands.Context, *, name: str):
         """Set the name for the bank's currency."""
         await bank.set_currency_name(name, ctx.guild)
         await ctx.send(_("Currency name has been set to: {name}").format(name=name))
 
+    @is_owner_if_bank_global()
+    @checks.guildowner_or_permissions(administrator=True)
     @bankset.command(name="maxbal")
-    @check_global_setting_guildowner()
     async def bankset_maxbal(self, ctx: commands.Context, *, amount: int):
         """Set the maximum balance a user can get."""
         try:

redbot/cogs/economy/economy.py

@@ -7,7 +7,7 @@ from typing import cast, Iterable, Union
 
 import discord
 
-from redbot.cogs.bank import check_global_setting_guildowner, check_global_setting_admin
+from redbot.cogs.bank import is_owner_if_bank_global
 from redbot.cogs.mod.converters import RawUserIds
 from redbot.core import Config, bank, commands, errors, checks
 from redbot.core.i18n import Translator, cog_i18n
@@ -191,8 +191,9 @@ class Economy(commands.Cog):
             )
         )
 
+    @is_owner_if_bank_global()
+    @checks.admin_or_permissions(manage_guild=True)
     @_bank.command(name="set")
-    @check_global_setting_admin()
     async def _set(self, ctx: commands.Context, to: discord.Member, creds: SetParser):
         """Set the balance of user's bank account.
 
@@ -236,8 +237,9 @@ class Economy(commands.Cog):
         else:
             await ctx.send(msg)
 
+    @is_owner_if_bank_global()
+    @checks.guildowner_or_permissions(administrator=True)
     @_bank.command()
-    @check_global_setting_guildowner()
     async def reset(self, ctx, confirmation: bool = False):
         """Delete all bank accounts."""
         if confirmation is False:
@@ -258,8 +260,9 @@ class Economy(commands.Cog):
                 )
             )
 
+    @is_owner_if_bank_global()
+    @checks.admin_or_permissions(manage_guild=True)
     @_bank.group(name="prune")
-    @check_global_setting_admin()
     async def _prune(self, ctx):
         """Prune bank accounts."""
         pass
@@ -646,9 +649,10 @@ class Economy(commands.Cog):
             )
         )
 
-    @commands.group()
     @guild_only_check()
-    @check_global_setting_admin()
+    @is_owner_if_bank_global()
+    @checks.admin_or_permissions(manage_guild=True)
+    @commands.group()
     async def economyset(self, ctx: commands.Context):
         """Manage Economy settings."""
         guild = ctx.guild

redbot/cogs/trivia/checks.py

@@ -0,0 +1,25 @@
+from redbot.core import commands
+from redbot.core.i18n import Translator
+
+__all__ = ("trivia_stop_check",)
+
+_ = Translator("Trivia", __file__)
+
+
+def trivia_stop_check():
+    async def predicate(ctx: commands.GuildContext) -> bool:
+        session = ctx.cog._get_trivia_session(ctx.channel)
+        if session is None:
+            raise commands.CheckFailure(_("There is no ongoing trivia session in this channel."))
+
+        author = ctx.author
+        auth_checks = (
+            await ctx.bot.is_owner(author),
+            await ctx.bot.is_mod(author),
+            await ctx.bot.is_admin(author),
+            author == ctx.guild.owner,
+            author == session.ctx.author,
+        )
+        return any(auth_checks)
+
+    return commands.permissions_check(predicate)

redbot/cogs/trivia/converters.py

@@ -0,0 +1,18 @@
+import math
+
+from redbot.core import commands
+from redbot.core.i18n import Translator
+
+__all__ = ("finite_float",)
+
+_ = Translator("Trivia", __file__)
+
+
+def finite_float(arg: str) -> float:
+    try:
+        ret = float(arg)
+    except ValueError:
+        raise commands.BadArgument(_("`{arg}` is not a number.").format(arg=arg))
+    if not math.isfinite(ret):
+        raise commands.BadArgument(_("`{arg}` is not a finite number.").format(arg=ret))
+    return ret

redbot/cogs/trivia/trivia.py

@@ -4,17 +4,21 @@ import math
 import pathlib
 from collections import Counter
 from typing import List
+
 import io
 import yaml
 import discord
-from redbot.core import commands
-from redbot.core import Config, checks
+
+from redbot.core import Config, commands, checks
+from redbot.cogs.bank import is_owner_if_bank_global
 from redbot.core.data_manager import cog_data_path
 from redbot.core.i18n import Translator, cog_i18n
 from redbot.core.utils.chat_formatting import box, pagify, bold
-from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate
 from redbot.core.utils.menus import start_adding_reactions
-from redbot.cogs.bank import check_global_setting_admin
+from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate
+
+from .checks import trivia_stop_check
+from .converters import finite_float
 from .log import LOG
 from .session import TriviaSession
 
@@ -25,16 +29,6 @@ UNIQUE_ID = 0xB3C0E453
 _ = Translator("Trivia", __file__)
 
 
-def finite_float(arg: str) -> float:
-    try:
-        ret = float(arg)
-    except ValueError:
-        raise commands.BadArgument(_("`{arg}` is not a number.").format(arg=arg))
-    if not math.isfinite(ret):
-        raise commands.BadArgument(_("`{arg}` is not a finite number.").format(arg=ret))
-    return ret
-
-
 class InvalidListError(Exception):
     """A Trivia list file is in invalid format."""
 
@@ -163,8 +157,9 @@ class Trivia(commands.Cog):
         else:
             await ctx.send(_("Alright, I won't reveal the answer to the questions anymore."))
 
+    @is_owner_if_bank_global()
+    @checks.admin_or_permissions(manage_guild=True)
     @triviaset.command(name="payout")
-    @check_global_setting_admin()
     async def triviaset_payout_multiplier(self, ctx: commands.Context, multiplier: finite_float):
         """Set the payout multiplier.
 
@@ -321,27 +316,14 @@ class Trivia(commands.Cog):
         self.trivia_sessions.append(session)
         LOG.debug("New trivia session; #%s in %d", ctx.channel, ctx.guild.id)
 
+    @trivia_stop_check()
     @trivia.command(name="stop")
     async def trivia_stop(self, ctx: commands.Context):
         """Stop an ongoing trivia session."""
         session = self._get_trivia_session(ctx.channel)
-        if session is None:
-            await ctx.send(_("There is no ongoing trivia session in this channel."))
-            return
-        author = ctx.author
-        auth_checks = (
-            await ctx.bot.is_owner(author),
-            await ctx.bot.is_mod(author),
-            await ctx.bot.is_admin(author),
-            author == ctx.guild.owner,
-            author == session.ctx.author,
-        )
-        if any(auth_checks):
-            await session.end_game()
-            session.force_stop()
-            await ctx.send(_("Trivia stopped."))
-        else:
-            await ctx.send(_("You are not allowed to do that."))
+        await session.end_game()
+        session.force_stop()
+        await ctx.send(_("Trivia stopped."))
 
     @trivia.command(name="list")
     async def trivia_list(self, ctx: commands.Context):