Merge pull request #155 from TotallyNotRobots/regex-rate-limit

Rate limit regex hooks, fixes #153

Author: linuxdaemon

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Refactor minecraft_ping plugin for updated mcstatus library
 - Expand youtube.py error information
 - Handle 'a' vs 'an' in drinks plugin
+- Apply rate limiting to regex hooks
 ### Fixed
 - Fix matching exception in horoscope test
 - Fix youtube.py ISO time parse

plugins/core/core_sieve.py

@@ -1,7 +1,11 @@
 import logging
 from time import time
+from typing import Optional
 
 from cloudbot import hook
+from cloudbot.bot import CloudBot
+from cloudbot.event import CommandEvent, Event
+from cloudbot.plugin_hooks import Hook
 from cloudbot.util.tokenbucket import TokenBucket
 
 ready = False
@@ -16,29 +20,36 @@ def task_clear():
             del buckets[uid]
 
 
-@hook.sieve(priority=100)
-async def sieve_suite(bot, event, _hook):
+# noinspection PyUnusedLocal
+@hook.sieve()
+def check_acls(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
+    """
+    Handle config ACLs
+    """
     conn = event.conn
 
     # check acls
-    acl = conn.config.get('acls', {}).get(_hook.function_name)
-    if acl:
-        if 'deny-except' in acl:
-            allowed_channels = list(map(str.lower, acl['deny-except']))
-            if event.chan.lower() not in allowed_channels:
-                return None
-        if 'allow-except' in acl:
-            denied_channels = list(map(str.lower, acl['allow-except']))
-            if event.chan.lower() in denied_channels:
-                return None
-
-    # check disabled_commands
-    if _hook.type == "command":
-        disabled_commands = conn.config.get('disabled_commands', [])
-        if event.triggered_command in disabled_commands:
+    acl = conn.config.get("acls", {}).get(_hook.function_name, {})
+    allowlist = acl.get("deny-except")
+    denylist = acl.get("allow-except")
+    chan = event.chan.lower()
+    if allowlist is not None:
+        allowed_channels = list(map(str.lower, allowlist))
+        if chan not in allowed_channels:
+            return None
+
+    if denylist is not None:
+        denied_channels = list(map(str.lower, denylist))
+        if chan in denied_channels:
             return None
 
-    # check permissions
+    return event
+
+
+# noinspection PyUnusedLocal
+@hook.sieve()
+async def perm_sieve(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
+    """check permissions"""
     allowed_permissions = _hook.permissions
     if allowed_permissions:
         allowed = False
@@ -51,33 +62,58 @@ async def sieve_suite(bot, event, _hook):
             event.notice("Sorry, you are not allowed to use this command.")
             return None
 
-    # check command spam tokens
+    return event
+
+
+# noinspection PyUnusedLocal
+@hook.sieve()
+def check_disabled(bot: CloudBot, event: CommandEvent, _hook: Hook) -> Optional[Event]:
+    """
+    check disabled_commands
+    """
+    conn = event.conn
     if _hook.type == "command":
+        disabled_commands = conn.config.get("disabled_commands", [])
+        if event.triggered_command in disabled_commands:
+            return None
+
+    return event
+
+
+# noinspection PyUnusedLocal
+@hook.sieve()
+def rate_limit(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
+    """
+    Handle rate limiting certain hooks
+    """
+    conn = event.conn
+    # check command spam tokens
+    if _hook.type in ("command", "regex"):
         uid = "!".join([conn.name, event.chan, event.nick]).lower()
 
-        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
-        restore_rate = conn.config.get('ratelimit', {}).get('restore_rate', 2.5)
-        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
-        strict = conn.config.get('ratelimit', {}).get('strict', True)
-
-        if uid not in buckets:
-            bucket = TokenBucket(tokens, restore_rate)
-            bucket.consume(message_cost)
-            buckets[uid] = bucket
-            return event
-
-        bucket = buckets[uid]
-        if bucket.consume(message_cost):
-            pass
-        else:
+        config = conn.config.get("ratelimit", {})
+        tokens = config.get("tokens", 17.5)
+        restore_rate = config.get("restore_rate", 2.5)
+        message_cost = config.get("message_cost", 5)
+        strict = config.get("strict", True)
+
+        try:
+            bucket = buckets[uid]
+        except KeyError:
+            buckets[uid] = bucket = TokenBucket(tokens, restore_rate)
+
+        if not bucket.consume(message_cost):
             logger.info(
-                "[%s|sieve] Refused command from %s. "
-                "Entity had %s tokens, needed %s.",
-                conn.name, uid, bucket.tokens, message_cost
+                "[%s|sieve] Refused command from %s. Entity had %s tokens, needed %s.",
+                conn.name,
+                uid,
+                bucket.tokens,
+                message_cost,
             )
             if strict:
                 # bad person loses all tokens
                 bucket.empty()
+
             return None
 
     return event

tests/plugin_tests/core_sieve_test.py

@@ -0,0 +1,147 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from cloudbot.event import CommandEvent, RegexEvent
+from cloudbot.util.tokenbucket import TokenBucket
+from plugins.core import core_sieve
+
+
+# noinspection PyUnusedFunction
+@pytest.fixture(autouse=True)
+def reset_buckets() -> None:
+    """
+    Clear bucket data after a test
+    """
+    try:
+        yield
+    finally:
+        core_sieve.buckets.clear()
+
+
+def test_rate_limit_command() -> None:
+    conn = MagicMock()
+    conn.name = "foobarconn"
+    conn.config = {}
+    conn.bot = MagicMock()
+
+    hook = MagicMock()
+    hook.type = "command"
+    event = CommandEvent(
+        text="bar",
+        cmd_prefix=".",
+        triggered_command="foo",
+        hook=hook,
+        bot=conn.bot,
+        conn=conn,
+        channel="#foo",
+        nick="foobaruser",
+    )
+    for _ in range(3):
+        res = core_sieve.rate_limit(event.bot, event, event.hook)
+        assert res is event
+
+    res = core_sieve.rate_limit(event.bot, event, event.hook)
+    assert res is None
+
+
+def test_rate_limit_regex() -> None:
+    conn = MagicMock()
+    conn.name = "foobarconn"
+    conn.config = {}
+    conn.bot = MagicMock()
+
+    hook = MagicMock()
+    hook.type = "regex"
+    event = RegexEvent(
+        hook=hook,
+        bot=conn.bot,
+        conn=conn,
+        channel="#foo",
+        nick="foobaruser",
+        match=MagicMock(),
+    )
+    for _ in range(3):
+        res = core_sieve.rate_limit(event.bot, event, event.hook)
+        assert res is event
+
+    res = core_sieve.rate_limit(event.bot, event, event.hook)
+    assert res is None
+
+
+def test_task_clear() -> None:
+    core_sieve.buckets["a"] = bucket = TokenBucket(10, 2)
+    bucket.timestamp = 0
+    assert len(core_sieve.buckets) == 1
+    core_sieve.task_clear()
+    assert len(core_sieve.buckets) == 0
+
+
+@pytest.mark.parametrize(
+    "config,allowed",
+    [
+        ({"foo": {"allow-except": ["#foo"]}}, False),
+        ({"foo": {"deny-except": ["#bar"]}}, False),
+        ({"foo": {"deny-except": ["#foo"]}}, True),
+        ({"foo": {"allow-except": ["#bar"]}}, True),
+        ({"bar": {"deny-except": ["#bar"]}}, True),
+        ({"foo": {"allow-except": []}}, True),
+        ({"foo": {"deny-except": []}}, False),
+    ],
+)
+def test_check_acls(config, allowed) -> None:
+    event = make_command_event()
+    event.conn.config["acls"] = config
+    res = core_sieve.check_acls(event.bot, event, event.hook)
+    if allowed:
+        assert res is event
+    else:
+        assert res is None
+
+
+@pytest.mark.asyncio()
+async def test_permissions() -> None:
+    event = make_command_event()
+    event.hook.permissions = ["admin"]
+
+    event.has_permission = perm = MagicMock()
+    res = await core_sieve.perm_sieve(event.bot, event, event.hook)
+    assert res is event
+
+    perm.return_value = False
+    res = await core_sieve.perm_sieve(event.bot, event, event.hook)
+    assert res is None
+
+
+def make_command_event():
+    conn = MagicMock()
+    conn.name = "foobarconn"
+    conn.config = {}
+    conn.bot = MagicMock()
+
+    hook = MagicMock()
+    hook.type = "command"
+    hook.function_name = "foo"
+    event = CommandEvent(
+        text="bar",
+        cmd_prefix=".",
+        triggered_command="foo",
+        hook=hook,
+        bot=conn.bot,
+        conn=conn,
+        channel="#foo",
+        nick="foobaruser",
+        user="user",
+        host="host",
+    )
+    return event
+
+
+def test_disabled():
+    event = make_command_event()
+    event.conn.config["disabled_commands"] = [event.triggered_command]
+    assert core_sieve.check_disabled(event.bot, event, event.hook) is None
+    event.conn.config["disabled_commands"] = ["random"]
+    assert core_sieve.check_disabled(event.bot, event, event.hook) is event
+    event.conn.config["disabled_commands"] = []
+    assert core_sieve.check_disabled(event.bot, event, event.hook) is event