Download logs and untruncate messages

.gitignore

@@ -1,3 +1,3 @@
 *.db
-logs_readable/
+logs_readable/*
 */__pycache__/

bugbot/__init__.py

@@ -0,0 +1,3 @@
+from .connection import Connection
+from .download import Downloader
+from .log import Log

bugbot/connection.py

@@ -1,5 +1,6 @@
 import json
 import time
+import logging
 import threading
 import websocket
 from websocket import WebSocketException as WSException
@@ -46,12 +47,16 @@ class Connection():
 					ROOM_FORMAT.format(self.room),
 					enable_multithread=True
 				)
+				logging.debug("Connected")
 				return True
+			
 			except WSException:
 				if tries > 0:
 					tries -= 1
 				if tries != 0:
 					time.sleep(delay)
+		
+		logging.debug("Failed to connect")
 		return False
 	
 	def disconnect(self):
@@ -65,14 +70,19 @@ class Connection():
 		if self.ws:
 			self.ws.close()
 			self.ws = None
+		
+		logging.debug("Disconnected")
 	
-	def launch(self):
+	def launch(self, func=None):
 		"""
-		launch() -> Thread
+		launch(function) -> Thread
 		
 		Connect to the room and spawn a new thread running run.
+		This also calls the function func in the new thread.
 		"""
 		
+		self.func = func
+		
 		if self.connect(tries=1):
 			self.thread = threading.Thread(target=self.run, name=self.room)
 			self.thread.start()
@@ -87,6 +97,9 @@ class Connection():
 		Receive messages.
 		"""
 		
+		if self.func:
+			self.func()
+		
 		while not self.stopping:
 			try:
 				self.handle_json(self.ws.recv())
@@ -104,6 +117,8 @@ class Connection():
 		
 		self.stopping = True
 		self.disconnect()
+		
+		logging.debug("Stopped")
 	
 	def join(self):
 		"""
@@ -231,5 +246,6 @@ class Connection():
 			"data": kwargs or None,
 			"id": str(self.send_id)
 		}
+		
 		self.send_id += 1
 		self.send_json(packet)

bugbot/dbaccess.py

@@ -0,0 +1,70 @@
+import sqlite3
+
+import yaboli
+
+class DBAccess():
+	"""
+	Takes care of opening and closing the connection to the db.
+	"""
+	
+	def __init__(self, db):
+		"""
+		db - path to the db, or ":memory:"
+		"""
+		
+		self._con = sqlite3.connect(db)
+		self._con.row_factory = sqlite3.Row
+	
+	def execute(self, *args, **kwargs):
+		return self._con.execute(*args, **kwargs)
+	
+	def close(self):
+		self._con.close()
+	
+	def __del__(self):
+		self.close()
+
+class Log(DBAccess):
+	"""
+	More abstract way to access a room's messages in the db.
+	"""
+	
+	def __init__(self, db, room):
+		"""
+		db   - path to the db, or ":memory:"
+		room - name of the room
+		"""
+		
+		super().__init__(self, db)
+		self._room = room
+	
+	def get_session(self, sid):
+		"""
+		get_session(session_id) -> session
+		
+		Returns the session with that id.
+		"""
+		
+		cur = self.execute("SELECT * FROM sessions WHERE id=?", (mid, self._room))
+		result = cur.fetchone()
+		if result:
+			return yaboli.Message(
+				# TODO: <Arguments go here>
+			)
+	
+	def get_message(self, mid):
+		"""
+		get_message(message_id) -> message
+		
+		Returns the message with that id.
+		"""
+		
+		cur = self.execute("SELECT * FROM messages WHERE id=? AND room=?", (mid, self._room))
+		result = cur.fetchone()
+		if result:
+			return yaboli.Message(
+				# TODO: <Arguments go here>
+			)
+
+class Rooms(DBAccess):
+	pass

bugbot/download.py

@@ -1,20 +1,34 @@
-import tempfile
+import logging
 
-import connection
+from . import log
+from . import connection
 
 class Downloader():
 	"""
 	Update or redownload a room's log.
 	"""
 	
