From b26f9cf6ebd4472224266fbdd0f5cc1e84ff3971 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 23 Jun 2019 14:20:25 +0300
Subject: [PATCH] Separate per-user and per-room flood limits

---
 base-config.yaml   |  8 ++++++--
 reactbot/bot.py    | 24 ++++++++++++++++++------
 reactbot/config.py |  6 ++++--
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/base-config.yaml b/base-config.yaml
index 27c6a78..f99e872 100644
--- a/base-config.yaml
+++ b/base-config.yaml
@@ -30,8 +30,12 @@ default_flags:
 - ignorecase
 
 antispam:
-  max: 2
-  delay: 60
+  room:
+    max: 1
+    delay: 60
+  user:
+    max: 2
+    delay: 60
 
 rules:
   twim_cookies:
diff --git a/reactbot/bot.py b/reactbot/bot.py
index eea1aea..a0f17c8 100644
--- a/reactbot/bot.py
+++ b/reactbot/bot.py
@@ -29,16 +29,17 @@ from .config import Config, ConfigError
 
 @dataclass
 class FloodInfo:
-    rb: 'ReactBot'
+    max: int
+    delay: int
     count: int
     last_message: int
 
     def bump(self) -> bool:
         now = int(time.time())
-        if self.last_message + self.rb.config["antispam.delay"] < now:
+        if self.last_message + self.delay < now:
             self.count = 0
         self.count += 1
-        if self.count > self.rb.config["antispam.max"]:
+        if self.count > self.max:
             return True
         self.last_message = now
         return False
@@ -66,10 +67,21 @@ class ReactBot(Plugin):
         except ConfigError:
             self.log.exception("Failed to load config")
 
+    def _make_flood_info(self, for_type: str) -> 'FloodInfo':
+        return FloodInfo(max=self.config[f"antispam.{for_type}.max"],
+                         delay=self.config[f"antispam.{for_type}.delay"],
+                         count=0, last_message=0)
+
+    def _get_flood_info(self, flood_map: dict, key: str, for_type: str) -> 'FloodInfo':
+        try:
+            return flood_map[key]
+        except KeyError:
+            fi = flood_map[key] = self._make_flood_info(for_type)
+            return fi
+
     def is_flood(self, evt: MessageEvent) -> bool:
-        uf = self.user_flood.setdefault(evt.sender, FloodInfo(rb=self, count=0, last_message=0))
-        rf = self.room_flood.setdefault(evt.room_id, FloodInfo(rb=self, count=0, last_message=0))
-        return uf.bump() or rf.bump()
+        return (self._get_flood_info(self.user_flood, evt.sender, "user").bump()
+                or self._get_flood_info(self.room_flood, evt.room_id, "room").bump())
 
     @event.on(EventType.ROOM_MESSAGE)
     async def event_handler(self, evt: MessageEvent) -> None:
diff --git a/reactbot/config.py b/reactbot/config.py
index 29d02b6..379dead 100644
--- a/reactbot/config.py
+++ b/reactbot/config.py
@@ -37,8 +37,10 @@ class Config(BaseProxyConfig):
         helper.copy("rules")
         helper.copy("templates")
         helper.copy("default_flags")
-        helper.copy("antispam.max")
-        helper.copy("antispam.delay")
+        helper.copy("antispam.user.max")
+        helper.copy("antispam.user.delay")
+        helper.copy("antispam.room.max")
+        helper.copy("antispam.room.delay")
 
     def parse_data(self) -> None:
         self.default_flags = re.RegexFlag(0)