anonkun/database.py

#!/usr/bin/env python3
"""
Database tools and functions.
"""
import MySQLdb
from flask import session
from passlib.hash import argon2

class Database():
	"""
	An interface to interact with the database.
	"""
	def __init__(self):
		"""
		Initalizes the database.
		"""
		with open("db_key", "r") as file:
			# TODO: encrypt this 
			self.user, self.db, self.key = file.read().strip().split()
		try:
			self.execute("SELECT * FROM `users`").fetchone()
		except MySQLdb.ProgrammingError: # database not initialized
			with open("anonkun.sql", "r") as file:
				commands = file.read().split(";")
			for cmd in commands:
				cmd = cmd.strip()
				if not cmd:
					continue
				self.execute(cmd)


	def execute(self, *args, **kwargs):
		"""
		Opens a connection to the app's database and executes the SQL
		statements passed to this function.
		"""
		with MySQLdb.connect(
				user=self.user,
				passwd=self.key,
				db=self.db,
				charset='utf8mb4') as cur:
			cur.execute(*args, **kwargs)
			return cur

_DB = Database()


def add_user(username, password, timestamp):
	"""
	Adds a user to the database.
	"""
	if verify_username(username): # username taken
		return "username_taken"
	elif len(username) > 20:
		return "username_too_long"

	pw_hash = argon2.hash(password)
	_DB.execute(
		"INSERT INTO `users` (`username`, `password_hash`, `signup_date`) " \
		+ "VALUES (%s, %s, %s)", (username, pw_hash, timestamp))
	return "success"


def verify_password(username, password):
	"""
	Verifies a user's password.
	"""
	user = verify_username(username)
	if not user:
		return False

	user_id, _, pw_hash, _ = user

	if argon2.verify(password, pw_hash):
		session["user_id"] = user_id
		return True
	else:
		return False


def verify_username(username):
	"""
	Checks to see if the given username is in the database.
	"""
	user = _DB.execute("SELECT * FROM `users` WHERE `username` = %s",
		(username,)).fetchone()
	if user:
		return user
	else:
		return False


def log_chat_message(data):
	"""
	Logs chat messages into the database.
	'data' should be a dict containing: message, date, room_id, name, and
	user_id (optional).
	"""
	message = data["message"]
	date = data["date"]
	room_id = int(data["room"])
	name = data["name"]
	user_id = data.get("user_id")
	_DB.execute(
		"INSERT INTO `chat_messages` (" \
		+ "`message`, `room_id`, `date`, `name`, `name_id`) VALUES (" \
		+ "%s, %s, %s, %s, %s)", (message, room_id, date, name, user_id))


def get_chat_messages(room_id):
	"""
	Retrieves all chat messages for the provided room_id.
	"""
	res = _DB.execute(
		"SELECT * FROM `chat_messages` WHERE `room_id` = %s " \
		+ "ORDER BY `date` ASC", (room_id,)).fetchall()
	return res


def insert_quest(canon_title, ident_title, owner_id):
	"""
	Creates a new quest entry.
	"""
	_DB.execute(
		"INSERT INTO `quest_meta` (`canon_title`, `ident_title`, `owner_id`) "\
		+ "VALUES (%s, %s, %s)", (canon_title, ident_title, owner_id))
	quest_id = _DB.execute(
		"SELECT `quest_id` FROM `quest_meta` WHERE `ident_title` = %s",
		(ident_title,)).fetchone()[0]
	return quest_id


def insert_quest_post(quest_id, post, timestamp):
	"""
	Insers a new quest post.
	"""
	_DB.execute(
		"INSERT INTO `quest_data` (`quest_id`, `post`, `timestamp`) " \
		+ "VALUES (%s, %s, %s)", (quest_id, post, timestamp))
	post_id = _DB.execute(
		"SELECT `post_id` FROM `quest_data` WHERE `quest_id` = %s " \
		+ "ORDER BY `post_id` DESC", (quest_id,)).fetchone()[0]
	return post_id


def get_quest_meta(quest_id=None, ident_title=None):
	"""
	Retrieves all meta info about a quest. Allows searching by either
	quest_id or ident_title.
	"""
	statement = "SELECT * FROM `quest_meta` WHERE "
	if quest_id:
		statement += "`quest_id` = %s"
		data = _DB.execute(statement, (quest_id,)).fetchone()
	elif ident_title:
		statement += "`ident_title` = %s"
		data = _DB.execute(statement, (ident_title,)).fetchone()
	else:
		return
	return data


def get_quest_data(quest_id):
	"""
	Retrieves all quest posts.
	"""
	data = _DB.execute("SELECT * FROM `quest_data` WHERE `quest_id` = %s",
		(quest_id,)).fetchall()
	return data


def get_user_info(username):
	"""
	Retrives relevant user data.
	"""
	data = _DB.execute(
		"SELECT `user_id`, `signup_date` FROM `users` WHERE `username` = %s",
		(username,)).fetchone()
	return data


def get_user_quests(user_id):
	"""
	Retrieves all quests ran by a particular user_id.
	"""
	data = _DB.execute(
		"SELECT * FROM `quest_meta` WHERE `owner_id` = %s",
		(user_id,)).fetchall()
	return data


def update_quest_post(post_id, new_post):
	"""
	Updates a quest post.
	"""
	_DB.execute(
		"UPDATE `quest_data` SET `post` = %s WHERE `post_id` = %s",
		(new_post, post_id))