-	def __init__(self, room, logfile, password=None):
+	def __init__(self, room, db_name, password=None):
 		"""
-		room     - name of the room to download the logs of
+		room     - name of the room to download the log of
-		logfile  - path to the file to save the log in
+		db_name  - name of the db to download the log to
 		password - password of said room, optional
 		"""
 		
-		pass
+		self.password = password
+		self.con = connection.Connection(room)
+		self.log = log.Log(db_name, room)
+		
+		self.downloading = True # still downloading new messages
+		self.truncated = 0 # messages still truncated
+		
+		self.con.add_callback("ping-event",        self._handle_ping_event       )
+		self.con.add_callback("bounce-event",      self._handle_bounce_event     )
+		self.con.add_callback("auth-reply",        self._handle_auth_reply       )
+		self.con.add_callback("disconnect-event",  self._handle_disconnect_event )
+		self.con.add_callback("snapshot-event",    self._handle_snapshot_event   )
+		self.con.add_callback("get-message-reply", self._handle_get_message_reply)
+		self.con.add_callback("log-reply",         self._handle_log_reply        )
 	
 	def _handle_ping_event(self, data):
 		"""
@@ -23,7 +37,8 @@ class Downloader():
 		Pong!
 		"""
 		
-		pass
+		self.con.send_packet("ping-reply", time=data["time"])
+		logging.debug("Ping-reply on {}, expected next on {}.".format(data["time"], data["next"]))
 	
 	def _handle_bounce_event(self, data):
 		"""
@@ -32,7 +47,13 @@ class Downloader():
 		Authenticate if possible, otherwise give up and stop.
 		"""
 		
-		pass
+		if self.password:
+			self.con.send_packet("auth", type="passcode", passcode=self.password)
+			logging.info("Bounce! Authenticating with {}".format(self.password))
+		else:
+			self.log.close()
+			self.con.stop()
+			logging.warn("Bounce! Could not authenticate :/")
 	
 	def _handle_auth_reply(self, data):
 		"""
@@ -41,7 +62,12 @@ class Downloader():
 		Disconnect if authentication unsucessful.
 		"""
 		
-		pass
+		if data["success"]:
+			logging.debug("Successfully authenticated")
+		else:
+			logging.warn("Error authenticating: '{}'".format(data["reason"]))
+			self.log.close()
+			self.con.stop()
 	
 	def _handle_disconnect_event(self, data):
 		"""
@@ -50,36 +76,78 @@ class Downloader():
 		Immediately disconnect.
 		"""
 		
-		pass
+		logging.warn("Disconnecting: '{}'".format(data["reason"]))
+		self.log.close()
+		self.con.stop()
 	
 	def _handle_snapshot_event(self, data):
 		"""
 		_handle_snapshot_event(data) -> None
 		
-		Save messages and request further messages
+		Save messages and request further messages.
 		"""
 		
-		pass
+		self.add_messages(data["log"])
 	
 	def _handle_get_message_reply(self, data):
 		"""
 		_handle_get_message_reply(data) -> None
 		
-		Append untruncated message to log file and then continue
+		Replace truncated message by untruncated message.
-		transferring the messages from the temp file to the
-		log file.
 		"""
 		
-		pass
+		logging.debug("Untruncate! {}".format(data["id"]))
-	
+		self.log.add_message(data)
+		self.truncated -= 1
+		
+		if self.truncated <= 0 and not self.downloading:
+			logging.debug("Last untruncated message received - stopping now")
+			self.log.close()
+			self.con.stop()
+		
 	def _handle_log_reply(self, data):
 		"""
 		_handle_log_reply(data) -> None
 		
-		Save messages received to temp file.
+		Save messages and request further messages.
 		"""
 		
