Only accept positive integers in [p]cleanup commands (#4115)

This commit is contained in:
jack1142 2020-07-29 02:44:36 +02:00 committed by GitHub
parent d73ad3115f
commit 8ada1ee152
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 36 additions and 14 deletions

View File

@ -1,6 +1,6 @@
import logging import logging
from datetime import datetime, timedelta from datetime import datetime, timedelta
from typing import Union, List, Callable, Set from typing import Callable, List, Optional, Set, Union
import discord import discord
@ -10,7 +10,7 @@ from redbot.core.i18n import Translator, cog_i18n
from redbot.core.utils.chat_formatting import humanize_number from redbot.core.utils.chat_formatting import humanize_number
from redbot.core.utils.mod import slow_deletion, mass_purge from redbot.core.utils.mod import slow_deletion, mass_purge
from redbot.core.utils.predicates import MessagePredicate from redbot.core.utils.predicates import MessagePredicate
from .converters import RawMessageIds from .converters import PositiveInt, RawMessageIds, positive_int
_ = Translator("Cleanup", __file__) _ = Translator("Cleanup", __file__)
@ -60,9 +60,9 @@ class Cleanup(commands.Cog):
async def get_messages_for_deletion( async def get_messages_for_deletion(
*, *,
channel: discord.TextChannel, channel: discord.TextChannel,
number: int = None, number: Optional[PositiveInt] = None,
check: Callable[[discord.Message], bool] = lambda x: True, check: Callable[[discord.Message], bool] = lambda x: True,
limit: int = None, limit: Optional[PositiveInt] = None,
before: Union[discord.Message, datetime] = None, before: Union[discord.Message, datetime] = None,
after: Union[discord.Message, datetime] = None, after: Union[discord.Message, datetime] = None,
delete_pinned: bool = False, delete_pinned: bool = False,
@ -105,7 +105,7 @@ class Cleanup(commands.Cog):
break break
if message_filter(message): if message_filter(message):
collected.append(message) collected.append(message)
if number and number <= len(collected): if number is not None and number <= len(collected):
break break
return collected return collected
@ -120,7 +120,7 @@ class Cleanup(commands.Cog):
@commands.guild_only() @commands.guild_only()
@commands.bot_has_permissions(manage_messages=True) @commands.bot_has_permissions(manage_messages=True)
async def text( async def text(
self, ctx: commands.Context, text: str, number: int, delete_pinned: bool = False self, ctx: commands.Context, text: str, number: positive_int, delete_pinned: bool = False
): ):
"""Delete the last X messages matching the specified text. """Delete the last X messages matching the specified text.
@ -169,7 +169,7 @@ class Cleanup(commands.Cog):
@commands.guild_only() @commands.guild_only()
@commands.bot_has_permissions(manage_messages=True) @commands.bot_has_permissions(manage_messages=True)
async def user( async def user(
self, ctx: commands.Context, user: str, number: int, delete_pinned: bool = False self, ctx: commands.Context, user: str, number: positive_int, delete_pinned: bool = False
): ):
"""Delete the last X messages from a specified user. """Delete the last X messages from a specified user.
@ -270,7 +270,7 @@ class Cleanup(commands.Cog):
self, self,
ctx: commands.Context, ctx: commands.Context,
message_id: RawMessageIds, message_id: RawMessageIds,
number: int, number: positive_int,
delete_pinned: bool = False, delete_pinned: bool = False,
): ):
"""Deletes X messages before specified message. """Deletes X messages before specified message.
@ -351,7 +351,9 @@ class Cleanup(commands.Cog):
@cleanup.command() @cleanup.command()
@commands.guild_only() @commands.guild_only()
@commands.bot_has_permissions(manage_messages=True) @commands.bot_has_permissions(manage_messages=True)
async def messages(self, ctx: commands.Context, number: int, delete_pinned: bool = False): async def messages(
self, ctx: commands.Context, number: positive_int, delete_pinned: bool = False
):
"""Delete the last X messages. """Delete the last X messages.
Example: Example:
@ -381,7 +383,9 @@ class Cleanup(commands.Cog):
@cleanup.command(name="bot") @cleanup.command(name="bot")
@commands.guild_only() @commands.guild_only()
@commands.bot_has_permissions(manage_messages=True) @commands.bot_has_permissions(manage_messages=True)
async def cleanup_bot(self, ctx: commands.Context, number: int, delete_pinned: bool = False): async def cleanup_bot(
self, ctx: commands.Context, number: positive_int, delete_pinned: bool = False
):
"""Clean up command messages and messages from the bot.""" """Clean up command messages and messages from the bot."""
channel = ctx.channel channel = ctx.channel
@ -458,7 +462,7 @@ class Cleanup(commands.Cog):
async def cleanup_self( async def cleanup_self(
self, self,
ctx: commands.Context, ctx: commands.Context,
number: int, number: positive_int,
match_pattern: str = None, match_pattern: str = None,
delete_pinned: bool = False, delete_pinned: bool = False,
): ):
@ -531,7 +535,7 @@ class Cleanup(commands.Cog):
@cleanup.command(name="spam") @cleanup.command(name="spam")
@commands.guild_only() @commands.guild_only()
@commands.bot_has_permissions(manage_messages=True) @commands.bot_has_permissions(manage_messages=True)
async def cleanup_spam(self, ctx: commands.Context, number: int = 50): async def cleanup_spam(self, ctx: commands.Context, number: positive_int = PositiveInt(50)):
"""Deletes duplicate messages in the channel from the last X messages and keeps only one copy. """Deletes duplicate messages in the channel from the last X messages and keeps only one copy.
Defaults to 50. Defaults to 50.

View File

@ -1,12 +1,30 @@
from redbot.core.commands import Converter, BadArgument from typing import NewType, TYPE_CHECKING
from redbot.core.commands import BadArgument, Context, Converter
from redbot.core.i18n import Translator from redbot.core.i18n import Translator
from redbot.core.utils.chat_formatting import inline
_ = Translator("Cleanup", __file__) _ = Translator("Cleanup", __file__)
class RawMessageIds(Converter): class RawMessageIds(Converter):
async def convert(self, ctx, argument): async def convert(self, ctx: Context, argument: str) -> int:
if argument.isnumeric() and len(argument) >= 17: if argument.isnumeric() and len(argument) >= 17:
return int(argument) return int(argument)
raise BadArgument(_("{} doesn't look like a valid message ID.").format(argument)) raise BadArgument(_("{} doesn't look like a valid message ID.").format(argument))
PositiveInt = NewType("PositiveInt", int)
if TYPE_CHECKING:
positive_int = PositiveInt
else:
def positive_int(arg: str) -> int:
try:
ret = int(arg)
except ValueError:
raise BadArgument(_("{arg} is not an integer.").format(arg=inline(arg)))
if ret <= 0:
raise BadArgument(_("{arg} is not a positive integer.").format(arg=inline(arg)))
return ret