izbi-dns/app/db.py

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]