izbi-dns/app/auth.py

import sqlite3
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
            )
        ''')

        connection.commit()
        cursor.close()

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) -> bool:
    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 False
        password_salt = password_salt[0]

        password_hash = sha256( (password_salt + password).encode('utf-8') ).hexdigest()

        cursor.execute('''
            SELECT COUNT(*) FROM Users
            WHERE username = ?
                AND pwhash = ?
        ''', (username, password_hash))

        return cursor.fetchone()[0] == 1