Juice/db.py

#!/usr/bin/env python3
"""
A module for interacting with Juice's database.
"""
import time
import sqlite3
import threading


DB_LOCK = threading.Lock()


def init_db():
	"""
	Checks to see if the database is initialized yet or not. If not,
	the appropriate tables are created.
	"""
	con = sqlite3.connect('juice.db')
	cur = con.cursor()
	try:
		cur.execute("SELECT * FROM user LIMIT 1").fetchone()
		cur.execute("SELECT * FROM credential LIMIT 1").fetchone()
		cur.execute("SELECT * FROM token LIMIT 1").fetchone()
	except sqlite3.OperationalError:
		cur.execute("CREATE TABLE user("
			"id INTEGER PRIMARY KEY, "
			"username TEXT UNIQUE, "
			"email TEXT UNIQUE"
			")"
		)
		cur.execute("CREATE TABLE credential("
			"id INTEGER PRIMARY KEY, "
			"user_id INTEGER REFERENCES user(id) ON UPDATE CASCADE, "
			"nick TEXT, "
			"credential BLOB"
			")"
		)
		cur.execute("CREATE TABLE token("
			"id INTEGER PRIMARY KEY, "
			"user_id INTEGER REFERENCES user(id) ON UPDATE CASCADE, "
			"user_agent TEXT, "
			"ip_address TEXT, "
			"token_hash TEXT, "
			"date_issued INTEGER, "
			"date_expired INTEGER"
			")"
		)
		con.commit()
	con.close()
init_db()


def db_execute(*args, **kwargs):
	"""
	Opens a connection to the app's database and executes the SQL
	statements passed to this function.
	"""
	with sqlite3.connect('juice.db') as con:
		DB_LOCK.acquire()
		con.row_factory = sqlite3.Row
		cur = con.cursor()
		res = cur.execute(*args, **kwargs)
		DB_LOCK.release()
		return res


def set_user(username, email):
	"""
	Adds a new user.
	"""
	user_id = db_execute(
		"INSERT INTO user(username, email) VALUES (?, ?)",
		(username, email)
	).lastrowid
	return user_id


def get_user(username):
	"""
	Returns a user entry.
	"""
	data = db_execute(
		"SELECT * FROM user WHERE username = ?",
		(username,)
	).fetchone()
	return data


def set_token(user_id, user_agent, ip_address, token_hash):
	"""
	Sets a user's token hash.
	"""
	date_issued = int(time.time())
	date_expired = date_issued + 30*24*60*60
	token_id = db_execute(
		"INSERT INTO "
		"token("
		"user_id, user_agent, ip_address, token_hash, date_issued,date_expired"
		") "
		"VALUES (?,?,?,?,?,?)",
		(user_id, user_agent, ip_address, token_hash, date_issued,date_expired)
	).lastrowid
	return token_id


def get_token(token_id):
	"""
	Returns the token of the specified id.
	"""
	data = db_execute(
		"SELECT * FROM token WHERE id = ?",
		(token_id,)
	).fetchone()
	return data


def delete_token(token_id):
	"""
	Deletes the token of the specified id.
	"""
	db_execute(
		"DELETE FROM token WHERE id = ?",
		(token_id,)
	).fetchone()
	return True


def get_tokens(user_id):
	"""
	Returns all tokens assigned to a user.
	"""
	data = db_execute(
		"SELECT * FROM token WHERE user_id = ?",
		(user_id,)
	).fetchall()
	return data


def refresh_token(token_id):
	"""
	Extends a token's expiration date.
	"""
	new_date_expired = int(time.time()) + 30*24*60*60
	db_execute(
		"UPDATE token SET date_expired = ? WHERE id = ?",
		(new_date_expired, token_id)
	)
	return True


def set_credential(user_id, nick, credential):
	"""
	Adds a credential to the database.
	"""
	cred_id = db_execute(
		"INSERT INTO credential(user_id, nick, credential) "
		"VALUES (?, ?, ?)",
		(user_id, nick, credential)
	).lastrowid
	return cred_id


def get_credential(cred_id):
	"""
	Returns the credential of the specified id.
	"""
	data = db_execute(
		"SELECT * FROM credential WHERE id = ?",
		(cred_id,)
	).fetchone()
	return data


def delete_credential(cred_id):
	"""
	Deletes the credential of the specified id.
	"""
	db_execute(
		"DELETE FROM credential WHERE id = ?",
		(cred_id,)
	).fetchone()
	return True


def get_credentials(user_id):
	"""
	Returns all credentials registered to a user.
	"""
	data = db_execute(
		"SELECT * FROM credential WHERE user_id = ?",
		(user_id,)
	).fetchall()
	return data