Implementing a buffering system for flood protection

boots/admin.py

@@ -4,6 +4,7 @@ import logging
 import pickle
 import hashlib
 import aiofiles
+
 import robot
 
 
@@ -284,36 +285,36 @@ class Admin:
                             kwargs['message'] = \
                                 "login <user> <password> <level> - Login as " \
                                 "an admin with your account"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                         if match_help_cmd.group(1) == 'logout':
                             kwargs['message'] = \
                                 "logout <user> - Log out from your account"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                         if match_help_cmd.group(1) == 'passwd':
                             kwargs['message'] = \
                                 "passwd <new password> - Change your" \
                                 " account\'s password"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                         if match_help_cmd.group(1) == 'add':
                             kwargs['message'] = \
                                 "add <user> <password> <level> - adds" \
                                 " an admin account to the list of admins" \
                                 " with provided level"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                         if match_help_cmd.group(1) == 'rm':
                             kwargs['message'] = \
                                 "rm <user> - removes an admin from the list" \
                                 " of admins"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                         if match_help_cmd.group(1) == 'list':
                             kwargs['message'] = "list - lists all the admins"
-                            self.client.send("PRIVMSG", **kwargs)
+                            self.client.msg(**kwargs)
                 elif match_help:
                     kwargs['message'] = "help [command]"
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
                     kwargs['message'] = \
                         "commands: login logout passwd add rm list"
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
 
     async def log_in(self,
                      nick: str,
@@ -369,7 +370,7 @@ class Admin:
                             nick, match.group(1))
                         self.logger.debug("We have logged in {} successfully, "
                                           "notifying".format(nick))
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
                     self.logger.debug("We are calling save_config()")
                     await self.save_config()
 
@@ -402,7 +403,7 @@ class Admin:
                                         " successfully".format(nick, admin)
                     self.logger.debug("We have successfully logged {}"
                                       " out, notifying".format(nick))
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
                     self.logger.debug("We are calling save_config")
                     await self.save_config()
 
@@ -440,7 +441,7 @@ class Admin:
                     self.logger.debug(
                         "We have successfully changed {}'s password,"
                         " notifying".format(nick))
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
                     self.logger.debug("We are calling save_config()")
                     await self.save_config()
                     kwargs['target'] = self.client.nick
@@ -488,7 +489,7 @@ class Admin:
                         self.logger.warn(
                             "We detected that {} has already been added,"
                             " notifying {}".format(match.group(1), nick))
-                    self.client.send("PRIVMSG", **kwargs)
+                    self.client.msg(**kwargs)
 
     async def admin_rm(self,
                        nick: str,
@@ -542,7 +543,7 @@ class Admin:
                                 "We detected that {0} does not have enough"
                                 " access to delete {1}, notifying {0}".format(
                                   nick, match.group(1)))
-                    self.client.send("PRIVMSG", **kwags)
+                    self.client.msg(**kwags)
 
     def admin_list(self, nick: str, target: str, **kwargs) -> None:
         """
@@ -566,11 +567,11 @@ class Admin:
                             key, self.admins[key]['level'])
                 kwargs['target'] = nick
                 kwargs['message'] = "List of Administrators:"
-                self.client.send("PRIVMSG", **kwargs)
+                self.client.msg(**kwargs)
                 kwargs['message'] = admins
                 self.logger.debug("We are returning admin list page to"
                                   " {}".format(kwargs))
-                self.client.send("PRIVMSG", **kwargs)
+                self.client.msg(**kwargs)
 
     def is_admin(self, user: str, host: str):
         """

boots/admin_commands.py

@@ -3,6 +3,7 @@ import re
 import functools
 import logging
 import asyncio
+
 import admin
 
 

boots/boots.py

@@ -21,21 +21,21 @@ async def plugins(bot: robot.Bot):
         bot.send("PRIVMSG", target=target, message=message)
         await bot.wait("NOTICE")
 
-#    Code below will not work, it is awaiting a PR upstream
+    # Code below is an example
-#    @bot.on("USERMODE")
+    @bot.on("USERMODE")
-#    def umode(**kwargs):
+    def umode(**kwargs):
-#        logger.debug("USERMODE {}".format(kwargs))
+       logger.debug("USERMODE {}".format(kwargs))
 
-#    @bot.on("CHANNELMODE")
+    @bot.on("CHANNELMODE")
-#    def cmode(**kwargs):
+    def cmode(**kwargs):
-#        logger.debug("CHANNELMODE {}".format(kwargs))
+       logger.debug("CHANNELMODE {}".format(kwargs))
 
     # 8Ball
     magic_ball = eightball.EightBall(bot)
 
     @magic_ball.on_keyword
     async def ball(target, message, **kwargs):
