diff --git a/config/settings.yml.example b/config/settings.yml.example index f1024e5..8cfa314 100644 --- a/config/settings.yml.example +++ b/config/settings.yml.example @@ -6,8 +6,9 @@ DEFAULT_PREFIX: USER_IDS: SYSADMINS: - - 123456789012345 - BOT_OWNER: 123456789012345 + - 123456789012345679 + - 123456789012345679 + BOT_OWNER: 123456789012345679 TEMPVC_CATEGORY_ID: 123456789012345 TEMPVC_CHANNEL_ID: 123456789012345 @@ -40,13 +41,13 @@ EMBED_ICONS: TIMEOUT: "https://github.com/allthingslinux/tux/blob/main/assets/emojis/timeout.png?raw=true" WARN: "https://github.com/allthingslinux/tux/blob/main/assets/emojis/warn.png?raw=true" -RECENT_GIF_AGE: 60 +GIF_LIMITER: + RECENT_GIF_AGE: 60 -GIF_LIMIT_EXCLUDE: - - 123456789012345 + GIF_LIMIT_EXCLUDE: + - 123456789012345 -GIF_LIMITS_USER: - "123456789012345": 2 - -GIF_LIMITS_CHANNEL: - "123456789012345": 3 + GIF_LIMITS_USER: + "123456789012345": 2 + GIF_LIMITS_CHANNEL: + "123456789012345": 3 diff --git a/tux/cogs/services/gif_limiter.py b/tux/cogs/services/gif_limiter.py index 790f142..922685f 100644 --- a/tux/cogs/services/gif_limiter.py +++ b/tux/cogs/services/gif_limiter.py @@ -1,41 +1,20 @@ -""" -This cog is a handler for GIF ratelimiting. -It keeps a list of GIF send times and routinely removes old times. -If a user posts a GIF, the message_handler function should be externally called. -It will delete the message if the user or channel quota is exceeded. -""" - import asyncio from collections import defaultdict from time import time import discord -from discord import Message from discord.ext import commands, tasks -from loguru import logger from tux.bot import Tux from tux.utils.constants import CONST -def convert_dict_str_to_int(original_dict: dict[str, int]) -> dict[int, int]: - """Helper function required as YAML keys are str. Channel and user IDs are int.""" - - converted_dict: dict[int, int] = {} - - for key, value in original_dict.items(): - try: - int_key: int = int(key) - converted_dict[int_key] = value - except ValueError: - logger.error("An error occurred when loading the GIF ratelimiter configuration.") - - return converted_dict - - class GifLimiter(commands.Cog): - """Main class with GIF tracking and message handlers""" - + """ + This class is a handler for GIF ratelimiting. + It keeps a list of GIF send times and routinely removes old times. + It will prevent people from posting GIFs if the quotas are exceeded. + """ def __init__(self, bot: Tux) -> None: self.bot = bot @@ -43,9 +22,9 @@ class GifLimiter(commands.Cog): self.recent_gif_age: int = CONST.RECENT_GIF_AGE # Max number of GIFs sent recently in a channel - self.channelwide_gif_limits: dict[int, int] = convert_dict_str_to_int(CONST.GIF_LIMITS_CHANNEL) + self.channelwide_gif_limits: dict[int, int] = CONST.GIF_LIMITS_CHANNEL # Max number of GIFs sent recently by a user to be able to post one in specified channels - self.user_gif_limits: dict[int, int] = convert_dict_str_to_int(CONST.GIF_LIMITS) + self.user_gif_limits: dict[int, int] = CONST.GIF_LIMITS # list of channels in which not to count GIFs self.gif_limit_exclude: list[int] = CONST.GIF_LIMIT_EXCLUDE @@ -57,68 +36,73 @@ class GifLimiter(commands.Cog): # Channel ID, list of timestamps self.recent_gifs_by_channel: defaultdict[int, list[int]] = defaultdict(list) - async def delete_message(self, message: discord.Message, epilogue: str) -> None: + # Lock to prevent race conditions + self.gif_lock = asyncio.Lock() + + self.old_gif_remover.start() + + async def _should_process_message(self, message: discord.Message) -> bool: + """ Checks if a message contains a GIF and was not sent in a blacklisted channel """ + return not (len(message.embeds) == 0 + or "gif" not in message.content.lower() + or message.channel.id in self.gif_limit_exclude) + + async def _handle_gif_message(self, message: discord.Message) -> None: + """ Checks for ratelimit infringements """ + async with self.gif_lock: + channel: int = message.channel.id + user: int = message.author.id + + if ( + channel in self.channelwide_gif_limits + and channel in self.recent_gifs_by_channel + and len(self.recent_gifs_by_channel[channel]) >= self.channelwide_gif_limits[channel] + ): + await self._delete_message(message, "for channel") + return + + if ( + user in self.recent_gifs_by_user + and channel in self.user_gif_limits + and len(self.recent_gifs_by_user[user]) >= self.user_gif_limits[channel] + ): + await self._delete_message(message, "for user") + return + + # Add message to recent GIFs if it doesn't infringe on ratelimits + current_time: int = int(time()) + self.recent_gifs_by_channel[channel].append(current_time) + self.recent_gifs_by_user[user].append(current_time) + + async def _delete_message(self, message: discord.Message, epilogue: str) -> None: """ Deletes the message passed as an argument, and sends a self-deleting message with the reason """ - sent_message: Message = await message.channel.send(f"-# GIF ratelimit exceeded {epilogue}") await message.delete() - await asyncio.sleep(3) - await sent_message.delete() + await message.channel.send(f"-# GIF ratelimit exceeded {epilogue}", delete_after=3) @commands.Cog.listener() async def on_message(self, message: discord.Message) -> None: """Checks for GIFs in every sent message""" - - # Nothing to do if the message doesn't have a .gif embed, - # or if it was sent in a blacklisted channel - if ( - len(message.embeds) == 0 - or "gif" not in message.content.lower() - or message.channel.id in self.gif_limit_exclude - ): - return - - channel: int = message.channel.id - user: int = message.author.id - - # Check if the message infringes on any ratelimits - if ( - channel in self.channelwide_gif_limits - and channel in self.recent_gifs_by_channel - and len(self.recent_gifs_by_channel[channel]) >= self.channelwide_gif_limits[channel] - ): - await self.delete_message(message, "for channel") - return - - if ( - user in self.recent_gifs_by_user - and channel in self.user_gif_limits - and len(self.recent_gifs_by_user[user]) >= self.user_gif_limits[channel] - ): - await self.delete_message(message, "for user") - return - - # If it doesn't, add it to recent GIFs - current_time: int = int(time()) - self.recent_gifs_by_channel[channel].append(current_time) - self.recent_gifs_by_user[user].append(current_time) + + if (await self._should_process_message(message)): + await self._handle_gif_message(message) @tasks.loop(seconds=20) async def old_gif_remover(self) -> None: """Regularly cleans old GIF timestamps""" current_time: int = int(time()) - for channel_id, timestamps in self.recent_gifs_by_channel.items(): - self.recent_gifs_by_channel[channel_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] + async with self.gif_lock: + for channel_id, timestamps in self.recent_gifs_by_channel.items(): + self.recent_gifs_by_channel[channel_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] - for user_id, timestamps in self.recent_gifs_by_user.items(): - self.recent_gifs_by_user[user_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] - - # Delete user key if no GIF has recently been sent by them - if len(self.recent_gifs_by_user[user_id]) == 0: - del self.recent_gifs_by_user[user_id] + for user_id, timestamps in self.recent_gifs_by_user.items(): + self.recent_gifs_by_user[user_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] + # Delete user key if no GIF has recently been sent by them + if len(self.recent_gifs_by_user[user_id]) == 0: + del self.recent_gifs_by_user[user_id] async def setup(bot: Tux) -> None: await bot.add_cog(GifLimiter(bot)) diff --git a/tux/utils/constants.py b/tux/utils/constants.py index f0e0b97..76bc2cb 100644 --- a/tux/utils/constants.py +++ b/tux/utils/constants.py @@ -6,6 +6,8 @@ from typing import Final import yaml from dotenv import load_dotenv, set_key +from tux.utils.functions import convert_dict_str_to_int + load_dotenv(verbose=True) config_file = Path("config/settings.yml") @@ -78,11 +80,11 @@ class Constants: EMBED_ICONS: Final[dict[str, str]] = config["EMBED_ICONS"] # GIF ratelimit constants - RECENT_GIF_AGE: Final[int] = config["RECENT_GIF_AGE"] - GIF_LIMIT_EXCLUDE: Final[list[int]] = config["GIF_LIMIT_EXCLUDE"] - # Ideally would be int, int but YAML doesn't support integer keys - GIF_LIMITS: Final[dict[str, int]] = config["GIF_LIMITS_USER"] - GIF_LIMITS_CHANNEL: Final[dict[str, int]] = config["GIF_LIMITS_CHANNEL"] + RECENT_GIF_AGE: Final[int] = config["GIF_LIMITER"]["RECENT_GIF_AGE"] + GIF_LIMIT_EXCLUDE: Final[list[int]] = config["GIF_LIMITER"]["GIF_LIMIT_EXCLUDE"] + + GIF_LIMITS: Final[dict[int, int]] = convert_dict_str_to_int(config["GIF_LIMITER"]["GIF_LIMITS_USER"]) + GIF_LIMITS_CHANNEL: Final[dict[int, int]] = convert_dict_str_to_int(config["GIF_LIMITER"]["GIF_LIMITS_CHANNEL"]) # Embed limit constants EMBED_MAX_NAME_LENGTH = 256 diff --git a/tux/utils/functions.py b/tux/utils/functions.py index 1245e99..e90a559 100644 --- a/tux/utils/functions.py +++ b/tux/utils/functions.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime, timedelta from typing import Any import discord +from loguru import logger harmful_command_pattern = r"(?:sudo\s+|doas\s+|run0\s+)?rm\s+(-[frR]*|--force|--recursive|--no-preserve-root|\s+)*([/\∕~]\s*|\*|/bin|/boot|/etc|/lib|/proc|/root|/sbin|/sys|/tmp|/usr|/var|/var/log|/network.|/system)(\s+--no-preserve-root|\s+\*)*|:\(\)\{ :|:& \};:" # noqa: RUF001 @@ -300,3 +301,19 @@ def extract_member_attrs(member: discord.Member) -> dict[str, Any]: "status": member.status, "activity": member.activity, } + + +def convert_dict_str_to_int(original_dict: dict[str, int]) -> dict[int, int]: + """Helper function used for GIF Limiter constants. + Required as YAML keys are str. Channel and user IDs are int.""" + + converted_dict: dict[int, int] = {} + + for key, value in original_dict.items(): + try: + int_key: int = int(key) + converted_dict[int_key] = value + except ValueError: + logger.exception(f"An error occurred when loading the GIF ratelimiter configuration at key {key}") + + return converted_dict