rewrite backend
This commit is contained in:
@@ -0,0 +1,429 @@
|
||||
import sqlite3
|
||||
from os import urandom
|
||||
from hashlib import sha256
|
||||
|
||||
# Initialize database
|
||||
def init(db_path: str) -> None:
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
# Auth token table
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS Tokens (
|
||||
TKHASH TEXT PRIMARY KEY NOT NULL,
|
||||
USERID INTEGER NOT NULL,
|
||||
NAME TEXT NOT NULL
|
||||
)
|
||||
''')
|
||||
|
||||
# Registered domains table
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS Domains (
|
||||
DOMAIN TEXT PRIMARY KEY NOT NULL,
|
||||
USERID INTEGER NOT NULL
|
||||
)
|
||||
''')
|
||||
|
||||
# Users table
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS Users (
|
||||
USERID INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
|
||||
USERNAME TEXT NOT NULL UNIQUE,
|
||||
PWHASH TEXT NOT NULL,
|
||||
PWSALT TEXT NOT NULL,
|
||||
EMAIL TEXT,
|
||||
IS_ADMIN INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
''')
|
||||
|
||||
# Create a new admin if there are no accounts
|
||||
cursor.execute('''
|
||||
SELECT COUNT(*) FROM Users
|
||||
''')
|
||||
|
||||
if cursor.fetchone()[0] == 0:
|
||||
username = "root"
|
||||
password = urandom(8).hex()
|
||||
|
||||
User.new(db_path, username, password, is_admin=True)
|
||||
|
||||
print(f"No users detected, created new admin\nUser: {username}\nPass: {password}", flush=True)
|
||||
|
||||
|
||||
def list_users(db_path: str) -> list[dict]:
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT userid, username, is_admin, email FROM Users
|
||||
''')
|
||||
|
||||
users = cursor.fetchall()
|
||||
|
||||
users = list(map(lambda x: {"userid": x[0], "username": x[1], "domains": [], "email": x[3], "is_admin": x[2]}, users))
|
||||
|
||||
for user in users:
|
||||
cursor.execute('''
|
||||
SELECT domain FROM Domains
|
||||
WHERE userid = ?
|
||||
''', (user["userid"], ))
|
||||
|
||||
domains = cursor.fetchall()
|
||||
|
||||
domains = list(map(lambda x: x[0], domains))
|
||||
|
||||
user['domains'] = domains
|
||||
|
||||
return users
|
||||
|
||||
|
||||
# Collection of getters and setters
|
||||
class User:
|
||||
def exists(db_path: str, userid: int) -> bool:
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT COUNT(*) FROM Users
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
return cursor.fetchone()[0] == 1
|
||||
|
||||
|
||||
def new(db_path: str, username: str, password: str, domains: list = [], is_admin: bool = False, email: str = None) -> int:
|
||||
username = username.strip()
|
||||
|
||||
if not username.isalnum():
|
||||
raise ValueError("Username has to be alphanumeric")
|
||||
|
||||
if len(username) < 1:
|
||||
raise ValueError("Username has to be at least 1 character long")
|
||||
|
||||
if len(password) > 16:
|
||||
raise ValueError("Username has to be at most 16 characters long")
|
||||
|
||||
if len(password) < 8:
|
||||
raise ValueError("Password has to be at least 8 characters long")
|
||||
|
||||
if len(password) > 128:
|
||||
raise ValueError("Password has to be at most 128 characters long")
|
||||
|
||||
if email is not None:
|
||||
email = email.strip()
|
||||
|
||||
if len(email) > 128:
|
||||
raise ValueError("Email has to be at most 128 characters long, be reasonable")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
pwsalt = urandom(16).hex()
|
||||
pwhash = sha256( (pwsalt + password).encode('utf-8') ).hexdigest()
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO Users (username, pwsalt, pwhash, email, is_admin) VALUES (?, ?, ?, ?, ?)
|
||||
''', (username, pwsalt, pwhash, email, is_admin))
|
||||
|
||||
# Get new user's ID
|
||||
cursor.execute('''
|
||||
SELECT userid FROM Users
|
||||
WHERE username = ?
|
||||
''', (username, ))
|
||||
|
||||
userid = cursor.fetchone()[0]
|
||||
|
||||
User.set_domains(db_path, userid, domains)
|
||||
|
||||
return userid
|
||||
|
||||
|
||||
def delete(db_path: str, userid: int):
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('DELETE FROM Domains WHERE userid = ?', (userid, ))
|
||||
cursor.execute('DELETE FROM Tokens WHERE userid = ?', (userid, ))
|
||||
cursor.execute('DELETE FROM Users WHERE userid = ?', (userid, ))
|
||||
|
||||
|
||||
def is_admin(db_path: str, userid: int) -> bool:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT is_admin FROM Users
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
return cursor.fetchone()[0]
|
||||
|
||||
def set_admin(db_path: str, userid: int, is_admin: bool):
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE Users
|
||||
SET is_admin = ?
|
||||
WHERE userid = ?
|
||||
''', (is_admin, userid))
|
||||
|
||||
def get_username(db_path: str, userid: int) -> str:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT username FROM Users
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
return cursor.fetchone()[0]
|
||||
|
||||
|
||||
def get_id(db_path: str, username: str) -> int:
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT userid FROM Users
|
||||
WHERE username = ?
|
||||
''', (username, ))
|
||||
|
||||
userid = cursor.fetchone()
|
||||
|
||||
if userid is None:
|
||||
raise ValueError(f"User {username} does not exist")
|
||||
|
||||
return userid[0]
|
||||
|
||||
def get_domains(db_path: str, userid: int) -> list[str]:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT domain FROM Domains
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
domains = cursor.fetchall()
|
||||
domains = list(map(lambda x: x[0], domains))
|
||||
|
||||
return domains
|
||||
|
||||
|
||||
def set_domains(db_path: str, userid: int, domains: list[str]) -> None:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
#TODO domain regex
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
DELETE FROM Domains
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
for domain in domains:
|
||||
domain = domain.strip()
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO Domains (domain, userid) VALUES (?, ?)
|
||||
''', (domain, userid))
|
||||
|
||||
|
||||
def get_token_names(db_path: str, userid: int) -> list[str]:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT name FROM Tokens
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
token_names = cursor.fetchall()
|
||||
token_names = list(map(lambda x: x[0], token_names))
|
||||
|
||||
return token_names
|
||||
|
||||
|
||||
def generate_token(db_path: str, userid: int, token_name: str) -> str | None:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
if not token_name.isalnum():
|
||||
raise ValueError("Token name has to be alphanumeric")
|
||||
|
||||
if len(token_name) < 1:
|
||||
raise ValueError("Token name has to be at least 1 character long")
|
||||
|
||||
if len(token_name) > 16:
|
||||
raise ValueError("Token name has to be at most 16 characters long")
|
||||
|
||||
token_name_list = User.get_token_names(db_path, userid)
|
||||
|
||||
if len(token_name_list) >= 5:
|
||||
raise ValueError("One user can have at most 5 tokens")
|
||||
|
||||
if token_name in token_name_list:
|
||||
raise ValueError("Token name is not unique")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
# Try 10 times in case of collisions
|
||||
for _ in range(10):
|
||||
token = urandom(32).hex()
|
||||
token_hash = sha256(token.encode("utf-8")).hexdigest()
|
||||
|
||||
try:
|
||||
cursor.execute('''
|
||||
INSERT INTO Tokens (tkhash, userid, name) VALUES (?, ?, ?)
|
||||
''', (token_hash, userid, token_name))
|
||||
except sqlite3.IntegrityError:
|
||||
continue
|
||||
|
||||
return token
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def revoke_token(db_path: str, userid: int, token_name: str) -> None:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
if token_name not in User.get_token_names(db_path, userid):
|
||||
raise ValueError("Token name not found")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
DELETE FROM Tokens
|
||||
WHERE userid = ?
|
||||
AND name = ?
|
||||
''', (userid, token_name))
|
||||
|
||||
|
||||
def get_password_hash(db_path: str, userid: int) -> tuple[str, str]:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT pwsalt, pwhash FROM Users
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
return cursor.fetchone()
|
||||
|
||||
|
||||
def set_password(db_path: str, userid: int, password: str) -> None:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
if len(password) < 8:
|
||||
raise ValueError("Password has to be at least 8 characters long")
|
||||
|
||||
if len(password) > 128:
|
||||
raise ValueError("Password has to be at most 128 characters long")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
pwsalt = urandom(16).hex()
|
||||
pwhash = sha256( (pwsalt + password).encode('utf-8') ).hexdigest()
|
||||
|
||||
cursor.execute('''
|
||||
UPdATE Users
|
||||
SET pwsalt = ?,
|
||||
pwhash = ?
|
||||
WHERE userid = ?
|
||||
''', (pwsalt, pwhash, userid))
|
||||
|
||||
def get_email(db_path: str, userid: int) -> str:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT email FROM Users
|
||||
WHERE userid = ?
|
||||
''', (userid, ))
|
||||
|
||||
return cursor.fetchone()[0]
|
||||
|
||||
def set_email(db_path: str, userid: int, email: str | None) -> None:
|
||||
if not User.exists(db_path, userid):
|
||||
raise ValueError(f"User {userid} does not exist")
|
||||
|
||||
if email is not None:
|
||||
email = email.strip()
|
||||
|
||||
if len(email) > 128:
|
||||
raise ValueError("Email has to be at most 128 characters long, be reasonable")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE Users
|
||||
SET email = ?
|
||||
WHERE userid = ?
|
||||
''', (email, userid, ))
|
||||
|
||||
|
||||
# Collection of getters
|
||||
class Token:
|
||||
def exists(db_path: str, token: str) -> bool:
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
tkhash = sha256(token.encode('utf-8')).hexdigest()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT COUNT(*) FROM Tokens
|
||||
WHERE tkhash = ?
|
||||
''', (tkhash, ))
|
||||
|
||||
return cursor.fetchone()[0] == 1
|
||||
|
||||
|
||||
def get_user(db_path: str, token: str) -> int:
|
||||
if not Token.exists(db_path, userid):
|
||||
raise ValueError("Token does not exist")
|
||||
|
||||
with sqlite3.connect(db_path) as connection:
|
||||
cursor = connection.cursor()
|
||||
|
||||
tkhash = sha256(token.encode('utf-8')).hexdigest()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT userid FROM Tokens
|
||||
WHERE tkhash = ?
|
||||
''', (tkhash, ))
|
||||
|
||||
return cursor.fetcone()[0]
|
||||
Reference in New Issue
Block a user