[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.
This commit is contained in:
Michael H 2018-12-15 18:09:18 -05:00 committed by Toby Harradine
parent 2d9912cea7
commit 351749dff6
2 changed files with 135 additions and 6 deletions

View File

@ -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"^<?(?:(?:@[!&]?)?|#)(\d{15,21})>?$")
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

View File

@ -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.