-		pass
+		self.add_messages(data["log"])
+	
+	def add_messages(self, msgs):
+		"""
+		add_mesages(messages) -> None
+		
+		Save messages to the db and request further messages.
+		"""
+		
+		logging.info("Processing messages")
+		
+		if len(msgs) == 0:
+			logging.info("End of log - empty")
+			self.log.close()
+			self.con.stop()
+			return
+		
+		for msg in msgs[::-1]:
+			logging.debug("Testing '{}' from {}".format(msg["id"], msg["sender"]["name"]))
+			
+			if msg["id"] <= self.newmsg:
+				logging.info("End of log - too old")
+				self.log.close()
+				self.con.stop()
+				return
+			
+			else:
+				logging.debug("Adding message: {}".format(msg["id"]))
+				self.log.add_message(msg)
+				logging.info("Untruncating message: {}".format(msg["id"]))
+				self.con.send_packet("get-message", id=msg["id"])
+		
+		else:
+			self.log.commit()
+			logging.info("Requesting more messages")
+			self.con.send_packet("log", n=1000, before=msgs[0]["id"])
 	
 	def launch(self):
 		"""
@@ -88,13 +156,14 @@ class Downloader():
 		Start the download in a separate thread.
 		"""
 		
-		pass
+		self.con.launch(self._on_launch)
 	
-	def transfer(self):
+	def _on_launch(self):
 		"""
-		transfer() -> None
+		_on_launch() -> None
 		
-		Transfer the messages from the temporary file to the log file.
+		Gets called in the new thread.
 		"""
 		
-		pass
+		self.log.open()
+		self.newmsg = self.log.get_newest() or ""

bugbot/log.py

@@ -1,3 +1,4 @@
+import logging
 import sqlite3
 
 class Log():
@@ -9,10 +10,30 @@ class Log():
 		"""
 		name - name of the db
 		room - name of the room
+		
+		This also opens a connection to the db - make sure to close that later!
 		"""
 		self.name = name
 		self.room = room
 	
+	def open(self):
+		"""
+		open() -> None
+		
+		Open the connection to the db.
+		"""
+		
+		self.con = sqlite3.connect(self.name)
+	
+	def close(self):
+		"""
+		close() -> None
+		
+		Close the connection to the db.
+		"""
+		self.con.commit()
+		self.con.close()
+	
 	def get_newest(self):
 		"""
 		get_newest() -> message_id
@@ -21,14 +42,13 @@ class Log():
 		Returns None if no message was found.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		message = self.con.execute(
-			message = db.execute(
+			"SELECT id FROM messages WHERE room=? ORDER BY id DESC LIMIT 1",
-				"SELECT id FROM messages WHERE room=? ORDER BY id DESC LIMIT 1",
+			(self.room,)
-				(self.room,)
+		)
-			)
 		
 		result = message.fetchone()
-		return result[0] if result
+		return result[0] if result else None
 	
 	def get_top_level(self):
 		"""
@@ -37,14 +57,13 @@ class Log():
 		Returns a full list of top-level messages' ids.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		message = self.con.execute(
-			message = db.execute(
+			"SELECT id FROM messages WHERE parent ISNULL AND room=?",
-				"SELECT id FROM messages WHERE parent ISNULL AND room=?",
+			(self.room,)
-				(self.room,)
+		)
-			)
 		
 		result = message.fetchall()
-		return [entry[0] for entry in result] if result
+		return [entry[0] for entry in result] if result else None
 	
 	def get_message(self, mid):
 		"""
@@ -53,11 +72,10 @@ class Log():
 		Returns message with that id.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		message = self.con.execute(
-			message = db.execute(
+			"SELECT * FROM messages WHERE id=? AND room=?",
-				"SELECT * FROM messages WHERE id=? AND room=?",
+			(mid, self.room)
-				(mid, self.room)
+		)
-			)
 		
 		result = message.fetchone()
 		return {
@@ -77,14 +95,13 @@ class Log():
 		Returns the message's parent's id.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		message = self.con.execute(
-			message = db.execute(
+			"SELECT parent FROM messages WHERE id=? AND room=?",
-				"SELECT parent FROM messages WHERE id=? AND room=?",
+			(mid, self.room)
-				(mid, self.room)
+		)
-			)
 		
 		result = message.fetchone()
-		return result[0] if result
+		return result[0] if result else None
 	
 	def get_children(self, mid):
 		"""
