106 lines
2.0 KiB
Python
106 lines
2.0 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
A module for interacting with Juice's database.
|
||
|
|
"""
|
||
|
|
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()
|
||
|
|
except sqlite3.OperationalError:
|
||
|
|
cur.execute("CREATE TABLE user("
|
||
|
|
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||
|
|
"name TEXT UNIQUE, "
|
||
|
|
"token_hash TEXT"
|
||
|
|
")"
|
||
|
|
)
|
||
|
|
cur.execute("CREATE TABLE credential("
|
||
|
|
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||
|
|
"user_id INTEGER REFERENCES user(id) ON UPDATE CASCADE, "
|
||
|
|
"credential BLOB"
|
||
|
|
")"
|
||
|
|
)
|
||
|
|
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):
|
||
|
|
"""
|
||
|
|
Adds a new user.
|
||
|
|
"""
|
||
|
|
db_execute(
|
||
|
|
"INSERT INTO user(name) VALUES (?)",
|
||
|
|
(username,)
|
||
|
|
)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def set_user_token(username, token_hash):
|
||
|
|
"""
|
||
|
|
Sets a user's token hash.
|
||
|
|
"""
|
||
|
|
db_execute(
|
||
|
|
"UPDATE user SET token_hash = ? WHERE name = ?",
|
||
|
|
(token_hash, username)
|
||
|
|
)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def get_user(username):
|
||
|
|
"""
|
||
|
|
Returns a user entry.
|
||
|
|
"""
|
||
|
|
data = db_execute(
|
||
|
|
"SELECT * FROM user WHERE name = ?",
|
||
|
|
(username,)
|
||
|
|
).fetchone()
|
||
|
|
return data
|
||
|
|
|
||
|
|
|
||
|
|
def set_credential(user_id, credential):
|
||
|
|
"""
|
||
|
|
Adds a credential to the database.
|
||
|
|
"""
|
||
|
|
db_execute(
|
||
|
|
"INSERT INTO credential(user_id, credential) VALUES (?, ?)",
|
||
|
|
(user_id, credential)
|
||
|
|
)
|
||
|
|
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
|