bugbot/main.py

import sys
import time
import logging
import sqlite3
import threading

import bugbot


THREAD_LIMIT = 10
DB_NAME = "logs.db"
DB_SETUP = """
CREATE TABLE IF NOT EXISTS messages(
	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,
	PRIMARY KEY (id, room)
);

CREATE TABLE IF NOT EXISTS sessions(
	id         STRING PRIMARY KEY,
	user_id    STRING NOT NULL,
	is_staff   INTEGER,
	is_manager INTEGER
);

CREATE TABLE IF NOT EXISTS rooms(
	name     STRING PRIMARY KEY,
	password STRING
);
"""
HELP_TEXT = """
Usage: python3 main.py action[ parameters]

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(*sys.argv[1:])
	
	print("Done")