145 lines
2.9 KiB
Python
145 lines
2.9 KiB
Python
#!/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, "
|
|
"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()
|
|
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, 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, token_hash, date_issued, date_expired) "
|
|
"VALUES (?,?,?,?)",
|
|
(user_id, token_hash, date_issued, date_expired)
|
|
).lastrowid
|
|
return token_id
|
|
|
|
|
|
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_credentials(user_id):
|
|
"""
|
|
Returns all credentials registered to a user.
|
|
"""
|
|
data = db_execute(
|
|
"SELECT * FROM credential WHERE user_id = ?",
|
|
(user_id,)
|
|
).fetchall()
|
|
return data
|