From d46207961b16af0c38f85594835a52935f6ef524 Mon Sep 17 00:00:00 2001 From: Saksham Date: Tue, 15 Sep 2026 21:03:52 +0530 Subject: [PATCH 1/2] Add configurable command cooldown to prevent spam --- .env.example | 3 ++ bot.py | 89 ++++++++++++++++++++++++++++++++++-------- tests/test_cooldown.py | 87 +++++++++++++++++++++++++++++++++++++++++ utils/config.py | 1 + utils/cooldown.py | 59 ++++++++++++++++++++++++++++ 5 files changed, 222 insertions(+), 17 deletions(-) create mode 100644 tests/test_cooldown.py create mode 100644 utils/cooldown.py diff --git a/.env.example b/.env.example index 969cea3..d8d13c0 100644 --- a/.env.example +++ b/.env.example @@ -9,6 +9,9 @@ DATABASE_URL=sqlite+aiosqlite:///fun2oosh.db # For development # Logging LOG_LEVEL=INFO +# Command Cooldown +COMMAND_COOLDOWN_SECONDS=3.0 # Per-user cooldown between commands in seconds + # TTS Voice Channel Leave timeout in seconds TTS_VC_LEAVE_TIMEOUT=90 diff --git a/bot.py b/bot.py index 3c1f980..58d7bdc 100644 --- a/bot.py +++ b/bot.py @@ -15,6 +15,7 @@ from dotenv import load_dotenv from utils.config import Config +from utils.cooldown import CooldownManager # Load environment variables load_dotenv() @@ -52,6 +53,9 @@ def __init__(self, config: Config): self.start_time = discord.utils.utcnow() self.config = config + self.cooldown_manager = CooldownManager(config.command_cooldown_seconds) + self.tree.on_error = self.on_app_command_error + self._setup_cooldowns() # Discover available cog modules from the cogs directory from pathlib import Path @@ -63,6 +67,57 @@ def __init__(self, config: Config): self.available_cogs.append(p.stem) logger.info(f"Available cogs discovered: {self.available_cogs}") + def _setup_cooldowns(self) -> None: + """Register global cooldown checks for prefix and slash commands.""" + + @self.check + async def prefix_cooldown_check(ctx: commands.Context) -> bool: + if self.config.owner_id and ctx.author.id == self.config.owner_id: + return True + + # If command defines its own cooldown, let discord.py handle it + cmd = ctx.command + if cmd is not None and cmd._buckets and cmd._buckets.valid: + return True + + retry_after = self.cooldown_manager.trigger(ctx.author.id) + if retry_after is not None: + raise commands.CommandOnCooldown( + commands.Cooldown(1, self.cooldown_manager.cooldown_seconds), + retry_after, + commands.BucketType.user, + ) + return True + + async def app_command_cooldown_check( + interaction: discord.Interaction, + ) -> bool: + # Only rate-limit application command executions (not autocomplete or components) + if interaction.type != discord.InteractionType.application_command: + return True + + if self.config.owner_id and interaction.user.id == self.config.owner_id: + return True + + cmd = interaction.command + if cmd is not None: + # Let commands with custom checks or hybrid cooldowns handle their own rate limits + if getattr(cmd, "checks", None): + return True + wrapped = getattr(cmd, "wrapped", None) + if wrapped and getattr(wrapped, "_buckets", None) and wrapped._buckets.valid: + return True + + retry_after = self.cooldown_manager.trigger(interaction.user.id) + if retry_after is not None: + raise app_commands.CommandOnCooldown( + app_commands.Cooldown(1, self.cooldown_manager.cooldown_seconds), + retry_after, + ) + return True + + self.tree.interaction_check = app_command_cooldown_check + async def setup_hook(self) -> None: """Setup hook called before the bot starts.""" # Initialize CodeBuddy database @@ -230,23 +285,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/tests/test_cooldown.py b/tests/test_cooldown.py new file mode 100644 index 0000000..185c1c9 --- /dev/null +++ b/tests/test_cooldown.py @@ -0,0 +1,87 @@ +"""Tests for the command cooldown system.""" + +from __future__ import annotations + +import unittest +from unittest.mock import patch + +from utils.config import Config +from utils.cooldown import CooldownManager + + +class TestCooldownManager(unittest.TestCase): + def test_first_invocation_is_allowed(self) -> None: + manager = CooldownManager(cooldown_seconds=3.0) + self.assertIsNone(manager.trigger(12345)) + + def test_subsequent_invocation_within_cooldown_is_blocked(self) -> None: + manager = CooldownManager(cooldown_seconds=3.0) + + with patch("time.monotonic", side_effect=[100.0, 101.0, 101.0]): + manager.trigger(12345) + retry_after = manager.get_retry_after(12345) + self.assertIsNotNone(retry_after) + self.assertAlmostEqual(retry_after, 2.0, places=2) # type: ignore[arg-type] + + retry = manager.trigger(12345) + self.assertIsNotNone(retry) + self.assertAlmostEqual(retry, 2.0, places=2) # type: ignore[arg-type] + + def test_invocation_after_cooldown_window_is_allowed(self) -> None: + manager = CooldownManager(cooldown_seconds=3.0) + + with patch("time.monotonic", side_effect=[100.0, 103.5]): + manager.trigger(12345) + self.assertIsNone(manager.trigger(12345)) + + def test_different_users_have_independent_cooldowns(self) -> None: + manager = CooldownManager(cooldown_seconds=3.0) + + with patch("time.monotonic", side_effect=[100.0, 100.5, 101.0]): + self.assertIsNone(manager.trigger(111)) + self.assertIsNone(manager.trigger(222)) + self.assertIsNotNone(manager.trigger(111)) + + def test_zero_or_negative_cooldown_disables_rate_limiting(self) -> None: + zero_manager = CooldownManager(cooldown_seconds=0.0) + neg_manager = CooldownManager(cooldown_seconds=-1.0) + + self.assertIsNone(zero_manager.trigger(123)) + self.assertIsNone(zero_manager.trigger(123)) + self.assertIsNone(neg_manager.trigger(123)) + self.assertIsNone(neg_manager.trigger(123)) + + def test_reset_and_clear(self) -> None: + manager = CooldownManager(cooldown_seconds=5.0) + manager.trigger(100) + manager.trigger(200) + + manager.reset(100) + self.assertIsNone(manager.get_retry_after(100)) + self.assertIsNotNone(manager.get_retry_after(200)) + + manager.clear() + self.assertIsNone(manager.get_retry_after(200)) + + def test_pruning_expired_entries_when_cache_large(self) -> None: + manager = CooldownManager(cooldown_seconds=2.0) + # Pre-populate map with 5001 old entries + manager._last_used = {i: 10.0 for i in range(5001)} + + with patch("time.monotonic", return_value=50.0): + # Triggering a new user should prune the old entries older than 50.0 - 2.0 + manager.trigger(999999) + + # Only the new user should remain + self.assertEqual(len(manager._last_used), 1) + self.assertIn(999999, manager._last_used) + + +class TestCooldownConfig(unittest.TestCase): + def test_default_cooldown_config(self) -> None: + config = Config() + self.assertEqual(config.command_cooldown_seconds, 3.0) + + def test_custom_cooldown_config(self) -> None: + config = Config(command_cooldown_seconds=5.5) + self.assertEqual(config.command_cooldown_seconds, 5.5) diff --git a/utils/config.py b/utils/config.py index f7ba725..9b17e1a 100644 --- a/utils/config.py +++ b/utils/config.py @@ -22,6 +22,7 @@ class Config(BaseSettings): topgg_token: str | None = Field(default=None) topgg_webhook_secret: str | None = Field(default=None) redis_url: str | None = Field(default=None) + command_cooldown_seconds: float = Field(default=3.0) # CodeBuddy settings question_channel_id: int | None = Field(default=None) diff --git a/utils/cooldown.py b/utils/cooldown.py new file mode 100644 index 0000000..ba68bc4 --- /dev/null +++ b/utils/cooldown.py @@ -0,0 +1,59 @@ +"""Command cooldown manager to prevent command spam.""" + +from __future__ import annotations + +import time + + +class CooldownManager: + """Tracks per-user command execution timestamps to enforce a cooldown.""" + + def __init__(self, cooldown_seconds: float = 3.0) -> None: + self.cooldown_seconds = cooldown_seconds + self._last_used: dict[int, float] = {} + + def get_retry_after(self, user_id: int, now: float | None = None) -> float | None: + """Return remaining cooldown seconds if user is on cooldown, else None.""" + if self.cooldown_seconds <= 0: + return None + + last_time = self._last_used.get(user_id) + if last_time is None: + return None + + current_time = now if now is not None else time.monotonic() + elapsed = current_time - last_time + if elapsed < self.cooldown_seconds: + return self.cooldown_seconds - elapsed + + return None + + def trigger(self, user_id: int) -> float | None: + """Check if user is on cooldown. + + If on cooldown, returns remaining seconds. + If not on cooldown, records current timestamp and returns None. + """ + now = time.monotonic() + retry_after = self.get_retry_after(user_id, now=now) + if retry_after is not None: + return retry_after + + self._last_used[user_id] = now + + # Prevent unbounded memory growth by pruning expired entries when large + if len(self._last_used) > 5000: + cutoff = now - self.cooldown_seconds + self._last_used = { + uid: ts for uid, ts in self._last_used.items() if ts > cutoff + } + + return None + + def reset(self, user_id: int) -> None: + """Clear cooldown for a specific user.""" + self._last_used.pop(user_id, None) + + def clear(self) -> None: + """Clear all active cooldowns.""" + self._last_used.clear() From 9f430add07cbbde91b246b5a0db9cdbe18149c24 Mon Sep 17 00:00:00 2001 From: Saksham Date: Tue, 15 Sep 2026 21:33:00 +0530 Subject: [PATCH 2/2] Resolve CI action deprecation and test warnings --- .github/workflows/ci.yml | 4 ++-- mypy.ini | 3 +-- pytest.ini | 3 +++ tests/test_cooldown.py | 10 ++++++---- utils/cooldown.py | 4 +--- 5 files changed, 13 insertions(+), 11 deletions(-) create mode 100644 pytest.ini diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 70d1e1a..c113cdd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,10 +9,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.14" diff --git a/mypy.ini b/mypy.ini index 881f933..333da23 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,5 +1,4 @@ [mypy] python_version = 3.14 ignore_missing_imports = True -disable_error_code = arg-type, assignment, attr-defined, union-attr, var-annotated, override - +disable_error_code = arg-type, assignment, attr-defined, union-attr, var-annotated, override, annotation-unchecked diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a4a3a65 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +filterwarnings = + ignore::DeprecationWarning:discord.* diff --git a/tests/test_cooldown.py b/tests/test_cooldown.py index 185c1c9..ef9a8c1 100644 --- a/tests/test_cooldown.py +++ b/tests/test_cooldown.py @@ -20,12 +20,14 @@ def test_subsequent_invocation_within_cooldown_is_blocked(self) -> None: with patch("time.monotonic", side_effect=[100.0, 101.0, 101.0]): manager.trigger(12345) retry_after = manager.get_retry_after(12345) - self.assertIsNotNone(retry_after) - self.assertAlmostEqual(retry_after, 2.0, places=2) # type: ignore[arg-type] + if retry_after is None: + self.fail("Expected retry_after to be not None") + self.assertAlmostEqual(retry_after, 2.0, places=2) retry = manager.trigger(12345) - self.assertIsNotNone(retry) - self.assertAlmostEqual(retry, 2.0, places=2) # type: ignore[arg-type] + if retry is None: + self.fail("Expected retry to be not None") + self.assertAlmostEqual(retry, 2.0, places=2) def test_invocation_after_cooldown_window_is_allowed(self) -> None: manager = CooldownManager(cooldown_seconds=3.0) diff --git a/utils/cooldown.py b/utils/cooldown.py index ba68bc4..3b69e01 100644 --- a/utils/cooldown.py +++ b/utils/cooldown.py @@ -44,9 +44,7 @@ def trigger(self, user_id: int) -> float | None: # Prevent unbounded memory growth by pruning expired entries when large if len(self._last_used) > 5000: cutoff = now - self.cooldown_seconds - self._last_used = { - uid: ts for uid, ts in self._last_used.items() if ts > cutoff - } + self._last_used = {uid: ts for uid, ts in self._last_used.items() if ts > cutoff} return None