Implementing a buffering system for flood protection

This commit is contained in:
Elijah Lazkani 2018-02-24 23:40:44 -05:00
parent 7da5650153
commit 86111fb8ce
5 changed files with 201 additions and 38 deletions

View file

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

View file

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

View file

@ -21,21 +21,21 @@ async def plugins(bot: robot.Bot):
bot.send("PRIVMSG", target=target, message=message) bot.send("PRIVMSG", target=target, message=message)
await bot.wait("NOTICE") 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 # 8Ball
magic_ball = eightball.EightBall(bot) magic_ball = eightball.EightBall(bot)
@magic_ball.on_keyword @magic_ball.on_keyword
async def ball(target, message, **kwargs): async def ball(target, message, **kwargs):
bot.send("PRIVMSG", target=target, message=message) await bot.msg(target=target, message=message)
administrator = admin.Admin(bot) administrator = admin.Admin(bot)
await administrator.init() await administrator.init()
@ -43,29 +43,28 @@ async def plugins(bot: robot.Bot):
admin_cmd = admin_commands.AdminCmd(administrator) admin_cmd = admin_commands.AdminCmd(administrator)
@admin_cmd.on_command("join") @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") @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") @admin_cmd.on_command("msg")
def msg(target, message, **kwargs): async def msg(target, message, **kwargs):
kwargs['target'] = message.split(' ')[0] kwargs['target'] = message.split(' ')[0]
kwargs['message'] = " ".join(message.split(' ')[1:]) kwargs['message'] = " ".join(message.split(' ')[1:])
bot.send("PRIVMSG", **kwargs) await bot.msg(**kwargs)
@admin_cmd.on_command("action") @admin_cmd.on_command("action")
def action(target, message, **kwargs): async def action(target, message, **kwargs):
kwargs['target'] = message.split(' ')[0] 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") @admin_cmd.on_command("quit")
async def quit(target, message, **kwargs): async def quit(target, message, **kwargs):
bot.send("QUIT", message=message) await bot.quit(message=message)
await bot.disconnect() await bot.disconnect()
# Exit the event loop cleanly # Exit the event loop cleanly
bot.loop.stop() bot.loop.stop()
@ -78,11 +77,11 @@ def main():
ssl = False ssl = False
nick = "LuckyBoots" nick = "LuckyBoots"
channel = "#boots" channel = ["#boots"]
bot = robot.Bot(host=host, port=port, bot = robot.Bot(host=host, port=port,
ssl=ssl, nick=nick, ssl=ssl, nick=nick,
channels=[channel]) channels=channel)
bot.loop.create_task(bot.connect()) bot.loop.create_task(bot.connect())

View file

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

View file

@ -1,7 +1,9 @@
import logging import logging
import asyncio import asyncio
import types import types
import typing
import functools import functools
import datetime
import bottom import bottom
@ -17,8 +19,15 @@ class Bot(bottom.Client):
nick: str, nick: str,
user: str = None, user: str = None,
realname: 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 = logging.getLogger(self.__class__.__name__)
self.logger.debug("Initializing...") self.logger.debug("Initializing...")
self.nick = nick self.nick = nick
@ -26,6 +35,15 @@ class Bot(bottom.Client):
self.realname = realname or self.nick self.realname = realname or self.nick
self.channels = channels or [] self.channels = channels or []
self.pre_pub = [] 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("ping", self.keepalive)
self.on("CLIENT_CONNECT", self.on_connect) self.on("CLIENT_CONNECT", self.on_connect)
self.on("PRIVMSG", self.privmsg) self.on("PRIVMSG", self.privmsg)
@ -103,6 +121,149 @@ class Bot(bottom.Client):
self.pre_pub.extend([wrapped]) self.pre_pub.extend([wrapped])
return func 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: def privmsg(self, **kwargs) -> None:
self.logger.debug("PRIVMSG {}".format(kwargs)) self.logger.debug("PRIVMSG {}".format(kwargs))