user_management.py

import sqlite3
import bcrypt
import os
from datetime import datetime

DB_FILE = "user_auth.db"

# Initialize the database
def initialize_database():
    if not os.path.exists(DB_FILE):
        conn = sqlite3.connect(DB_FILE)
        c = conn.cursor()
        # Create users table
        c.execute("""
            CREATE TABLE users (
                username TEXT PRIMARY KEY,
                password_hash TEXT
            )
        """)
        # Create settings table with timestamp
        c.execute("""
            CREATE TABLE user_settings (
                username TEXT PRIMARY KEY,
                topic TEXT,
                instructions TEXT,
                timestamp TEXT,
                FOREIGN KEY (username) REFERENCES users (username)
            )
        """)
        conn.commit()
        conn.close()

# Add a new user to the database
def add_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
    try:
        c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", (username, password_hash))
        conn.commit()
    except sqlite3.IntegrityError:
        conn.close()
        return False  # Username already exists
    conn.close()
    return True

# Authenticate a user
def authenticate_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute("SELECT password_hash FROM users WHERE username = ?", (username,))
    row = c.fetchone()
    conn.close()
    if row and bcrypt.checkpw(password.encode(), row[0].encode()):
        return True
    return False

# Load settings for a user
def load_settings(username):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute("SELECT topic, instructions FROM user_settings WHERE username = ?", (username,))
    row = c.fetchone()
    conn.close()
    if row:
        return {"topic": row[0], "instructions": row[1]}
    return {"topic": "", "instructions": ""}

# Save settings for a user with timestamp
def save_settings(username, settings):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    c.execute("""
        INSERT INTO user_settings (username, topic, instructions, timestamp)
        VALUES (?, ?, ?, ?)
        ON CONFLICT(username)
        DO UPDATE SET topic = excluded.topic, instructions = excluded.instructions, timestamp = excluded.timestamp
    """, (username, settings["topic"], settings["instructions"], timestamp))
    conn.commit()
    conn.close()

# Dummy function to simulate prompt generation
def generate_prompt(topic, instructions):
    return f"Generated prompt based on topic '{topic}' and instructions: '{instructions}'."