Download logs and untruncate messages

main.py

@@ -1,15 +1,24 @@
+import sys
+import time
+import logging
 import sqlite3
+import threading
 
-db_name = "bugbot.db"
-db_setup = """
+import bugbot
+
+
+THREAD_LIMIT = 10
+DB_NAME = "logs.db"
+DB_SETUP = """
 CREATE TABLE IF NOT EXISTS messages(
-	id      STRING  PRIMARY KEY,
+	id      STRING  NOT NULL,
 	room    STRING  NOT NULL,
 	time    INTEGER NOT NULL,
 	session STRING  NOT NULL,
 	name    STRING  NOT NULL,
 	content STRING  NOT NULL,
-	parent  STRING
+	parent  STRING,
+	PRIMARY KEY (id, room)
 );
 
 CREATE TABLE IF NOT EXISTS sessions(
@@ -20,15 +29,166 @@ CREATE TABLE IF NOT EXISTS sessions(
 );
 
 CREATE TABLE IF NOT EXISTS rooms(
-	name STRING PRIMARY KEY
+	name     STRING PRIMARY KEY,
+	password STRING
 );
 """
+HELP_TEXT = """
+Usage: python3 main.py action[ parameters]
 
-def main():
-	# make sure the tables are set up correctly
-	with sqlite3.connect(db_name) as db:
-		db.executescript(db_setup)
+Actions:
+    help[ action name]          -> Display help.
+    list                        -> List rooms saved in the db.
+    add roomname[ password]     -> Add a room to the db.
+    remove roomname[ roomnames] -> Remove a room and all its messages from the db.
+                                   WARNING: This action is irreversible!
+    reset roomname[ roomnames]  -> Remove a room's messages, but not the room.
+                                   This way, the room's log will be downloaded again
+                                   the next time you update it.
+                                   WARNING: This action is irreversible!
+    update[ roomnames]          -> Update a room's log.
+                                   If no room is specified, all logs will be updated.
+    redownload[ roomnames]      -> Redownload a room's log
+                                   If no room is specified, all logs will be redownloaded.
+    readable[ roomnames]        -> Convert a room's log to a readable format.
+                                   If no room is specified, all logs will be converted.
+"""
+
+
+def listrooms():
+	"""
+	List all rooms and passwords.
+	"""
+	
+	with sqlite3.connect(DB_NAME) as db:
+		rooms = db.execute("SELECT * FROM rooms")
+		for room in rooms:
+			if room[1] is not None:
+				print("name: {:20} pw: {}".format(room[0], room[1]))
+			else:
+				print("name: {}".format(room[0]))
+
+
+def loadrooms(names):
+	"""
+	Load rooms/passwords from db.
+	"""
+	
+	rooms = {}
+	
+	if names:
+		with sqlite3.connect(DB_NAME) as db:
+			for name in names:
+				pw = db.execute("SELECT password FROM rooms WHERE name=?", (name,)).fetchone()
+				rooms[name] = pw[0] if pw else None
+	else:
+		with sqlite3.connect(DB_NAME) as db:
+			r = db.execute("SELECT * FROM rooms")
+			for room in r:
+				rooms[room[0]] = room[1]
+	
+	return rooms
+
+
+def addroom(room, pw=None):
+	"""
+	Add a room and pw to the db.
+	"""
+	
+	with sqlite3.connect(DB_NAME) as db:
+		db.execute("INSERT OR REPLACE INTO rooms VALUES(?,?)", (room, pw))
 		db.commit()
 
+
+def removerooms(rooms):
+	"""
+	Remove rooms from the db.
+	"""
+	
+	resetrooms(rooms)
+	
+	with sqlite3.connect(DB_NAME) as db:
+		for room in rooms:
+			db.execute("DELETE FROM rooms WHERE name=?", (room,))
+			db.commit()
+
+
+def resetrooms(rooms):
+	"""
+	Remove all messages of the rooms from the db.
+	"""
+	
+	with sqlite3.connect(DB_NAME) as db:
+		for room in rooms:
+			db.execute("DELETE FROM messages WHERE room=?", (room,))
+			db.commit()
+
+
+def updaterooms(rooms):
+	"""
+	Update rooms' logs.
+	"""
+	
+	for room in rooms:
+		while not threading.active_count() < THREAD_LIMIT:
+			time.sleep(1)
+		bugbot.download.Downloader(room, DB_NAME, password=rooms[room]).launch()
+		print("Started download: {}".format(room))
+	
+	print("Started all downloads")
+
+
+def readable(rooms):
+	print("This action is currently not available.")
+
+
+def main(action, *argv):
+	# initialize logging for all other modules
+	logging.basicConfig(level=logging.INFO,
+	                    format="[%(levelname)s] (%(threadName)-20s) %(message)s")
+	
+	# make sure the tables are set up correctly
+	with sqlite3.connect(DB_NAME) as db:
+		db.executescript(DB_SETUP)
+		db.commit()
+	
+	if action == "help":
+		print(HELP_TEXT)
+	
+	elif action == "list":
+		listrooms()
+	
+	elif action == "add":
+		if len(argv) == 1:
+			addroom(argv[0])
+		elif len(argv) == 2:
+			addroom(argv[0], pw=argv[1])
+		else:
+			print("Usage: addroom roomname[ password]")
+	
+	elif action == "remove":
+		removerooms(argv)
+	
+	elif action == "reset":
+		resetrooms(argv)
+	
+	else:
+		rooms = loadrooms(argv)
+		
+		if action == "update":
+			updaterooms(rooms)
+		
+		elif action == "redownload":
+			resetrooms(rooms)
+			updaterooms(rooms)
+		
+		elif action == "readable":
+			readable(rooms)
+		
+		else:
+			print(HELP_TEXT)
+
 if __name__ == "__main__":
-	main()
+	main(*sys.argv[1:])
+	
+	print("Done")