anonkun/database.py

195 lines
4.4 KiB
Python

#!/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))
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))