From 85a7967888f6d4afd3b1a798e452974f5b2b8384 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 23 Jun 2019 14:14:28 +0300
Subject: [PATCH] Add initial flood prevention

---
 base-config.yaml   |  4 ++++
 reactbot/bot.py    | 35 +++++++++++++++++++++++++++++++++--
 reactbot/config.py |  2 ++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/base-config.yaml b/base-config.yaml
index a523deb..27c6a78 100644
--- a/base-config.yaml
+++ b/base-config.yaml
@@ -29,6 +29,10 @@ templates:
 default_flags:
 - ignorecase
 
+antispam:
+  max: 2
+  delay: 60
+
 rules:
   twim_cookies:
     rooms: ["!FPUfgzXYWTKgIrwKxW:matrix.org"]
diff --git a/reactbot/bot.py b/reactbot/bot.py
index 03f57c7..eea1aea 100644
--- a/reactbot/bot.py
+++ b/reactbot/bot.py
@@ -13,9 +13,12 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-from typing import Type, Tuple
+from typing import Type, Tuple, Dict
+import time
 
-from mautrix.types import EventType, MessageType
+from attr import dataclass
+
+from mautrix.types import EventType, MessageType, UserID, RoomID
 from mautrix.util.config import BaseProxyConfig
 
 from maubot import Plugin, MessageEvent
@@ -24,8 +27,27 @@ from maubot.handlers import event
 from .config import Config, ConfigError
 
 
+@dataclass
+class FloodInfo:
+    rb: 'ReactBot'
+    count: int
+    last_message: int
+
+    def bump(self) -> bool:
+        now = int(time.time())
+        if self.last_message + self.rb.config["antispam.delay"] < now:
+            self.count = 0
+        self.count += 1
+        if self.count > self.rb.config["antispam.max"]:
+            return True
+        self.last_message = now
+        return False
+
+
 class ReactBot(Plugin):
     allowed_msgtypes: Tuple[MessageType, ...] = (MessageType.TEXT, MessageType.EMOTE)
+    user_flood: Dict[UserID, FloodInfo]
+    room_flood: Dict[RoomID, FloodInfo]
 
     @classmethod
     def get_config_class(cls) -> Type[BaseProxyConfig]:
@@ -33,6 +55,8 @@ class ReactBot(Plugin):
 
     async def start(self) -> None:
         await super().start()
+        self.user_flood = {}
+        self.room_flood = {}
         self.on_external_config_update()
 
     def on_external_config_update(self) -> None:
@@ -42,6 +66,11 @@ class ReactBot(Plugin):
         except ConfigError:
             self.log.exception("Failed to load config")
 
+    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()
+
     @event.on(EventType.ROOM_MESSAGE)
     async def event_handler(self, evt: MessageEvent) -> None:
         if evt.sender == self.client.mxid or evt.content.msgtype not in self.allowed_msgtypes:
@@ -49,6 +78,8 @@ class ReactBot(Plugin):
         for name, rule in self.config.rules.items():
             match = rule.match(evt)
             if match is not None:
+                if self.is_flood(evt):
+                    return
                 try:
                     await rule.execute(evt, match)
                 except Exception:
diff --git a/reactbot/config.py b/reactbot/config.py
index 754b51b..29d02b6 100644
--- a/reactbot/config.py
+++ b/reactbot/config.py
@@ -37,6 +37,8 @@ class Config(BaseProxyConfig):
         helper.copy("rules")
         helper.copy("templates")
         helper.copy("default_flags")
+        helper.copy("antispam.max")
+        helper.copy("antispam.delay")
 
     def parse_data(self) -> None:
         self.default_flags = re.RegexFlag(0)