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() if len(domain) < 1: continue 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 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") if not token_name.isalnum(): raise ValueError("Token name has to be alphanumeric") 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]