@@ -93,14 +110,13 @@ class Log():
 		Returns a list of the message's childrens' ids.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		message = self.con.execute(
-			message = db.execute(
+			"SELECT id FROM messages WHERE parent=? AND room=?",
-				"SELECT id FROM messages WHERE parent=? AND room=?",
+			(mid, self.room)
-				(mid, self.room)
+		)
-			)
 		
 		result = message.fetchall()
-		return [entry[0] for entry in result] if result
+		return [entry[0] for entry in result] if result else None
 	
 	def add_message(self, msg):
 		"""
@@ -109,28 +125,36 @@ class Log():
 		Add a message to the db.
 		"""
 		
-		with sqlite3.connect(self.name) as db:
+		# insert or update message
-			# insert or update message
+		self.con.execute(
-			db.execute(
+			"INSERT OR REPLACE INTO messages VALUES(?,?,?,?,?,?,?)",
-				"INSERT OR REPLACE INTO messages VALUES(?,?,?,?,?,?,?)",
+			(
-				(
+				msg["id"],
-					msg["id"],
+				self.room,
-					self.room,
+				msg["time"],
-					msg["time"],
+				msg["sender"]["session_id"],
-					msg["sender"]["session_id"],
+				msg["sender"]["name"],
-					msg["sender"]["name"],
+				msg["content"],
-					msg["content"],
+				msg["parent"] if "parent" in msg else None
-					msg["parent"] if "parent" in msg else None
-				)
 			)
-			
+		)
-			# insert or update session
+		
-			db.execute(
+		# insert or update session
-				"INSERT OR REPLACE INTO sessions VALUES(?,?,?,?)",
+		self.con.execute(
-				(
+			"INSERT OR REPLACE INTO sessions VALUES(?,?,?,?)",
-					msg["sender"]["session_id"],
+			(
-					msg["sender"]["id"],
+				msg["sender"]["session_id"],
-					1 if "is_staff" in msg["sender"] and msg["sender"]["is_staff"] else None,
+				msg["sender"]["id"],
-					1 if "is_manager" in msg["sender"] and msg["sender"]["is_manager"] else None
+				1 if "is_staff" in msg["sender"] and msg["sender"]["is_staff"] else None,
-				)
+				1 if "is_manager" in msg["sender"] and msg["sender"]["is_manager"] else None
 			)
+		)
+	
+	def commit(self):
+		"""
+		commit() -> None
+		
+		Write all the changes to the db.
+		"""
+		
+		self.con.commit()

bugbot/rooms.py

@@ -0,0 +1,6 @@
+import sqlite3
+
+class Rooms():
+	"""
+	
+	"""

convert.py

@@ -0,0 +1,39 @@
+import sys
+import json
+import sqlite3
+
+def main(filename, roomname):
+	with open(filename) as f:
+		log = json.load(f)
+	
+	with sqlite3.connect("logs.db") as db:
+		for msg in log:
+			print("Adding {}".format(msg))
+			
+			# insert or update message
+			db.execute(
+				"INSERT OR REPLACE INTO messages VALUES(?,?,?,?,?,?,?)",
+				(
+					msg["id"],
+					roomname,
+					msg["time"],
+					msg["sender"]["session_id"],
+					msg["sender"]["name"],
+					msg["content"],
+					msg["parent"] if "parent" in msg else None
+				)
+			)
+			
+			# insert or update session
+			db.execute(
+				"INSERT OR REPLACE INTO sessions VALUES(?,?,?,?)",
+				(
+					msg["sender"]["session_id"],
+					msg["sender"]["id"],
+					1 if "is_staff" in msg["sender"] and msg["sender"]["is_staff"] else None,
+					1 if "is_manager" in msg["sender"] and msg["sender"]["is_manager"] else None
+				)
+			)
+
+if __name__ == "__main__":
+	main(sys.argv[1], sys.argv[2])

main.py

@@ -1,15 +1,24 @@
+import sys
+import time
+import logging
 import sqlite3
+import threading
 
-db_name = "bugbot.db"
+import bugbot
-db_setup = """
+
+
+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():
+Actions:
-	# make sure the tables are set up correctly
+    help[ action name]          -> Display help.
-	with sqlite3.connect(db_name) as db:
+    list                        -> List rooms saved in the db.
-		db.executescript(db_setup)
+    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")