import sqlite3 from os import urandom from hashlib import sha256 def init_db(path: str) -> None: with sqlite3.connect(path) as connection: cursor = connection.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS Tokens ( TKHASH TEXT PRIMARY KEY NOT NULL, USERID INTEGER NOT NULL, NAME TEXT NOT NULL ) ''') cursor.execute(''' CREATE TABLE IF NOT EXISTS Domains ( DOMAIN TEXT PRIMARY KEY NOT NULL, USERID INTEGER NOT NULL ) ''') 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 ) ''') cursor.execute(''' SELECT COUNT(*) FROM Users ''') if cursor.fetchone()[0] == 0: username = "root" password = urandom(8).hex() create_user(path, username, password, is_admin=True) print(f"No users detected, created new admin\nUser: {username}\nPass: {password}", flush=True) def check_token_domain(db_path: str, token: str, domain: str) -> bool: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() token_hash = sha256(token.encode('utf-8')).hexdigest() cursor.execute(''' SELECT COUNT(*) FROM Tokens JOIN Domains ON Tokens.userid = Domains.userid WHERE tkhash = ? AND domain = ? ''', (token_hash, domain)) return cursor.fetchone()[0] == 1 def user_login(db_path: str, username: str, password: str) -> int | None: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT pwsalt FROM Users WHERE username = ? ''', (username,)) password_salt = cursor.fetchone() if password_salt is None: return None password_salt = password_salt[0] password_hash = sha256( (password_salt + password).encode('utf-8') ).hexdigest() cursor.execute(''' SELECT userid FROM Users WHERE username = ? AND pwhash = ? ''', (username, password_hash)) result = cursor.fetchone() if result is None: return None return result[0] def get_username(db_path: str, userid: int) -> str | None: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT username FROM Users WHERE userid = ? ''', (userid, )) username = cursor.fetchone() if username is None: return None return username[0] def get_user_domains(db_path: str, userid: int) -> list[str]: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT domain FROM Domains WHERE userid = ? ''', (userid, )) domains = cursor.fetchall() domains = map(lambda x: x[0], domains) return domains def get_user_tokens(db_path: str, userid: int) -> list[str]: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT name FROM Tokens WHERE userid = ? ''', (userid, )) tokens = cursor.fetchall() tokens = map(lambda x: x[0], tokens) return tokens def generate_user_token(db_path: str, userid: int, token_name: str) -> str | None: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? ''', (userid, )) if cursor.fetchone()[0] != 1: return None cursor.execute(''' SELECT COUNT(*) FROM Tokens WHERE userid = ? ''', (userid, )) if cursor.fetchone()[0] >= 5: return None cursor.execute(''' SELECT COUNT(*) FROM Tokens WHERE userid = ? AND name = ? ''', (userid, token_name)) if cursor.fetchone()[0] != 0: return None if not token_name.isalnum(): return None if len(token_name) >= 16: return None for _ in range(10): try: token = urandom(32).hex() token_hash = sha256(token.encode("utf-8")).hexdigest() cursor.execute(''' INSERT INTO Tokens (tkhash, userid, name) VALUES (?, ?, ?) ''', (token_hash, userid, token_name)) return token except sqlite3.IntegrityError: pass return None def revoke_user_token(db_path: str, userid: int, token_name: str) -> None: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' DELETE FROM Tokens WHERE userid = ? AND name = ? ''', (userid, token_name)) def change_user_password(db_path: str, userid: int, oldpass: str, newpass: str) -> bool: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? ''', (userid, )) if cursor.fetchone()[0] != 1: return False cursor.execute(''' SELECT pwsalt FROM Users WHERE userid = ? ''', (userid, )) salt = cursor.fetchone() if salt is None: return False salt = salt[0] oldpass_hash = sha256( (salt + oldpass).encode('utf-8')).hexdigest() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? AND pwhash = ? ''', (userid, oldpass_hash)) if cursor.fetchone()[0] != 1: return False return set_user_password(db_path, userid, newpass) def set_user_password(db_path: str, userid: int, newpass: str) -> bool: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? ''', (userid, )) if cursor.fetchone()[0] != 1: return False salt = urandom(16).hex() newpass_hash = sha256( (salt + newpass).encode('utf-8')).hexdigest() cursor.execute(''' UPDATE Users SET pwsalt = ?, pwhash = ? WHERE userid = ? ''', (salt, newpass_hash, userid)) return True def is_admin(db_path: str, userid: int) -> bool: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? AND IS_ADMIN = 1 ''', (userid, )) return cursor.fetchone()[0] == 1 def get_users(db_path: str) -> 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": bool(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 def create_user(db_path: str, username: str, password: str, domains: list = [], is_admin: bool = False, email: str = None): with sqlite3.connect(db_path) as connection: cursor = connection.cursor() username.strip() if email is not None: email.strip() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE username = ? ''', (username, )) if cursor.fetchone()[0] != 0: return "Username exists" if username == "": return "Username is empty" if not username.isalnum(): return "Username is not alphanumeric" if password == "": return "Password empty" if len(password) < 8: return "Password is shorter than 8 characters" salt = urandom(16).hex() password_hash = sha256( (salt + password).encode('utf-8')).hexdigest() cursor.execute(''' INSERT INTO Users (username, pwsalt, pwhash, email, is_admin) VALUES (?, ?, ?, ?, ?) ''', (username, salt, password_hash, email, is_admin)) cursor.execute(''' SELECT userid FROM Users WHERE username = ? ''', (username, )) userid = cursor.fetchone()[0] for domain in domains: domain = domain.strip() if domain == "": continue cursor.execute(''' SELECT COUNT(*) FROM Domains WHERE domain = ? ''', (domain, )) if cursor.fetchone()[0] != 0: continue cursor.execute(''' INSERT INTO Domains (userid, domain) VALUES (?, ?) ''', (userid, domain)) def update_user(db_path: str, userid: int, domains: list[str] = None, password: str = None, is_admin: bool = None, email: str = None) -> tuple[bool, str]: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() cursor.execute(''' SELECT COUNT(*) FROM Users WHERE userid = ? ''', (userid, )) if cursor.fetchone()[0] != 1: return False, "Uerid not found" if domains is not None: # crude check if domains is iterable for domain in domains: break cursor.execute(''' DELETE FROM Domains WHERE userid = ? ''', (userid, )) for domain in domains: cursor.execute(''' INSERT INTO Domains (userid, domain) VALUES (?, ?) ''', (userid, domain)) if password is not None: set_user_password(db_path, userid, password) if is_admin is not None: cursor.execute(''' UPDATE Users SET is_admin = ? WHERE userid = ? ''', (is_admin, userid)) if email is not None: if email == "": email = None cursor.execute(''' UPDATE Users SET email = ? WHERE userid = ? ''', (email, userid)) def delete_user(db_path: str, userid: int): 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,))