diff --git a/.env.example b/.env.example index 969cea3..218acb0 100644 --- a/.env.example +++ b/.env.example @@ -9,6 +9,9 @@ DATABASE_URL=sqlite+aiosqlite:///fun2oosh.db # For development # Logging LOG_LEVEL=INFO +# Fallback per-user cooldown (seconds) for commands with no cooldown of their own +DEFAULT_COMMAND_COOLDOWN_SECONDS=3 + # TTS Voice Channel Leave timeout in seconds TTS_VC_LEAVE_TIMEOUT=90 diff --git a/bot.py b/bot.py index 3c1f980..54263ae 100644 --- a/bot.py +++ b/bot.py @@ -15,6 +15,11 @@ from dotenv import load_dotenv from utils.config import Config +from utils.cooldowns import ( + GlobalCooldown, + app_command_interaction_check, + prefix_command_check, +) # Load environment variables load_dotenv() @@ -28,6 +33,17 @@ logger = logging.getLogger(__name__) +class EigenCommandTree(app_commands.CommandTree): + """Command tree with a fallback per-user cooldown for slash-only commands.""" + + def __init__(self, bot: commands.Bot, cooldown: GlobalCooldown): + super().__init__(bot) + self._cooldown_check = app_command_interaction_check(cooldown) + + async def interaction_check(self, interaction: discord.Interaction) -> bool: + return await self._cooldown_check(interaction) + + class Fun2OoshBot(commands.Bot): """Main bot class for Eigen Bot.""" @@ -38,6 +54,10 @@ def __init__(self, config: Config): intents.presences = True # For seeing user activities (Spotify, games, etc.) intents.voice_states = True # For join/leave voice channel features + self.default_cooldown = GlobalCooldown( + rate=1, per=config.default_command_cooldown_seconds + ) + # Disable the built-in help_command so a custom help cog can register `?helpmenu` and `/help` super().__init__( command_prefix="?", @@ -48,8 +68,11 @@ def __init__(self, config: Config): allowed_mentions=discord.AllowedMentions( everyone=False, roles=False, users=True, replied_user=False ), + tree_cls=lambda bot: EigenCommandTree(bot, self.default_cooldown), ) + self.add_check(prefix_command_check(self.default_cooldown)) + self.start_time = discord.utils.utcnow() self.config = config # Discover available cog modules from the cogs directory @@ -230,23 +253,23 @@ async def _safe_ctx_send(message: str) -> None: except Exception: return - if isinstance(error, commands.CommandOnCooldown): - await _safe_ctx_send( - f"This command is on cooldown. Try again in {error.retry_after:.2f} seconds." - ) - elif isinstance(error, commands.MissingRequiredArgument): - command = getattr(ctx, "command", None) - cog = getattr(command, "cog", None) - if command and command.qualified_name == "resources" and cog is not None: - send_usage = getattr(cog, "_send_resource_usage", None) - if callable(send_usage): - await send_usage(ctx) - return - await _safe_ctx_send("A required argument is missing.") - elif isinstance(error, commands.MissingPermissions): - await _safe_ctx_send("You don't have permission to use this command.") - elif isinstance(error, commands.BadArgument): - await _safe_ctx_send("Invalid argument provided.") + if isinstance(error, commands.CommandOnCooldown): + await _safe_ctx_send( + f"This command is on cooldown. Try again in {error.retry_after:.2f} seconds." + ) + elif isinstance(error, commands.MissingRequiredArgument): + command = getattr(ctx, "command", None) + cog = getattr(command, "cog", None) + if command and command.qualified_name == "resources" and cog is not None: + send_usage = getattr(cog, "_send_resource_usage", None) + if callable(send_usage): + await send_usage(ctx) + return + await _safe_ctx_send("A required argument is missing.") + elif isinstance(error, commands.MissingPermissions): + await _safe_ctx_send("You don't have permission to use this command.") + elif isinstance(error, commands.BadArgument): + await _safe_ctx_send("Invalid argument provided.") else: logger.error(f"Command error: {error}") await _safe_ctx_send("An error occurred while processing your command.") diff --git a/utils/config.py b/utils/config.py index f7ba725..859d797 100644 --- a/utils/config.py +++ b/utils/config.py @@ -18,6 +18,8 @@ class Config(BaseSettings): # Preferred: comma-separated list (or JSON list) of guild ids for fast per-guild slash-command sync guild_ids: list[int] = Field(default_factory=list) log_level: str = Field(default='INFO') + # Fallback per-user cooldown (seconds) for commands with no cooldown of their own. + default_command_cooldown_seconds: float = Field(default=3.0) owner_id: int | None = Field(default=None) topgg_token: str | None = Field(default=None) topgg_webhook_secret: str | None = Field(default=None) diff --git a/utils/cooldowns.py b/utils/cooldowns.py new file mode 100644 index 0000000..4e4eb4d --- /dev/null +++ b/utils/cooldowns.py @@ -0,0 +1,76 @@ +"""Fallback per-user cooldown applied to commands with no cooldown of their own. + +Commands that already declare an explicit `@commands.cooldown(...)` / +`@app_commands.checks.cooldown(...)` keep using that instead — this only +gates commands that would otherwise have no rate limit at all. +""" + +from __future__ import annotations + +import time + +from discord import Interaction, app_commands +from discord.ext import commands + + +class GlobalCooldown: + """Sliding-window per-user cooldown bucket.""" + + def __init__(self, rate: int, per: float) -> None: + self.rate = rate + self.per = per + self._uses: dict[int, list[float]] = {} + + def update_rate_limit(self, user_id: int) -> float | None: + """Record a use for `user_id`, returning seconds to wait if rate-limited.""" + now = time.monotonic() + window_start = now - self.per + uses = [t for t in self._uses.get(user_id, []) if t > window_start] + + if len(uses) >= self.rate: + self._uses[user_id] = uses + return self.per - (now - uses[0]) + + uses.append(now) + self._uses[user_id] = uses + return None + + +def prefix_command_check(cooldown: GlobalCooldown): + """Bot-wide `commands.check` used as a fallback for prefix/hybrid commands.""" + + async def predicate(ctx: commands.Context) -> bool: + command = ctx.command + if command is not None and command._buckets.valid: + # Command already has its own cooldown bucket; don't double-gate it. + return True + + retry_after = cooldown.update_rate_limit(ctx.author.id) + if retry_after is not None: + raise commands.CommandOnCooldown( + commands.Cooldown(cooldown.rate, cooldown.per), + retry_after, + commands.BucketType.user, + ) + return True + + return predicate + + +def app_command_interaction_check(cooldown: GlobalCooldown): + """Fallback cooldown check for slash-only commands via `CommandTree.interaction_check`.""" + + async def predicate(interaction: Interaction) -> bool: + command = interaction.command + if command is not None and getattr(command, "checks", None): + # Command already has its own check(s), e.g. a cooldown. + return True + + retry_after = cooldown.update_rate_limit(interaction.user.id) + if retry_after is not None: + raise app_commands.CommandOnCooldown( + app_commands.Cooldown(cooldown.rate, cooldown.per), retry_after + ) + return True + + return predicate