-        bot.send("PRIVMSG", target=target, message=message)
+       await bot.msg(target=target, message=message)
 
     administrator = admin.Admin(bot)
     await administrator.init()
@@ -43,29 +43,28 @@ async def plugins(bot: robot.Bot):
     admin_cmd = admin_commands.AdminCmd(administrator)
 
     @admin_cmd.on_command("join")
-    def join(target, message, **kwargs):
+    async def join(target, message, **kwargs):
-        bot.send("JOIN", channel=message)
+        await bot.join(channel=message)
 
     @admin_cmd.on_command("part")
-    def part(target, message, **kwargs):
+    async def part(target, message, **kwargs):
-        bot.send("PART", channel=message)
+        await bot.part(channel=message)
 
     @admin_cmd.on_command("msg")
-    def msg(target, message, **kwargs):
+    async def msg(target, message, **kwargs):
         kwargs['target'] = message.split(' ')[0]
         kwargs['message'] = " ".join(message.split(' ')[1:])
-        bot.send("PRIVMSG", **kwargs)
+        await bot.msg(**kwargs)
 
     @admin_cmd.on_command("action")
-    def action(target, message, **kwargs):
+    async def action(target, message, **kwargs):
         kwargs['target'] = message.split(' ')[0]
-        kwargs['message'] = \
+        kwargs['message'] = " ".join(message.split(' ')[1:])
-            "\x01ACTION {}\x01".format(" ".join(message.split(' ')[1:]))
+        await bot.action(**kwargs)
-        bot.send("PRIVMSG", **kwargs)
 
     @admin_cmd.on_command("quit")
     async def quit(target, message, **kwargs):
-        bot.send("QUIT", message=message)
+        await bot.quit(message=message)
         await bot.disconnect()
         # Exit the event loop cleanly
         bot.loop.stop()
@@ -78,11 +77,11 @@ def main():
     ssl = False
 
     nick = "LuckyBoots"
-    channel = "#boots"
+    channel = ["#boots"]
 
     bot = robot.Bot(host=host, port=port,
                     ssl=ssl, nick=nick,
-                    channels=[channel])
+                    channels=channel)
 
     bot.loop.create_task(bot.connect())
 

boots/eightball.py

@@ -4,6 +4,7 @@ import random
 import re
 import functools
 import asyncio
+
 import robot
 
 

boots/robot.py

@@ -1,7 +1,9 @@
 import logging
 import asyncio
 import types
+import typing
 import functools
+import datetime
 import bottom
 
 
@@ -17,8 +19,15 @@ class Bot(bottom.Client):
                  nick: str,
                  user: str = None,
                  realname: str = None,
-                 channels: list = None) -> None:
+                 channels: list = None,
-        super().__init__(host=host, port=port, ssl=ssl)
+                 msg_flood_size_threshold = 3,
+                 msg_flood_time_threshold = 5,
+                 action_flood_size_threshold=3,
+                 action_flood_time_threshold=5,
+                 global_flood_size_threshold=3,
+                 global_flood_time_threshold=5,
+                 loop: typing.Optional[asyncio.AbstractEventLoop] = None) -> None:
+        super().__init__(host=host, port=port, ssl=ssl, loop=loop)
         self.logger = logging.getLogger(self.__class__.__name__)
         self.logger.debug("Initializing...")
         self.nick = nick
@@ -26,6 +35,15 @@ class Bot(bottom.Client):
         self.realname = realname or self.nick
         self.channels = channels or []
         self.pre_pub = []
+        self.msg_flood_size_threshold = msg_flood_size_threshold
+        self.msg_flood_time_threshold = msg_flood_time_threshold
+        self.action_flood_size_threshold = action_flood_size_threshold
+        self.action_flood_time_threshold = action_flood_time_threshold
+        self.global_flood_size_threshold = global_flood_size_threshold
+        self.global_flood_time_threshold = global_flood_time_threshold
+        self.msg_timestamp_queue = []
+        self.action_timestamp_queue = []
+        self.global_timestamp_queue = []
         self.on("ping", self.keepalive)
         self.on("CLIENT_CONNECT", self.on_connect)
         self.on("PRIVMSG", self.privmsg)
@@ -103,6 +121,149 @@ class Bot(bottom.Client):
         self.pre_pub.extend([wrapped])
         return func
 
