From d9326881b3d2d8f1fa7d29e3d279a12cb1833e19 Mon Sep 17 00:00:00 2001 From: matthew-robertson Date: Tue, 15 Jan 2019 20:55:29 -0500 Subject: [PATCH] Massive Streamlining Reworks the bot to actually make use of the additional DB related powers. Everything has been refactored to avoid a lot of duplicated code and weirdly implemented things. Fixes an issue where the bot would be immediately able to respond if it had just been restarted. Probably we do other stuff too. --- Procfile | 1 - bot.py | 235 ++++++++++++++---------------------- dao/server_dao.py | 22 +++- migrations/server-table.sql | 1 + runtime.txt | 1 - upgradeToDB.py | 34 +++--- 6 files changed, 123 insertions(+), 171 deletions(-) delete mode 100644 Procfile delete mode 100644 runtime.txt diff --git a/Procfile b/Procfile deleted file mode 100644 index 29cff6d..0000000 --- a/Procfile +++ /dev/null @@ -1 +0,0 @@ -worker: python bot.py diff --git a/bot.py b/bot.py index 34c839d..68de16e 100644 --- a/bot.py +++ b/bot.py @@ -1,6 +1,7 @@ import asyncio import datetime import discord +from enum import Enum import math import re import sys @@ -10,67 +11,20 @@ # A pattern to match the word vore, and only the single word vore. pattern = re.compile(r'\b[\*|_|~|`|-|\.]*[V|v][\*|_|~|`|-|\.]*[O|Ò|Ó|Ô|Õ|Ö|o|ò|ó|ô|õ|ö|ᴑ|о][\*|_|~|`|-|\.]*[R|r][\*|_|~|`|-|\.]*[E|È|É|Ê|Ë|Е|e|è|é|ê|ë|е][\*|_|~|`|-|\.]*[S|s]?\b') -serverAndDate = {} -botStartup = datetime.datetime.now() -lastMention = {} -awake = {} - -def readTimesFromFile(): - servers = ServerDao.get_all_servers() - - global serverAndDate - for server in servers: - serverAndDate[server[0]] = datetime.datetime.strptime(server[1], "%Y-%m-%d %H:%M:%S") - awake[server[0]] = bool(server[2]) - - -def writeTimesToFile(): - for serverId in serverAndDate: - ServerDao.insert_server(serverId, serverAndDate[serverId].strftime("%Y-%m-%d %H:%M:%S"), awake[serverId]) - client = discord.Client(shard_id=int(sys.argv[1]), shard_count=int(sys.argv[2])) -readTimesFromFile() -print('Stored server info:') -count = 0 -for id in serverAndDate: - print ("{}: id: {}, time: {}".format(count, id, serverAndDate[id])) - count = count + 1 - -@client.event -async def on_ready(): - print('Logged in as') - print(client.user.name) - print(client.user.id) - print('------') - -@client.event -async def on_message_edit(before, message): - global botStartup - global serverAndDate - global lastMention - global awake - currentTime = datetime.datetime.now() - serverId = message.server.id - - lastReferenced = botStartup - if serverId in serverAndDate: - lastReferenced = serverAndDate[serverId] - if serverId not in lastMention: - bot = None - for x in message.server.members: - if client.user.id == x.id: - bot = x - break - lastMention[serverId] = bot.joined_at - if serverId not in awake: - awake[serverId] = True +class Commands(Enum): + NOCOMMAND = 0 + VTSILENCE = 1 + VTALERT = 2 + VTHELP = 3 + VT = 4 - # Begin time formatting +def format_time(currentTime, lastReferenced): diff = currentTime - lastReferenced hours = math.floor(diff.seconds/3600) - minutes = math.floor((diff.seconds - hours * 3600)/60) - seconds = diff.seconds - hours * 3600 - minutes * 60 + minutes = math.floor((diff.seconds - hours*3600)/60) + seconds = diff.seconds - hours*3600 - minutes*60 dt = "{} days, ".format(diff.days) ht = "{} hours, ".format(hours) mt = "{} minutes, and ".format(minutes) @@ -95,108 +49,95 @@ async def on_message_edit(before, message): mt = "1 minute, and " if seconds == 1: st = "1 second" - # End Time formatting stuff - - if (hasattr(message.author, 'server_permissions')): - permission = message.author.server_permissions - if message.content.startswith('!vtsilence') and permission.administrator: - await client.send_message(message.channel, "Ok {}, I'll be quiet now. use '!vtalert' to wake me back up!".format(message.author.mention)) - awake[serverId] = False - writeTimesToFile() - elif message.content.startswith('!vtalert') and permission.administrator: - await client.send_message(message.channel, "Ok {}, I'm scanning now.".format(message.author.mention)) - awake[serverId] = True - writeTimesToFile() - if message.content.startswith('!vtsilence') or message.content.startswith('!vtalert'): + + return dt+ht+mt+st + +# Only returns a command if the command is in the message and also valid +def parse_for_command(msg, msg_author): + # Check if any admin only commands have been entered + if (hasattr(msg_author, 'server_permissions')): + permission = msg_author.server_permissions + if msg.startswith('!vtsilence') and permission.administrator: + return Commands.VTSILENCE + elif msg.startswith('!vtalert') and permission.administrator: + return Commands.VTALERT + + # Check the other commands + if msg.startswith('!vtsilence') or msg.startswith('!vtalert'): pass - elif message.content.startswith('!vthelp'): - await client.send_message(message.channel, "You can ask me how long we've made it with '!vt'.\n If you're an admin you can silence me with '!vtsilence' and wake me back up with '!vtalert'") - elif message.content.startswith('!vt'): - await client.send_message(message.channel, 'The server has gone {}{}{}{} without mentioning the forbidden word.'.format(dt, ht, mt, st)) - if ((pattern.search(message.content) is not None) and (message.author.id != client.user.id)): - serverAndDate[serverId] = currentTime - writeTimesToFile() - print ("{}::: {} lasted {} seconds.".format(currentTime, serverId, (currentTime - lastMention[serverId]).total_seconds())) - if (awake[serverId] and (currentTime - lastMention[serverId]).total_seconds() >= 1800): - await client.send_message(message.channel, '{} referenced the forbidden word, setting the counter back to 0. I\'ll wait a half hour before warning you again.\n The server went {}{}{}{} without mentioning it.'.format(message.author.mention, dt, ht, mt, st)) - lastMention[serverId] = currentTime + elif msg.startswith('!vthelp'): + return Commands.VTHELP + elif msg.startswith('!vt'): + return Commands.VT -@client.event -async def on_message(message): - global botStartup - global serverAndDate - global lastMention - global awake +def handle_message(message, botID): + msg_to_send = False currentTime = datetime.datetime.now() - serverId = message.server.id + if not message.server: + return msg_to_send - lastReferenced = botStartup - if serverId in serverAndDate: - lastReferenced = serverAndDate[serverId] - if serverId not in lastMention: + serverId = message.server.id + currentServer = ServerDao.get_server(serverId) + if currentServer is None: + currentServer = {} + currentServer['server_id'] = serverId bot = None for x in message.server.members: - if client.user.id == x.id: + if botID == x.id: bot = x break - lastMention[serverId] = bot.joined_at - if serverId not in awake: - awake[serverId] = True + currentServer['infracted_at'] = bot.joined_at + currentServer['calledout_at'] = bot.joined_at + currentServer['awake'] = True + ServerDao.insert_server(currentServer) + + timeLasted = format_time(currentTime, currentServer['infracted_at']) + + foundCommand = parse_for_command(message.content, message.author) + if foundCommand is Commands.VTSILENCE: + msg_to_send = "Ok {}, I'll be quiet now. use '!vtalert' to wake me back up!".format(message.author.mention) + currentServer['awake'] = False + ServerDao.insert_server(currentServer) + elif foundCommand is Commands.VTALERT: + msg_to_send = "Ok {}, I'm scanning now.".format(message.author.mention) + currentServer['awake'] = True + ServerDao.insert_server(currentServer) + elif foundCommand is Commands.VTHELP: + msg_to_send = "You can ask me how long we've made it with '!vt'.\n If you're an admin you can silence me with '!vtsilence' and wake me back up with '!vtalert'" + elif foundCommand is Commands.VT: + msg_to_send = "The server has gone {} without mentioning the forbidden word.".format(timeLasted) + + # Check if they've said the forbidden word + if ((pattern.search(message.content) is not None) and (message.author.id != botID)): + currentServer['infracted_at'] = currentTime + + if (currentServer['awake'] and (currentTime - currentServer['calledout_at']).total_seconds() >= 1800): + currentServer['calledout_at'] = currentTime + msg_to_send = "{} referenced the forbidden word, setting the counter back to 0. I'll wait a half hour before warning you again.\n The server went {} without mentioning it.".format(message.author.mention, timeLasted) + + ServerDao.insert_server(currentServer) + print("{}::: {} lasted {} seconds.".format(currentTime, serverId, (currentTime - currentServer['infracted_at']).total_seconds())) + + return msg_to_send - # Begin time formatting - diff = currentTime - lastReferenced - hours = math.floor(diff.seconds/3600) - minutes = math.floor((diff.seconds - hours * 3600)/60) - seconds = diff.seconds - hours * 3600 - minutes * 60 - dt = "{} days, ".format(diff.days) - ht = "{} hours, ".format(hours) - mt = "{} minutes, and ".format(minutes) - st = "{} seconds".format(seconds) +@client.event +async def on_ready(): + print("Logged in as") + print(client.user.name) + print(client.user.id) + print("------") - if diff.days == 1: - dt = "1 day, " - elif diff.days == 0: - dt = "" - if hours == 0: - ht = "" - mt = "{} minutes and ".format(minutes) - if minutes == 0: - mt = "" +@client.event +async def on_message_edit(before, message): + msg_to_send = handle_message(message, client.user.id) + if msg_to_send: + await client.send_message(message.channel, msg_to_send) - if hours == 1: - ht = "1 hour, " - if minutes == 1: - if ht == "": - mt = "1 minute and " - else: - mt = "1 minute, and " - if seconds == 1: - st = "1 second" - # End Time formatting stuff - - if (hasattr(message.author, 'server_permissions')): - permission = message.author.server_permissions - if message.content.startswith('!vtsilence') and permission.administrator: - await client.send_message(message.channel, "Ok {}, I'll be quiet now. use '!vtalert' to wake me back up!".format(message.author.mention)) - awake[serverId] = False - writeTimesToFile() - elif message.content.startswith('!vtalert') and permission.administrator: - await client.send_message(message.channel, "Ok {}, I'm scanning now.".format(message.author.mention)) - awake[serverId] = True - writeTimesToFile() - if message.content.startswith('!vtsilence') or message.content.startswith('!vtalert'): - pass - elif message.content.startswith('!vthelp'): - await client.send_message(message.channel, "You can ask me how long we've made it with '!vt'.\n If you're an admin you can silence me with '!vtsilence' and wake me back up with '!vtalert'") - elif message.content.startswith('!vt'): - await client.send_message(message.channel, 'The server has gone {}{}{}{} without mentioning the forbidden word.'.format(dt, ht, mt, st)) - if ((pattern.search(message.content) is not None) and (message.author.id != client.user.id)): - serverAndDate[serverId] = currentTime - writeTimesToFile() - print ("{}::: {} lasted {} seconds.".format(currentTime, serverId, (currentTime - lastMention[serverId]).total_seconds())) - if (awake[serverId] and (currentTime - lastMention[serverId]).total_seconds() >= 1800): - await client.send_message(message.channel, '{} referenced the forbidden word, setting the counter back to 0. I\'ll wait a half hour before warning you again.\n The server went {}{}{}{} without mentioning it.'.format(message.author.mention, dt, ht, mt, st)) - lastMention[serverId] = currentTime +@client.event +async def on_message(message): + msg_to_send = handle_message(message, client.user.id) + if msg_to_send: + await client.send_message(message.channel, msg_to_send) while True: try: diff --git a/dao/server_dao.py b/dao/server_dao.py index 4d73731..0a87adf 100644 --- a/dao/server_dao.py +++ b/dao/server_dao.py @@ -1,3 +1,5 @@ +import datetime + from dao import conn class ServerDao: @@ -8,9 +10,25 @@ def get_all_servers(): return c.fetchall() - def insert_server(server_id, infracted_at, awake): + def get_server(server_id): + c = conn.cursor() + c.execute('SELECT * FROM `server` WHERE server_id='+str(server_id)) + + row = c.fetchone() + if row is None: + return None + + server = {'server_id': row[0], + 'infracted_at': datetime.datetime.strptime(row[1], "%Y-%m-%d %H:%M:%S"), + 'calledout_at': datetime.datetime.strptime(row[2], "%Y-%m-%d %H:%M:%S"), + 'awake': row[3]} + return server + + def insert_server(server_row): c = conn.cursor() - res = c.execute('INSERT OR REPLACE INTO `server` (server_id, infracted_at, awake) VALUES (?, ?, ?)', (server_id, infracted_at, int(awake))) + res = c.execute('INSERT OR REPLACE INTO `server` (server_id, infracted_at, calledout_at, awake) VALUES (?, ?, ?, ?)', + (server_row['server_id'], server_row['infracted_at'].strftime("%Y-%m-%d %H:%M:%S"), + server_row['calledout_at'].strftime("%Y-%m-%d %H:%M:%S"), int(server_row['awake']))) conn.commit() return res diff --git a/migrations/server-table.sql b/migrations/server-table.sql index 9fb291a..7a3fee2 100644 --- a/migrations/server-table.sql +++ b/migrations/server-table.sql @@ -1,5 +1,6 @@ CREATE TABLE `server` ( server_id Varchar PRIMARY KEY NOT NULL, infracted_at Varchar NOT NULL, + calledout_at Varchar NOT NULL, awake Int DEFAULT 1 ); diff --git a/runtime.txt b/runtime.txt deleted file mode 100644 index cfa5aa5..0000000 --- a/runtime.txt +++ /dev/null @@ -1 +0,0 @@ -python-3.6.2 diff --git a/upgradeToDB.py b/upgradeToDB.py index 2f9e3da..5936fd1 100644 --- a/upgradeToDB.py +++ b/upgradeToDB.py @@ -2,25 +2,19 @@ from dao.server_dao import ServerDao -serverAndDate = {} -awake = {} -def readTimesFromFile(): - global serverAndDate - with open("timeStamps.txt", "r") as target: - for line in target: - tmp = line.split(',') - tmp[1] = tmp[1][0:-1] - serverAndDate[tmp[0]] = datetime.datetime.strptime(tmp[1], "%Y-%m-%d %H:%M:%S") - if (len(tmp) >= 3): - awake[tmp[0]] = tmp[2].strip() == 'True' - else: - awake[tmp[0]] = True - print (awake) - print(serverAndDate) -def writeTimesToDB(): - for serverId in serverAndDate: - ServerDao.insert_server(serverId, serverAndDate[serverId].strftime("%Y-%m-%d %H:%M:%S"), awake[serverId]) +with open("timeStamps.txt", "r") as target: + for line in target: + currentServer = {} + tmp = line.split(',') + tmp[1] = tmp[1][0:-1] -readTimesFromFile() -writeTimesToDB() \ No newline at end of file + currentServer['server_id'] = tmp[0] + currentServer['infracted_at'] = datetime.datetime.strptime(tmp[1], "%Y-%m-%d %H:%M:%S") + currentServer['calledout_at'] = currentServer['infracted_at'] + if (len(tmp) >= 3): + currentServer['awake'] = tmp[2].strip() == 'True' + else: + currentServer['awake'] = True + + ServerDao.insert_server(currentServer) \ No newline at end of file