1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
import logging
import random
import typing
from asyncio import Lock
from functools import wraps
from weakref import WeakValueDictionary
from discord import Colour, Embed
from discord.ext import commands
from discord.ext.commands import CheckFailure, Context
from bot.constants import ERROR_REPLIES
log = logging.getLogger(__name__)
class InChannelCheckFailure(CheckFailure):
"""Check failure when the user runs a command in a non-whitelisted channel."""
pass
def with_role(*role_ids: int):
"""Check to see whether the invoking user has any of the roles specified in role_ids."""
async def predicate(ctx: Context):
if not ctx.guild: # Return False in a DM
log.debug(
f"{ctx.author} tried to use the '{ctx.command.name}'command from a DM. "
"This command is restricted by the with_role decorator. Rejecting request."
)
return False
for role in ctx.author.roles:
if role.id in role_ids:
log.debug(f"{ctx.author} has the '{role.name}' role, and passes the check.")
return True
log.debug(
f"{ctx.author} does not have the required role to use "
f"the '{ctx.command.name}' command, so the request is rejected."
)
return False
return commands.check(predicate)
def without_role(*role_ids: int):
"""Check whether the invoking user does not have all of the roles specified in role_ids."""
async def predicate(ctx: Context):
if not ctx.guild: # Return False in a DM
log.debug(
f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM. "
"This command is restricted by the without_role decorator. Rejecting request."
)
return False
author_roles = [role.id for role in ctx.author.roles]
check = all(role not in author_roles for role in role_ids)
log.debug(
f"{ctx.author} tried to call the '{ctx.command.name}' command. "
f"The result of the without_role check was {check}."
)
return check
return commands.check(predicate)
def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None) -> typing.Callable[[Context], bool]:
"""Checks that the message is in a whitelisted channel or optionally has a bypass role."""
def predicate(ctx: Context) -> bool:
if not ctx.guild:
log.debug(f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM.")
return True
if ctx.channel.id in channels:
log.debug(
f"{ctx.author} tried to call the '{ctx.command.name}' command "
f"and the command was used in a whitelisted channel."
)
return True
if hasattr(ctx.command.callback, "in_channel_override"):
log.debug(
f"{ctx.author} called the '{ctx.command.name}' command "
f"and the command was whitelisted to bypass the in_channel check."
)
return True
if bypass_roles and any(r.id in bypass_roles for r in ctx.author.roles):
log.debug(
f"{ctx.author} called the '{ctx.command.name}' command and "
f"had a role to bypass the in_channel check."
)
return True
log.debug(
f"{ctx.author} tried to call the '{ctx.command.name}' command. "
f"The in_channel check failed."
)
channels_str = ', '.join(f"<#{c_id}>" for c_id in channels)
raise InChannelCheckFailure(
f"Sorry, but you may only use this command within {channels_str}."
)
return predicate
in_channel = commands.check(in_channel_check)
def override_in_channel(func: typing.Callable) -> typing.Callable:
"""
Set command callback attribute for detection in `in_channel_check`.
This decorator has to go before (below) below the `command` decorator.
"""
func.in_channel_override = True
return func
def locked():
"""
Allows the user to only run one instance of the decorated command at a time.
Subsequent calls to the command from the same author are ignored until the command has completed invocation.
This decorator has to go before (below) the `command` decorator.
"""
def wrap(func):
func.__locks = WeakValueDictionary()
@wraps(func)
async def inner(self, ctx, *args, **kwargs):
lock = func.__locks.setdefault(ctx.author.id, Lock())
if lock.locked():
embed = Embed()
embed.colour = Colour.red()
log.debug(f"User tried to invoke a locked command.")
embed.description = (
"You're already using this command. Please wait until "
"it is done before you use it again."
)
embed.title = random.choice(ERROR_REPLIES)
await ctx.send(embed=embed)
return
async with func.__locks.setdefault(ctx.author.id, Lock()):
return await func(self, ctx, *args, **kwargs)
return inner
return wrap
|
