From f779c699cf3a0a55f252862dca0a8e53132d721c Mon Sep 17 00:00:00 2001 From: klemek Date: Thu, 7 Jan 2021 10:50:25 +0100 Subject: [PATCH] overload security --- emotes.py | 41 +++++++++++++++++++++++------------------ log_manager.py | 8 +++++++- 2 files changed, 30 insertions(+), 19 deletions(-) diff --git a/emotes.py b/emotes.py index 873bb30..bfdaa40 100644 --- a/emotes.py +++ b/emotes.py @@ -69,25 +69,30 @@ async def compute(client: discord.client, message: discord.Message, *args: str): async with message.channel.typing(): progress = await message.channel.send("```Starting analysis...```") total_msg, total_chan = await logs.load(progress, channels) - msg_count = 0 - chan_count = 0 - for id in logs.channels: - count = analyse_channel( - logs.channels[id], emotes, raw_members, all_emojis="all" in args + if total_msg == -1: + await message.channel.send( + "f{message.author} An analysis is already running on this server, please be patient." + ) + else: + msg_count = 0 + chan_count = 0 + for id in logs.channels: + count = analyse_channel( + logs.channels[id], emotes, raw_members, all_emojis="all" in args + ) + msg_count += count + chan_count += 1 if count > 0 else 0 + await progress.edit(content=f"```Computing results...```") + # Display results + await tell_results( + get_intro(emotes, full, channels, members, msg_count, chan_count), + emotes, + message.channel, + total_msg, + top=top, + allow_unused=full and len(members) == 0, + show_life=False, ) - msg_count += count - chan_count += 1 if count > 0 else 0 - await progress.edit(content=f"```Computing results...```") - # Display results - await tell_results( - get_intro(emotes, full, channels, members, msg_count, chan_count), - emotes, - message.channel, - total_msg, - top=top, - allow_unused=full and len(members) == 0, - show_life=False, - ) # Delete custom progress message await progress.delete() diff --git a/log_manager.py b/log_manager.py index a707eba..d5a4e21 100644 --- a/log_manager.py +++ b/log_manager.py @@ -15,6 +15,8 @@ if not os.path.exists(LOG_DIR): CHUNK_SIZE = 1000 FORMAT = 2 +current_analysis = [] + class FakeMessage: def __init__(self, id: int): @@ -151,7 +153,10 @@ class GuildLogs: async def load( self, progress: discord.Message, target_channels: List[discord.TextChannel] = [] - ): + ) -> Tuple[int, int]: + if self.log_file in current_analysis: + return -1, -1 + current_analysis += [self.log_file] await progress.edit( content=f"```Reading history...\n(this might take a while)```" ) @@ -220,4 +225,5 @@ class GuildLogs: await progress.edit( content=f"```Analysing...\n{total_msg} messages in {total_chan} channels```" ) + current_analysis.remove(self.log_file) return total_msg, total_chan