Skip to content

Commit ccea19f

Browse files
committed
Rate limit regex hooks
1 parent 1918109 commit ccea19f

File tree

3 files changed

+221
-37
lines changed

3 files changed

+221
-37
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1515
- Refactor minecraft_ping plugin for updated mcstatus library
1616
- Expand youtube.py error information
1717
- Handle 'a' vs 'an' in drinks plugin
18+
- Apply rate limiting to regex hooks
1819
### Fixed
1920
- Fix matching exception in horoscope test
2021
- Fix youtube.py ISO time parse

plugins/core/core_sieve.py

Lines changed: 73 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
import logging
22
from time import time
3+
from typing import Optional
34

45
from cloudbot import hook
6+
from cloudbot.bot import CloudBot
7+
from cloudbot.event import CommandEvent, Event
8+
from cloudbot.plugin_hooks import Hook
59
from cloudbot.util.tokenbucket import TokenBucket
610

711
ready = False
@@ -16,29 +20,36 @@ def task_clear():
1620
del buckets[uid]
1721

1822

19-
@hook.sieve(priority=100)
20-
async def sieve_suite(bot, event, _hook):
23+
# noinspection PyUnusedLocal
24+
@hook.sieve()
25+
def check_acls(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
26+
"""
27+
Handle config ACLs
28+
"""
2129
conn = event.conn
2230

2331
# check acls
24-
acl = conn.config.get('acls', {}).get(_hook.function_name)
25-
if acl:
26-
if 'deny-except' in acl:
27-
allowed_channels = list(map(str.lower, acl['deny-except']))
28-
if event.chan.lower() not in allowed_channels:
29-
return None
30-
if 'allow-except' in acl:
31-
denied_channels = list(map(str.lower, acl['allow-except']))
32-
if event.chan.lower() in denied_channels:
33-
return None
34-
35-
# check disabled_commands
36-
if _hook.type == "command":
37-
disabled_commands = conn.config.get('disabled_commands', [])
38-
if event.triggered_command in disabled_commands:
32+
acl = conn.config.get("acls", {}).get(_hook.function_name, {})
33+
allowlist = acl.get("deny-except")
34+
denylist = acl.get("allow-except")
35+
chan = event.chan.lower()
36+
if allowlist is not None:
37+
allowed_channels = list(map(str.lower, allowlist))
38+
if chan not in allowed_channels:
39+
return None
40+
41+
if denylist is not None:
42+
denied_channels = list(map(str.lower, denylist))
43+
if chan in denied_channels:
3944
return None
4045

41-
# check permissions
46+
return event
47+
48+
49+
# noinspection PyUnusedLocal
50+
@hook.sieve()
51+
async def perm_sieve(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
52+
"""check permissions"""
4253
allowed_permissions = _hook.permissions
4354
if allowed_permissions:
4455
allowed = False
@@ -51,33 +62,58 @@ async def sieve_suite(bot, event, _hook):
5162
event.notice("Sorry, you are not allowed to use this command.")
5263
return None
5364

54-
# check command spam tokens
65+
return event
66+
67+
68+
# noinspection PyUnusedLocal
69+
@hook.sieve()
70+
def check_disabled(bot: CloudBot, event: CommandEvent, _hook: Hook) -> Optional[Event]:
71+
"""
72+
check disabled_commands
73+
"""
74+
conn = event.conn
5575
if _hook.type == "command":
76+
disabled_commands = conn.config.get("disabled_commands", [])
77+
if event.triggered_command in disabled_commands:
78+
return None
79+
80+
return event
81+
82+
83+
# noinspection PyUnusedLocal
84+
@hook.sieve()
85+
def rate_limit(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
86+
"""
87+
Handle rate limiting certain hooks
88+
"""
89+
conn = event.conn
90+
# check command spam tokens
91+
if _hook.type in ("command", "regex"):
5692
uid = "!".join([conn.name, event.chan, event.nick]).lower()
5793

58-
tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
59-
restore_rate = conn.config.get('ratelimit', {}).get('restore_rate', 2.5)
60-
message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
61-
strict = conn.config.get('ratelimit', {}).get('strict', True)
62-
63-
if uid not in buckets:
64-
bucket = TokenBucket(tokens, restore_rate)
65-
bucket.consume(message_cost)
66-
buckets[uid] = bucket
67-
return event
68-
69-
bucket = buckets[uid]
70-
if bucket.consume(message_cost):
71-
pass
72-
else:
94+
config = conn.config.get("ratelimit", {})
95+
tokens = config.get("tokens", 17.5)
96+
restore_rate = config.get("restore_rate", 2.5)
97+
message_cost = config.get("message_cost", 5)
98+
strict = config.get("strict", True)
99+
100+
try:
101+
bucket = buckets[uid]
102+
except KeyError:
103+
buckets[uid] = bucket = TokenBucket(tokens, restore_rate)
104+
105+
if not bucket.consume(message_cost):
73106
logger.info(
74-
"[%s|sieve] Refused command from %s. "
75-
"Entity had %s tokens, needed %s.",
76-
conn.name, uid, bucket.tokens, message_cost
107+
"[%s|sieve] Refused command from %s. Entity had %s tokens, needed %s.",
108+
conn.name,
109+
uid,
110+
bucket.tokens,
111+
message_cost,
77112
)
78113
if strict:
79114
# bad person loses all tokens
80115
bucket.empty()
116+
81117
return None
82118

83119
return event

tests/plugin_tests/core_sieve_test.py

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
from unittest.mock import MagicMock
2+
3+
import pytest
4+
5+
from cloudbot.event import CommandEvent, RegexEvent
6+
from cloudbot.util.tokenbucket import TokenBucket
7+
from plugins.core import core_sieve
8+
9+
10+
# noinspection PyUnusedFunction
11+
@pytest.fixture(autouse=True)
12+
def reset_buckets() -> None:
13+
"""
14+
Clear bucket data after a test
15+
"""
16+
try:
17+
yield
18+
finally:
19+
core_sieve.buckets.clear()
20+
21+
22+
def test_rate_limit_command() -> None:
23+
conn = MagicMock()
24+
conn.name = "foobarconn"
25+
conn.config = {}
26+
conn.bot = MagicMock()
27+
28+
hook = MagicMock()
29+
hook.type = "command"
30+
event = CommandEvent(
31+
text="bar",
32+
cmd_prefix=".",
33+
triggered_command="foo",
34+
hook=hook,
35+
bot=conn.bot,
36+
conn=conn,
37+
channel="#foo",
38+
nick="foobaruser",
39+
)
40+
for _ in range(3):
41+
res = core_sieve.rate_limit(event.bot, event, event.hook)
42+
assert res is event
43+
44+
res = core_sieve.rate_limit(event.bot, event, event.hook)
45+
assert res is None
46+
47+
48+
def test_rate_limit_regex() -> None:
49+
conn = MagicMock()
50+
conn.name = "foobarconn"
51+
conn.config = {}
52+
conn.bot = MagicMock()
53+
54+
hook = MagicMock()
55+
hook.type = "regex"
56+
event = RegexEvent(
57+
hook=hook,
58+
bot=conn.bot,
59+
conn=conn,
60+
channel="#foo",
61+
nick="foobaruser",
62+
match=MagicMock(),
63+
)
64+
for _ in range(3):
65+
res = core_sieve.rate_limit(event.bot, event, event.hook)
66+
assert res is event
67+
68+
res = core_sieve.rate_limit(event.bot, event, event.hook)
69+
assert res is None
70+
71+
72+
def test_task_clear() -> None:
73+
core_sieve.buckets["a"] = bucket = TokenBucket(10, 2)
74+
bucket.timestamp = 0
75+
assert len(core_sieve.buckets) == 1
76+
core_sieve.task_clear()
77+
assert len(core_sieve.buckets) == 0
78+
79+
80+
@pytest.mark.parametrize(
81+
"config,allowed",
82+
[
83+
({"foo": {"allow-except": ["#foo"]}}, False),
84+
({"foo": {"deny-except": ["#bar"]}}, False),
85+
({"foo": {"deny-except": ["#foo"]}}, True),
86+
({"foo": {"allow-except": ["#bar"]}}, True),
87+
({"bar": {"deny-except": ["#bar"]}}, True),
88+
({"foo": {"allow-except": []}}, True),
89+
({"foo": {"deny-except": []}}, False),
90+
],
91+
)
92+
def test_check_acls(config, allowed) -> None:
93+
event = make_command_event()
94+
event.conn.config["acls"] = config
95+
res = core_sieve.check_acls(event.bot, event, event.hook)
96+
if allowed:
97+
assert res is event
98+
else:
99+
assert res is None
100+
101+
102+
@pytest.mark.asyncio()
103+
async def test_permissions() -> None:
104+
event = make_command_event()
105+
event.hook.permissions = ["admin"]
106+
107+
event.has_permission = perm = MagicMock()
108+
res = await core_sieve.perm_sieve(event.bot, event, event.hook)
109+
assert res is event
110+
111+
perm.return_value = False
112+
res = await core_sieve.perm_sieve(event.bot, event, event.hook)
113+
assert res is None
114+
115+
116+
def make_command_event():
117+
conn = MagicMock()
118+
conn.name = "foobarconn"
119+
conn.config = {}
120+
conn.bot = MagicMock()
121+
122+
hook = MagicMock()
123+
hook.type = "command"
124+
hook.function_name = "foo"
125+
event = CommandEvent(
126+
text="bar",
127+
cmd_prefix=".",
128+
triggered_command="foo",
129+
hook=hook,
130+
bot=conn.bot,
131+
conn=conn,
132+
channel="#foo",
133+
nick="foobaruser",
134+
user="user",
135+
host="host",
136+
)
137+
return event
138+
139+
140+
def test_disabled():
141+
event = make_command_event()
142+
event.conn.config["disabled_commands"] = [event.triggered_command]
143+
assert core_sieve.check_disabled(event.bot, event, event.hook) is None
144+
event.conn.config["disabled_commands"] = ["random"]
145+
assert core_sieve.check_disabled(event.bot, event, event.hook) is event
146+
event.conn.config["disabled_commands"] = []
147+
assert core_sieve.check_disabled(event.bot, event, event.hook) is event

0 commit comments

Comments
 (0)