diff --git a/redbot/cogs/bank/__init__.py b/redbot/cogs/bank/__init__.py index e467452e2..9e59c17f8 100644 --- a/redbot/cogs/bank/__init__.py +++ b/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): diff --git a/redbot/cogs/bank/bank.py b/redbot/cogs/bank/bank.py index 48fb059e7..cc07bc6f7 100644 --- a/redbot/cogs/bank/bank.py +++ b/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: diff --git a/redbot/cogs/economy/economy.py b/redbot/cogs/economy/economy.py index 4fc05cd0e..fd3a30b08 100644 --- a/redbot/cogs/economy/economy.py +++ b/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 diff --git a/redbot/cogs/trivia/checks.py b/redbot/cogs/trivia/checks.py new file mode 100644 index 000000000..9cd8509fb --- /dev/null +++ b/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) diff --git a/redbot/cogs/trivia/converters.py b/redbot/cogs/trivia/converters.py new file mode 100644 index 000000000..211ed5d4b --- /dev/null +++ b/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 diff --git a/redbot/cogs/trivia/trivia.py b/redbot/cogs/trivia/trivia.py index 6ed9ded72..599df34f5 100644 --- a/redbot/cogs/trivia/trivia.py +++ b/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):