From 86111fb8ce2594468bea1935a855d23751068f4b Mon Sep 17 00:00:00 2001 From: Elijah Lazkani Date: Sat, 24 Feb 2018 23:40:44 -0500 Subject: [PATCH] Implementing a buffering system for flood protection --- boots/admin.py | 31 ++++---- boots/admin_commands.py | 1 + boots/boots.py | 41 +++++----- boots/eightball.py | 1 + boots/robot.py | 165 +++++++++++++++++++++++++++++++++++++++- 5 files changed, 201 insertions(+), 38 deletions(-) diff --git a/boots/admin.py b/boots/admin.py index 8dc2fb2..5f4972a 100644 --- a/boots/admin.py +++ b/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 - 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 - Log out from your account" - self.client.send("PRIVMSG", **kwargs) + self.client.msg(**kwargs) if match_help_cmd.group(1) == 'passwd': kwargs['message'] = \ "passwd - Change your" \ " account\'s password" - self.client.send("PRIVMSG", **kwargs) + self.client.msg(**kwargs) if match_help_cmd.group(1) == 'add': kwargs['message'] = \ "add - 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 - 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): """ diff --git a/boots/admin_commands.py b/boots/admin_commands.py index 8422f91..90ff5ab 100644 --- a/boots/admin_commands.py +++ b/boots/admin_commands.py @@ -3,6 +3,7 @@ import re import functools import logging import asyncio + import admin diff --git a/boots/boots.py b/boots/boots.py index 3597fa8..eb7153a 100644 --- a/boots/boots.py +++ b/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 -# @bot.on("USERMODE") -# def umode(**kwargs): -# logger.debug("USERMODE {}".format(kwargs)) + # Code below is an example + @bot.on("USERMODE") + def umode(**kwargs): + logger.debug("USERMODE {}".format(kwargs)) -# @bot.on("CHANNELMODE") -# def cmode(**kwargs): -# logger.debug("CHANNELMODE {}".format(kwargs)) + @bot.on("CHANNELMODE") + def cmode(**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): - bot.send("JOIN", channel=message) + async def join(target, message, **kwargs): + await bot.join(channel=message) @admin_cmd.on_command("part") - def part(target, message, **kwargs): - bot.send("PART", channel=message) + async def part(target, message, **kwargs): + 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'] = \ - "\x01ACTION {}\x01".format(" ".join(message.split(' ')[1:])) - bot.send("PRIVMSG", **kwargs) + kwargs['message'] = " ".join(message.split(' ')[1:]) + await bot.action(**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()) diff --git a/boots/eightball.py b/boots/eightball.py index c73f362..7ba6f09 100644 --- a/boots/eightball.py +++ b/boots/eightball.py @@ -4,6 +4,7 @@ import random import re import functools import asyncio + import robot diff --git a/boots/robot.py b/boots/robot.py index 11eaf6a..3eaeb0a 100644 --- a/boots/robot.py +++ b/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: - super().__init__(host=host, port=port, ssl=ssl) + channels: list = None, + 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))