diff --git a/cogs/dynamicconfig/cog.py b/cogs/dynamicconfig/cog.py index 702591b07..88d192b13 100644 --- a/cogs/dynamicconfig/cog.py +++ b/cogs/dynamicconfig/cog.py @@ -12,7 +12,7 @@ from cogs.base import Base from config.app_config import CONFIG_PATH, config_get_keys, load_config -from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from .messages_cz import MessagesCZ @@ -27,7 +27,7 @@ def __init__(self, bot: Rubbergod): super().__init__() self.bot = bot - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="config") async def config_cmd(self, inter: disnake.ApplicationCommandInteraction): """ diff --git a/cogs/fitwide/cog.py b/cogs/fitwide/cog.py index b503b9c78..507e965a1 100644 --- a/cogs/fitwide/cog.py +++ b/cogs/fitwide/cog.py @@ -16,7 +16,8 @@ from features.prompt import PromptSession from features.verification import Verification from features.verify_helper import VerifyHelper -from permissions import permission_check, room_check +from permissions import room_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from utils import cooldowns from utils.general import edit_msg @@ -49,7 +50,7 @@ def get_all_logins(self): user_logins = sorted(user_logins) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="verify_check", description=MessagesCZ.verify_check_brief) async def verify_check(self, inter: disnake.ApplicationCommandInteraction): @@ -68,7 +69,7 @@ async def verify_check(self, inter: disnake.ApplicationCommandInteraction): await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="status_check", description=MessagesCZ.status_check_brief) async def status_check(self, inter: disnake.ApplicationCommandInteraction): @@ -99,7 +100,7 @@ async def status_check(self, inter: disnake.ApplicationCommandInteraction): await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="zapis_check", description=MessagesCZ.zapis_check_brief) async def zapis_check(self, inter: disnake.ApplicationCommandInteraction): @@ -126,7 +127,7 @@ async def zapis_check(self, inter: disnake.ApplicationCommandInteraction): await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="mit_check", description=MessagesCZ.mit_check_brief) async def mit_check(self, inter: disnake.ApplicationCommandInteraction, p_debug: bool = True): @@ -164,7 +165,7 @@ async def mit_check(self, inter: disnake.ApplicationCommandInteraction, p_debug: await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="dropout_check", description=MessagesCZ.dropout_check_brief) async def dropout_check(self, inter: disnake.ApplicationCommandInteraction, p_debug: bool = True): @@ -205,7 +206,7 @@ async def dropout_check(self, inter: disnake.ApplicationCommandInteraction, p_de await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.check(room_check.is_in_modroom) @commands.slash_command(name="role_check", description=MessagesCZ.role_check_brief) async def role_check(self, inter: disnake.ApplicationCommandInteraction): @@ -233,7 +234,7 @@ async def role_check(self, inter: disnake.ApplicationCommandInteraction): await message.reply(MessagesCZ.role_check_end) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="increment_roles", description=MessagesCZ.increment_roles_brief) async def increment_roles(self, inter: disnake.ApplicationCommandInteraction): await inter.send(MessagesCZ.increment_roles_start) diff --git a/cogs/ios/cog.py b/cogs/ios/cog.py index 6070ea042..22ae43db9 100644 --- a/cogs/ios/cog.py +++ b/cogs/ios/cog.py @@ -9,6 +9,7 @@ from cogs.base import Base from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from utils import cooldowns @@ -33,7 +34,7 @@ async def ios(self, inter: disnake.ApplicationCommandInteraction): async def _ios(self, inter: disnake.ApplicationCommandInteraction): pass - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @_ios.sub_command(name="start", description=MessagesCZ.task_start_brief) async def ios_task_start(self, inter: disnake.ApplicationCommandInteraction): try: @@ -42,7 +43,7 @@ async def ios_task_start(self, inter: disnake.ApplicationCommandInteraction): except RuntimeError: await inter.send(MessagesCZ.task_start_already_set) - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @_ios.sub_command(name="stop", description=MessagesCZ.task_stop_brief) async def ios_task_stop(self, inter: disnake.ApplicationCommandInteraction): if self.ios_task.is_running(): @@ -51,7 +52,7 @@ async def ios_task_stop(self, inter: disnake.ApplicationCommandInteraction): else: await inter.send(MessagesCZ.task_nothing_to_stop) - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @_ios.sub_command(name="cancel", description=MessagesCZ.task_cancel_brief) async def ios_task_cancel(self, inter: disnake.ApplicationCommandInteraction): if self.ios_task.is_running(): diff --git a/cogs/karma/cog.py b/cogs/karma/cog.py index e86ad8763..61c33b252 100644 --- a/cogs/karma/cog.py +++ b/cogs/karma/cog.py @@ -14,7 +14,8 @@ from database.karma import KarmaDB from features.leaderboard import LeaderboardPageSource from features.reaction_context import ReactionContext -from permissions import permission_check, room_check +from permissions import room_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from utils import cooldowns @@ -227,7 +228,7 @@ async def givingboard( await inter.edit_original_response(embed=embed, view=view) view.message = await inter.original_message() - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="karma_mod", description=MessagesCZ.karma_brief) async def _karma_mod(self, inter: disnake.ApplicationCommandInteraction): pass diff --git a/cogs/review/cog.py b/cogs/review/cog.py index 4f379beb1..a35f5b7d1 100644 --- a/cogs/review/cog.py +++ b/cogs/review/cog.py @@ -12,7 +12,7 @@ from buttons.embed import PaginationView from cogs.base import Base from database.review import ProgrammeDB, ReviewDB, SubjectDB, SubjectDetailsDB -from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from utils import cooldowns @@ -116,7 +116,7 @@ async def remove( """ await inter.response.defer() if id is not None: - if permission_check.is_bot_admin(inter, False): + if PermissionsCheck.is_bot_admin(inter, False): review = ReviewDB.get_review_by_id(id) if review: review.remove() @@ -143,7 +143,7 @@ async def author_list(self, inter: disnake.ApplicationCommandInteraction): @cooldowns.short_cooldown @commands.slash_command(name="subject") - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() async def subject(self, inter: disnake.ApplicationCommandInteraction): """Group of commands for managing subjects in DB""" await inter.response.defer() diff --git a/cogs/system/cog.py b/cogs/system/cog.py index 61da9d5d6..0b14b163d 100644 --- a/cogs/system/cog.py +++ b/cogs/system/cog.py @@ -15,7 +15,7 @@ from database.error import ErrorLogDB from features.error import ErrorLogger from features.git import Git -from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from utils import cooldowns from utils.colors import RubbergodColors @@ -43,7 +43,7 @@ def check_first_boot(self): if not start_streak: ErrorLogDB.set() - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="get_logs", description=MessagesCZ.get_logs_brief) async def get_logs( self, @@ -83,7 +83,7 @@ async def get_logs( await inter.send(files=files) - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="shutdown", description=MessagesCZ.shutdown_brief) async def shutdown(self, inter: disnake.ApplicationCommandInteraction): await inter.send("Shutting down...") @@ -92,7 +92,7 @@ async def shutdown(self, inter: disnake.ApplicationCommandInteraction): await self.bot.vutapi_session.close() await self.bot.close() - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="cogs", description=MessagesCZ.cogs_brief, guild_ids=[Base.config.guild_id]) async def cogs(self, inter: disnake.ApplicationCommandInteraction): """ @@ -165,7 +165,7 @@ async def uptime(self, inter: disnake.ApplicationCommandInteraction): await inter.edit_original_response(embed=embed) @cooldowns.default_cooldown - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() @commands.slash_command(name="command_checks", description=MessagesCZ.command_checks_brief) async def command_checks(self, inter: disnake.ApplicationCommandInteraction): await inter.response.defer() diff --git a/cogs/system/views.py b/cogs/system/views.py index 443371fd6..9d7d57c1e 100644 --- a/cogs/system/views.py +++ b/cogs/system/views.py @@ -3,7 +3,7 @@ import disnake from buttons.base import BaseView -from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from . import features @@ -46,7 +46,7 @@ async def on_timeout(self) -> None: pass async def interaction_check(self, inter: disnake.Interaction) -> bool: - if permission_check.is_bot_admin(inter): + if PermissionsCheck.is_bot_admin(inter): return True return False @@ -96,7 +96,7 @@ def create_select(self) -> list[disnake.SelectOption]: async def callback(self, inter: disnake.MessageInteraction) -> None: """React to user selecting cog(s).""" await inter.response.defer() - if not permission_check.is_bot_admin(inter): + if not PermissionsCheck.is_bot_admin(inter): return unloadable = [cog for cog in self.unloadable_cogs if cog in self.values] diff --git a/cogs/warden/cog.py b/cogs/warden/cog.py index 615a9b734..af6bb4e54 100644 --- a/cogs/warden/cog.py +++ b/cogs/warden/cog.py @@ -13,7 +13,7 @@ import utils from cogs.base import Base from database.image import ImageDB -from permissions import permission_check +from permissions.permission_check import PermissionsCheck from rubbergod import Rubbergod from . import features @@ -92,7 +92,7 @@ async def handle_reaction(self, ctx): pass @commands.group() - @commands.check(permission_check.is_bot_admin) + @PermissionsCheck.is_bot_admin() async def scan(self, ctx: commands.Context): """Scan for reposts""" if ctx.invoked_subcommand is None: diff --git a/permissions/permission_check.py b/permissions/permission_check.py index 714657861..1405b56c2 100644 --- a/permissions/permission_check.py +++ b/permissions/permission_check.py @@ -1,7 +1,5 @@ """All checks for role and user permissions""" -from typing import Union - import disnake from disnake.ext import commands @@ -20,22 +18,44 @@ def __init__(self) -> None: self.message = Messages.bot_admin_only -def is_bot_admin( - ctx: Union[commands.Context, disnake.ApplicationCommandInteraction], raise_exception: bool = True -): - """ - Check if user is bot admin. - if raise_exception=False return False instead of raising error - """ - if ctx.author.id in config.admin_ids: - return True +class PermissionsCheck: + """Class containing methods to check permissions""" - if not raise_exception: - return False - raise NotAdminError + @staticmethod + def is_bot_admin( + ctx: commands.Context | disnake.ApplicationCommandInteraction = None, raise_exception: bool = True + ) -> bool: + """ + Check if the user is a bot admin. Works as a function or a command decorator. + + When `ctx` is passed, it works as a function. + When `ctx` is not passed, it acts as a decorator for a command. + + Throws `NotAdminError` if user is not an admin. + If `raise_exception=False` this function returns only boolean value and doesn't raise exception. + """ + def predicate(ctx: commands.Context | disnake.ApplicationCommandInteraction) -> bool: + if ctx.author.id in config.admin_ids: + return True + if not raise_exception: + return False + raise NotAdminError -def role_check(ctx, roles, MissingPermission, raise_exception=True): + if ctx: + # If ctx is passed, return the result of the predicate + return predicate(ctx) + else: + # If ctx is not passed, act as a decorator for command + return commands.check(predicate) + + +def role_check( + ctx: commands.Context | disnake.ApplicationCommandInteraction, + roles, + MissingPermission, + raise_exception: bool = True, +): """ Check if user has permission for command. if raise_exception=False return False instead of raising error @@ -61,7 +81,7 @@ def __init__(self) -> None: self.message = Messages.helper_plus_only -def helper_plus(ctx, raise_exception=True): +def helper_plus(ctx: commands.Context | disnake.ApplicationCommandInteraction, raise_exception: bool = True): """Check if user has permission for command helper or above.""" return role_check(ctx, allowed_roles, NotHelperPlusError, raise_exception) @@ -75,7 +95,7 @@ def __init__(self) -> None: self.message = Messages.submod_plus_only -def submod_plus(ctx, raise_exception=True): +def submod_plus(ctx: commands.Context | disnake.ApplicationCommandInteraction, raise_exception: bool = True): """Check if user has permission for command submod or above.""" return role_check(ctx, allowed_roles[:-1], NotSubmodPlusError, raise_exception) @@ -89,6 +109,6 @@ def __init__(self) -> None: self.message = Messages.mod_plus_only -def mod_plus(ctx, raise_exception=True): +def mod_plus(ctx: commands.Context | disnake.ApplicationCommandInteraction, raise_exception: bool = True): """Check if user has permission for command mod or above.""" return role_check(ctx, allowed_roles[:-2], NotModPlusError, raise_exception)