+    async def global_buffer(self, command: str, **kwargs) -> None:
+        """
+        This method will buffer the communication sent to the server
+        to keep the client from getting disconnected
+
+        It uses global_* variables as configuration
+
+        NOTE: this method is to be called *only* by other buffer methods
+
+        :param command str: the command to send to the server
+        :param kwargs: the information required for the command
+        :return None
+        """
+        timestamp = datetime.datetime.now()
+        if self.global_timestamp_queue.__len__() < self.global_flood_size_threshold:
+            self.global_timestamp_queue.extend([timestamp])
+        else:
+            del self.global_timestamp_queue[0]
+            self.global_timestamp_queue.extend([timestamp])
+
+        time_diff = self.global_timestamp_queue[-1] - self.global_timestamp_queue[0]
+
+        if len(self.global_timestamp_queue) == self.global_flood_size_threshold and \
+           time_diff.total_seconds() < self.global_flood_time_threshold:
+            self.logger.info("Waiting {}s".format(self.global_flood_time_threshold - time_diff.total_seconds()))
+            await asyncio.sleep(self.global_flood_time_threshold - time_diff.total_seconds())
+
+        self.send(command, **kwargs)
+
+    async def msg_buffer(self, command: str, **kwargs) -> None:
+        """
+        This method will buffer the messages sent to the server
+        to keep the client from getting disconnected
+
+        messages are of type: PRIVMSG which includes ACTION
+
+        It uses msg_* variables as configuration
+
+        NOTE: this method is to be called *only* by bot methods
+
+        :param command str: the message command to send to the server
+        :param kwargs: the information required for the message command
+        :return None
+        """
+        timestamp = datetime.datetime.now()
+        if self.msg_timestamp_queue.__len__() < self.msg_flood_size_threshold:
+            self.msg_timestamp_queue.extend([timestamp])
+        else:
+            del self.msg_timestamp_queue[0]
+            self.msg_timestamp_queue.extend([timestamp])
+
+        time_diff = self.msg_timestamp_queue[-1] - self.msg_timestamp_queue[0]
+
+        if len(self.msg_timestamp_queue) == self.msg_flood_size_threshold and \
+           time_diff.total_seconds() < self.msg_flood_time_threshold:
+            self.logger.info("Waiting {}s".format(self.msg_flood_time_threshold - time_diff.total_seconds()))
+            await asyncio.sleep(self.msg_flood_time_threshold - time_diff.total_seconds())
+
+        await self.global_buffer(command, **kwargs)
+
+    async def action_buffer(self, command: str, **kwargs):
+        """
+        This method will buffer the actions sent to the server
+        to keep the client from getting disconnected
+
+        actions are of type: JOIN, PART and QUIT
+
+        It uses action_* variables as configuration
+
+        NOTE: this method is to be called *only* by bot methods
+
+        :param command str: the action command to send to the server
+        :param kwargs: the information required for the action command
+        :return None
+        """
+        timestamp = datetime.datetime.now()
+        if self.action_timestamp_queue.__len__() < self.action_flood_size_threshold:
+            self.action_timestamp_queue.extend([timestamp])
+        else:
+            del self.action_timestamp_queue[0]
+            self.action_timestamp_queue.extend([timestamp])
+
+        time_diff = self.action_timestamp_queue[-1] - self.action_timestamp_queue[0]
+
+        if len(self.action_timestamp_queue) == self.action_flood_size_threshold and \
+           time_diff.total_seconds() < self.action_flood_time_threshold:
+            self.logger.info("Waiting {}s".format(self.action_flood_time_threshold - time_diff.total_seconds()))
+            await asyncio.sleep(self.action_flood_time_threshold - time_diff.total_seconds())
+
+        await self.global_buffer(command, **kwargs)
+
+    async def msg(self, **kwargs) -> None:
+        """
+        This method will send a private message to the messages buffer
+
+        :param kwargs: the information required for the command
+        :return None
+        """
+        await self.msg_buffer("PRIVMSG", **kwargs)
+
+    async def action(self, **kwargs) -> None:
+        """
+        This method will send an action message to the messages buffer
+
+        :param kwargs: the information required for the command
+        :return None
+        """
+        if kwargs.get('message', None):
+            kwargs['message'] = \
+                "\x01ACTION {}\x01".format(kwargs['message'])
+
+            await self.msg_buffer("PRIVMSG", **kwargs)
+        else:
+            # TODO: better error handling for the future
+            self.logger.error("ACTION does not have a message\n{}".format(kwargs))
+
+    async def join(self, **kwargs) -> None:
+        """
+        This method will send a join command to the action buffer
+
+        :param kwargs: the information required for the command
+        :return None
+        """
+        await self.action_buffer("JOIN", **kwargs)
+
+    async def part(self, **kwargs) -> None:
+        """
+        This method will send a part command to the action buffer
+
+        :param kwargs: the information required for the command
+        :return None
+        """
+        await self.action_buffer("PART", **kwargs)
+
+    async def quit(self, **kwargs) -> None:
+        """
+        This method will send a quit command to the action buffer
+
+        :param kwargs: the information required for the command
+        :return None
+        """
+        await self.action_buffer("QUIT", **kwargs)
+
     def privmsg(self, **kwargs) -> None:
         self.logger.debug("PRIVMSG {}".format(kwargs))