diff --git a/boots/robot.py b/boots/robot.py index 3eaeb0a..1a26783 100644 --- a/boots/robot.py +++ b/boots/robot.py @@ -121,6 +121,35 @@ class Bot(bottom.Client): self.pre_pub.extend([wrapped]) return func + @staticmethod + async def _wait(timestamp_queue: list, flood_size_threshold: int, + flood_time_threshold: int, logger: logging.Logger) -> list: + """ + This method will figure out if the calling function needs to wait. + If so, it will wait the correct amount of time before returning. + + :param timestamp_queue list: list of timestamp of the size of flood_size_threshold + :param flood_size_threshold int: total size of the flood threshold + :param flood_time_threshold int: total time of the flood threshold + :param logger logging.Logger: logger to use + :return timetamp_queue list: the new timestamp_queue after modifications + """ + timestamp = datetime.datetime.now() + if timestamp_queue.__len__() < flood_size_threshold: + timestamp_queue.extend([timestamp]) + else: + del timestamp_queue[0] + timestamp_queue.extend([timestamp]) + + time_diff = timestamp_queue[-1] - timestamp_queue[0] + + if len(timestamp_queue) == flood_size_threshold and \ + time_diff.total_seconds() < flood_time_threshold: + logger.debug("Waiting {}s".format(flood_time_threshold - time_diff.total_seconds())) + await asyncio.sleep(flood_time_threshold - time_diff.total_seconds()) + + return timestamp_queue + async def global_buffer(self, command: str, **kwargs) -> None: """ This method will buffer the communication sent to the server @@ -134,20 +163,10 @@ class Bot(bottom.Client): :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.global_timestamp_queue = await self._wait(self.global_timestamp_queue, + self.global_flood_size_threshold, + self.global_flood_time_threshold, + self.logger) self.send(command, **kwargs) async def msg_buffer(self, command: str, **kwargs) -> None: @@ -165,19 +184,10 @@ class Bot(bottom.Client): :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()) + self.msg_timestamp_queue = await self._wait(self.msg_timestamp_queue, + self.msg_flood_size_threshold, + self.msg_flood_time_threshold, + self.logger) await self.global_buffer(command, **kwargs) @@ -196,19 +206,10 @@ class Bot(bottom.Client): :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()) + self.action_timestamp_queue = await self._wait(self.action_timestamp_queue, + self.action_flood_size_threshold, + self.action_flood_time_threshold, + self.logger) await self.global_buffer(command, **kwargs)