db_utils.py

import os
import sqlite3
import hashlib
import base64
from cryptography.fernet import Fernet

DB_PATH = os.path.join(os.path.dirname(__file__), 'users.db')

def get_connection():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            salt TEXT NOT NULL
        )
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS projects (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            name TEXT NOT NULL,
            encrypted_data BLOB NOT NULL,
            FOREIGN KEY(user_id) REFERENCES users(id)
        )
    ''')
    conn.commit()
    conn.close()

def hash_password(password: str, salt: bytes = None):
    if salt is None:
        salt = os.urandom(16)
    pwd_hash = hashlib.pbkdf2_hmac('sha256', password.encode(), salt, 100_000)
    return base64.b64encode(pwd_hash).decode(), base64.b64encode(salt).decode()

def verify_password(stored_hash: str, stored_salt: str, password: str) -> bool:
    salt = base64.b64decode(stored_salt.encode())
    pwd_hash, _ = hash_password(password, salt)
    return pwd_hash == stored_hash

def derive_key(password: str, salt: bytes) -> bytes:
    # Derive a Fernet key (32 url-safe base64-encoded bytes)
    pwd_hash = hashlib.pbkdf2_hmac('sha256', password.encode(), salt, 100_000)
    return base64.urlsafe_b64encode(pwd_hash)

def encrypt_data(key: bytes, data: str) -> bytes:
    f = Fernet(key)
    return f.encrypt(data.encode())

def decrypt_data(key: bytes, token: bytes) -> str:
    f = Fernet(key)
    return f.decrypt(token).decode()

# Initialize DB on import
init_db()