- New web UI with OpenWebUI-like interface using Tailwind CSS - SQLite-based authentication with session management - User registration, login, and profile pages - Chat interface with conversation history - Streaming responses with visible thinking phase - File attachments support - User document uploads in profile - Rate limiting (5 requests/day for free users) - Admin panel user management with promote/demote/delete - Custom color theme matching balloon logo design - Compatible with Nuitka build system
563 lines
18 KiB
Python
563 lines
18 KiB
Python
"""
|
|
Authentication Models and Database Management
|
|
SQLite-based user management with password hashing and rate limiting.
|
|
"""
|
|
import sqlite3
|
|
import json
|
|
import secrets
|
|
import hashlib
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional, List, Dict, Any
|
|
from pathlib import Path
|
|
from pydantic import BaseModel
|
|
from loguru import logger
|
|
|
|
from config import get_db_path
|
|
|
|
|
|
class User(BaseModel):
|
|
"""User model."""
|
|
id: str
|
|
username: str
|
|
email: str
|
|
password_hash: str
|
|
is_admin: bool = False
|
|
created_at: str
|
|
request_count: int = 0
|
|
request_limit: int = 5
|
|
last_request_date: Optional[str] = None
|
|
|
|
|
|
class Session(BaseModel):
|
|
"""Session model."""
|
|
token: str
|
|
user_id: str
|
|
created_at: str
|
|
expires_at: str
|
|
|
|
|
|
class AuthManager:
|
|
"""
|
|
Handles user authentication, sessions, and rate limiting.
|
|
|
|
Features:
|
|
- SQLite-based user storage
|
|
- Password hashing with salt
|
|
- Session token management
|
|
- Daily request rate limiting
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.db_path = get_db_path()
|
|
self._init_db()
|
|
self._ensure_admin_exists()
|
|
logger.info("Auth Manager initialized")
|
|
|
|
def _init_db(self) -> None:
|
|
"""Initialize the auth database schema."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
# Users table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id TEXT PRIMARY KEY,
|
|
username TEXT UNIQUE NOT NULL,
|
|
email TEXT UNIQUE NOT NULL,
|
|
password_hash TEXT NOT NULL,
|
|
is_admin INTEGER DEFAULT 0,
|
|
created_at TEXT,
|
|
request_count INTEGER DEFAULT 0,
|
|
request_limit INTEGER DEFAULT 5,
|
|
last_request_date TEXT
|
|
)
|
|
""")
|
|
|
|
# Sessions table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
token TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
created_at TEXT,
|
|
expires_at TEXT,
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
)
|
|
""")
|
|
|
|
# User documents table (for profile document uploads)
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS user_documents (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
filename TEXT NOT NULL,
|
|
file_type TEXT,
|
|
content BLOB,
|
|
created_at TEXT,
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def _ensure_admin_exists(self) -> None:
|
|
"""Create default admin user if not exists."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("SELECT id FROM users WHERE username = ?", ("admin",))
|
|
if not cursor.fetchone():
|
|
# Create default admin
|
|
admin_id = self._generate_id()
|
|
password_hash = self._hash_password("changeme")
|
|
now = datetime.utcnow().isoformat()
|
|
|
|
cursor.execute("""
|
|
INSERT INTO users (id, username, email, password_hash, is_admin, created_at, request_limit)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""", (admin_id, "admin", "admin@moxie.local", password_hash, 1, now, 999999))
|
|
|
|
conn.commit()
|
|
logger.info("Default admin user created (username: admin, password: changeme)")
|
|
|
|
conn.close()
|
|
|
|
def _generate_id(self) -> str:
|
|
"""Generate a unique ID."""
|
|
import uuid
|
|
return str(uuid.uuid4())
|
|
|
|
def _hash_password(self, password: str) -> str:
|
|
"""Hash a password with salt."""
|
|
salt = secrets.token_hex(16)
|
|
hash_val = hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000)
|
|
return f"{salt}:{hash_val.hex()}"
|
|
|
|
def _verify_password(self, password: str, stored_hash: str) -> bool:
|
|
"""Verify a password against stored hash."""
|
|
try:
|
|
salt, hash_val = stored_hash.split(':')
|
|
new_hash = hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000)
|
|
return new_hash.hex() == hash_val
|
|
except Exception:
|
|
return False
|
|
|
|
def _generate_token(self) -> str:
|
|
"""Generate a secure session token."""
|
|
return secrets.token_urlsafe(32)
|
|
|
|
def create_user(self, username: str, email: str, password: str) -> Optional[User]:
|
|
"""Create a new user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
user_id = self._generate_id()
|
|
password_hash = self._hash_password(password)
|
|
now = datetime.utcnow().isoformat()
|
|
|
|
cursor.execute("""
|
|
INSERT INTO users (id, username, email, password_hash, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""", (user_id, username, email, password_hash, now))
|
|
|
|
conn.commit()
|
|
logger.info(f"User created: {username}")
|
|
|
|
return User(
|
|
id=user_id,
|
|
username=username,
|
|
email=email,
|
|
password_hash=password_hash,
|
|
is_admin=False,
|
|
created_at=now,
|
|
request_count=0,
|
|
request_limit=5
|
|
)
|
|
except sqlite3.IntegrityError as e:
|
|
logger.warning(f"Failed to create user: {e}")
|
|
return None
|
|
finally:
|
|
conn.close()
|
|
|
|
def authenticate(self, username: str, password: str) -> Optional[User]:
|
|
"""Authenticate a user by username and password."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT id, username, email, password_hash, is_admin, created_at,
|
|
request_count, request_limit, last_request_date
|
|
FROM users WHERE username = ?
|
|
""", (username,))
|
|
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
if not self._verify_password(password, row[3]):
|
|
return None
|
|
|
|
return User(
|
|
id=row[0],
|
|
username=row[1],
|
|
email=row[2],
|
|
password_hash=row[3],
|
|
is_admin=bool(row[4]),
|
|
created_at=row[5],
|
|
request_count=row[6] or 0,
|
|
request_limit=row[7] or 5,
|
|
last_request_date=row[8]
|
|
)
|
|
|
|
def create_session(self, user_id: str, expires_days: int = 7) -> str:
|
|
"""Create a new session for a user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
token = self._generate_token()
|
|
now = datetime.utcnow()
|
|
expires = now + timedelta(days=expires_days)
|
|
|
|
cursor.execute("""
|
|
INSERT INTO sessions (token, user_id, created_at, expires_at)
|
|
VALUES (?, ?, ?, ?)
|
|
""", (token, user_id, now.isoformat(), expires.isoformat()))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return token
|
|
|
|
def validate_session(self, token: str) -> Optional[User]:
|
|
"""Validate a session token and return the user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
now = datetime.utcnow().isoformat()
|
|
|
|
cursor.execute("""
|
|
SELECT s.user_id, u.username, u.email, u.password_hash, u.is_admin,
|
|
u.created_at, u.request_count, u.request_limit, u.last_request_date
|
|
FROM sessions s
|
|
JOIN users u ON s.user_id = u.id
|
|
WHERE s.token = ? AND s.expires_at > ?
|
|
""", (token, now))
|
|
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
return User(
|
|
id=row[0],
|
|
username=row[1],
|
|
email=row[2],
|
|
password_hash=row[3],
|
|
is_admin=bool(row[4]),
|
|
created_at=row[5],
|
|
request_count=row[6] or 0,
|
|
request_limit=row[7] or 5,
|
|
last_request_date=row[8]
|
|
)
|
|
|
|
def delete_session(self, token: str) -> None:
|
|
"""Delete a session (logout)."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("DELETE FROM sessions WHERE token = ?", (token,))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def check_rate_limit(self, user_id: str) -> Dict[str, Any]:
|
|
"""
|
|
Check if user can make a request.
|
|
Returns dict with allowed status and remaining requests.
|
|
"""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
today = datetime.utcnow().date().isoformat()
|
|
|
|
cursor.execute("""
|
|
SELECT request_count, request_limit, last_request_date
|
|
FROM users WHERE id = ?
|
|
""", (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
if not row:
|
|
conn.close()
|
|
return {"allowed": False, "remaining": 0, "limit": 0}
|
|
|
|
request_count, request_limit, last_request_date = row
|
|
|
|
# Reset count if it's a new day
|
|
if last_request_date != today:
|
|
request_count = 0
|
|
|
|
# Admins have unlimited requests
|
|
cursor.execute("SELECT is_admin FROM users WHERE id = ?", (user_id,))
|
|
is_admin = cursor.fetchone()[0]
|
|
|
|
if is_admin:
|
|
conn.close()
|
|
return {"allowed": True, "remaining": 999999, "limit": 999999}
|
|
|
|
allowed = request_count < request_limit
|
|
remaining = max(0, request_limit - request_count)
|
|
|
|
conn.close()
|
|
|
|
return {
|
|
"allowed": allowed,
|
|
"remaining": remaining,
|
|
"limit": request_limit
|
|
}
|
|
|
|
def increment_request_count(self, user_id: str) -> None:
|
|
"""Increment the request count for a user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
today = datetime.utcnow().date().isoformat()
|
|
|
|
# Check if we need to reset for new day
|
|
cursor.execute("""
|
|
SELECT last_request_date FROM users WHERE id = ?
|
|
""", (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
if row and row[0] != today:
|
|
# New day, reset count
|
|
cursor.execute("""
|
|
UPDATE users SET request_count = 1, last_request_date = ?
|
|
WHERE id = ?
|
|
""", (today, user_id))
|
|
else:
|
|
# Same day, increment
|
|
cursor.execute("""
|
|
UPDATE users SET request_count = request_count + 1, last_request_date = ?
|
|
WHERE id = ?
|
|
""", (today, user_id))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_user(self, user_id: str) -> Optional[User]:
|
|
"""Get a user by ID."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT id, username, email, password_hash, is_admin, created_at,
|
|
request_count, request_limit, last_request_date
|
|
FROM users WHERE id = ?
|
|
""", (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
return User(
|
|
id=row[0],
|
|
username=row[1],
|
|
email=row[2],
|
|
password_hash=row[3],
|
|
is_admin=bool(row[4]),
|
|
created_at=row[5],
|
|
request_count=row[6] or 0,
|
|
request_limit=row[7] or 5,
|
|
last_request_date=row[8]
|
|
)
|
|
|
|
def get_all_users(self) -> List[Dict[str, Any]]:
|
|
"""Get all users (for admin)."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT id, username, email, is_admin, created_at, request_count, request_limit
|
|
FROM users ORDER BY created_at DESC
|
|
""")
|
|
|
|
users = []
|
|
for row in cursor.fetchall():
|
|
users.append({
|
|
"id": row[0],
|
|
"username": row[1],
|
|
"email": row[2],
|
|
"is_admin": bool(row[3]),
|
|
"created_at": row[4],
|
|
"request_count": row[5] or 0,
|
|
"request_limit": row[6] or 5
|
|
})
|
|
|
|
conn.close()
|
|
return users
|
|
|
|
def promote_to_admin(self, user_id: str) -> bool:
|
|
"""Promote a user to admin."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("UPDATE users SET is_admin = 1 WHERE id = ?", (user_id,))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
def demote_from_admin(self, user_id: str) -> bool:
|
|
"""Demote a user from admin."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
# Prevent demoting the last admin
|
|
cursor.execute("SELECT COUNT(*) FROM users WHERE is_admin = 1")
|
|
admin_count = cursor.fetchone()[0]
|
|
|
|
cursor.execute("SELECT username FROM users WHERE id = ?", (user_id,))
|
|
row = cursor.fetchone()
|
|
if row and row[0] == "admin":
|
|
conn.close()
|
|
return False # Cannot demote the default admin
|
|
|
|
if admin_count <= 1:
|
|
conn.close()
|
|
return False # Cannot demote the last admin
|
|
|
|
cursor.execute("UPDATE users SET is_admin = 0 WHERE id = ?", (user_id,))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
def update_user_limit(self, user_id: str, limit: int) -> bool:
|
|
"""Update a user's daily request limit."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("UPDATE users SET request_limit = ? WHERE id = ?", (limit, user_id))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
def change_password(self, user_id: str, new_password: str) -> bool:
|
|
"""Change a user's password."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
password_hash = self._hash_password(new_password)
|
|
cursor.execute("UPDATE users SET password_hash = ? WHERE id = ?", (password_hash, user_id))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
def delete_user(self, user_id: str) -> bool:
|
|
"""Delete a user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
# Prevent deleting the default admin
|
|
cursor.execute("SELECT username FROM users WHERE id = ?", (user_id,))
|
|
row = cursor.fetchone()
|
|
if row and row[0] == "admin":
|
|
conn.close()
|
|
return False
|
|
|
|
# Delete sessions first
|
|
cursor.execute("DELETE FROM sessions WHERE user_id = ?", (user_id,))
|
|
# Delete user documents
|
|
cursor.execute("DELETE FROM user_documents WHERE user_id = ?", (user_id,))
|
|
# Delete user
|
|
cursor.execute("DELETE FROM users WHERE id = ?", (user_id,))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
# User document management
|
|
def add_user_document(self, user_id: str, filename: str, file_type: str, content: bytes) -> str:
|
|
"""Add a document to a user's profile."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
doc_id = self._generate_id()
|
|
now = datetime.utcnow().isoformat()
|
|
|
|
cursor.execute("""
|
|
INSERT INTO user_documents (id, user_id, filename, file_type, content, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""", (doc_id, user_id, filename, file_type, content, now))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return doc_id
|
|
|
|
def get_user_documents(self, user_id: str) -> List[Dict[str, Any]]:
|
|
"""Get all documents for a user."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT id, filename, file_type, created_at, LENGTH(content) as size
|
|
FROM user_documents WHERE user_id = ? ORDER BY created_at DESC
|
|
""", (user_id,))
|
|
|
|
docs = []
|
|
for row in cursor.fetchall():
|
|
docs.append({
|
|
"id": row[0],
|
|
"filename": row[1],
|
|
"file_type": row[2],
|
|
"created_at": row[3],
|
|
"size": row[4]
|
|
})
|
|
|
|
conn.close()
|
|
return docs
|
|
|
|
def delete_user_document(self, user_id: str, doc_id: str) -> bool:
|
|
"""Delete a user document."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("DELETE FROM user_documents WHERE id = ? AND user_id = ?", (doc_id, user_id))
|
|
|
|
success = cursor.rowcount > 0
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return success
|
|
|
|
|
|
# Global auth manager instance
|
|
_auth_manager: Optional[AuthManager] = None
|
|
|
|
|
|
def get_auth_manager() -> AuthManager:
|
|
"""Get or create the global auth manager instance."""
|
|
global _auth_manager
|
|
if _auth_manager is None:
|
|
_auth_manager = AuthManager()
|
|
return _auth_manager
|