From 351749dff657729a2173f7b8133b14636280d185 Mon Sep 17 00:00:00 2001 From: Michael H Date: Sat, 15 Dec 2018 18:09:18 -0500 Subject: [PATCH] [Permissions] Find things uniquely for models (#2258) This is a safety measure to prevent accidentally passing a model which has the same name as another model, potentially modifying rules for the unwanted one. --- redbot/cogs/permissions/converters.py | 125 ++++++++++++++++++++++++- redbot/cogs/permissions/permissions.py | 16 +++- 2 files changed, 135 insertions(+), 6 deletions(-) diff --git a/redbot/cogs/permissions/converters.py b/redbot/cogs/permissions/converters.py index 1b1668dc6..cbf1db76d 100644 --- a/redbot/cogs/permissions/converters.py +++ b/redbot/cogs/permissions/converters.py @@ -1,10 +1,133 @@ -from typing import NamedTuple, Union, Optional, cast, Type +import itertools +import re +from typing import NamedTuple, Union, Optional + +import discord from redbot.core import commands from redbot.core.i18n import Translator _ = Translator("PermissionsConverters", __file__) +MENTION_RE = re.compile(r"^?$") + + +def _match_id(arg: str) -> Optional[int]: + m = MENTION_RE.match(arg) + if m: + return int(m.group(1)) + + +class GlobalUniqueObjectFinder(commands.Converter): + async def convert( + self, ctx: commands.Context, arg: str + ) -> Union[discord.Guild, discord.abc.GuildChannel, discord.abc.User, discord.Role]: + bot: commands.Bot = ctx.bot + _id = _match_id(arg) + + if _id is not None: + guild: discord.Guild = bot.get_guild(_id) + if guild is not None: + return guild + channel: discord.abc.GuildChannel = bot.get_channel(_id) + if channel is not None: + return channel + + user: discord.User = bot.get_user(_id) + if user is not None: + return user + + for guild in bot.guilds: + role: discord.Role = guild.get_role(_id) + if role is not None: + return role + + objects = itertools.chain( + bot.get_all_channels(), + bot.users, + bot.guilds, + *(filter(lambda r: not r.is_default(), guild.roles) for guild in bot.guilds), + ) + + maybe_matches = [] + for obj in objects: + if obj.name == arg or str(obj) == arg: + maybe_matches.append(obj) + + if ctx.guild is not None: + for member in ctx.guild.members: + if member.nick == arg and not any(obj.id == member.id for obj in maybe_matches): + maybe_matches.append(member) + + if not maybe_matches: + raise commands.BadArgument( + _( + '"{arg}" was not found. It must be the ID, mention, or name of a server, ' + "channel, user or role which the bot can see." + ).format(arg=arg) + ) + elif len(maybe_matches) == 1: + return maybe_matches[0] + else: + raise commands.BadArgument( + _( + '"{arg}" does not refer to a unique server, channel, user or role. Please use ' + "the ID for whatever/whoever you're trying to specify, or mention it/them." + ).format(arg=arg) + ) + + +class GuildUniqueObjectFinder(commands.Converter): + async def convert( + self, ctx: commands.Context, arg: str + ) -> Union[discord.abc.GuildChannel, discord.Member, discord.Role]: + guild: discord.Guild = ctx.guild + _id = _match_id(arg) + + if _id is not None: + channel: discord.abc.GuildChannel = guild.get_channel(_id) + if channel is not None: + return channel + + member: discord.Member = guild.get_member(_id) + if member is not None: + return member + + role: discord.Role = guild.get_role(_id) + if role is not None and not role.is_default(): + return role + + objects = itertools.chain( + guild.channels, guild.members, filter(lambda r: not r.is_default(), guild.roles) + ) + + maybe_matches = [] + for obj in objects: + if obj.name == arg or str(obj) == arg: + maybe_matches.append(obj) + try: + if obj.nick == arg: + maybe_matches.append(obj) + except AttributeError: + pass + + if not maybe_matches: + raise commands.BadArgument( + _( + '"{arg}" was not found. It must be the ID, mention, or name of a channel, ' + "user or role in this server." + ).format(arg=arg) + ) + elif len(maybe_matches) == 1: + return maybe_matches[0] + else: + raise commands.BadArgument( + _( + '"{arg}" does not refer to a unique channel, user or role. Please use the ID ' + "for whatever/whoever you're trying to specify, or mention it/them." + ).format(arg=arg) + ) + class CogOrCommand(NamedTuple): type: str diff --git a/redbot/cogs/permissions/permissions.py b/redbot/cogs/permissions/permissions.py index 00854d6f1..3376368c7 100644 --- a/redbot/cogs/permissions/permissions.py +++ b/redbot/cogs/permissions/permissions.py @@ -14,7 +14,13 @@ from redbot.core.utils.chat_formatting import box from redbot.core.utils.menus import start_adding_reactions from redbot.core.utils.predicates import ReactionPredicate, MessagePredicate -from .converters import CogOrCommand, RuleType, ClearableRuleType +from .converters import ( + CogOrCommand, + RuleType, + ClearableRuleType, + GuildUniqueObjectFinder, + GlobalUniqueObjectFinder, +) _ = Translator("Permissions", __file__) @@ -275,7 +281,7 @@ class Permissions(commands.Cog): ctx: commands.Context, allow_or_deny: RuleType, cog_or_command: CogOrCommand, - who_or_what: commands.GlobalPermissionModel, + who_or_what: GlobalUniqueObjectFinder, ): """Add a global rule to a command. @@ -303,7 +309,7 @@ class Permissions(commands.Cog): ctx: commands.Context, allow_or_deny: RuleType, cog_or_command: CogOrCommand, - who_or_what: commands.GuildPermissionModel, + who_or_what: GuildUniqueObjectFinder, ): """Add a rule to a command in this server. @@ -328,7 +334,7 @@ class Permissions(commands.Cog): self, ctx: commands.Context, cog_or_command: CogOrCommand, - who_or_what: commands.GlobalPermissionModel, + who_or_what: GlobalUniqueObjectFinder, ): """Remove a global rule from a command. @@ -351,7 +357,7 @@ class Permissions(commands.Cog): ctx: commands.Context, cog_or_command: CogOrCommand, *, - who_or_what: commands.GuildPermissionModel, + who_or_what: GuildUniqueObjectFinder, ): """Remove a server rule from a command.