""" SQLite storage for authentication data (user credentials + JWT secret). """ import hashlib import sqlite3 from datetime import datetime, timezone from pathlib import Path from typing import Optional, Tuple DB_PATH = Path(__file__).parent / "auth.db" def _hash_token(token: str) -> str: """SHA-256 hash a setup token for safe storage.""" return hashlib.sha256(token.encode("utf-8")).hexdigest() def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" conn = sqlite3.connect(DB_PATH) conn.row_factory = sqlite3.Row conn.execute( """ CREATE TABLE IF NOT EXISTS auth_user ( id INTEGER PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_salt TEXT NOT NULL, password_hash TEXT NOT NULL, jwt_secret TEXT NOT NULL ); """ ) conn.execute( """ CREATE TABLE IF NOT EXISTS setup_tokens ( id INTEGER PRIMARY KEY, token_hash TEXT NOT NULL ); """ ) conn.execute( """ CREATE TABLE IF NOT EXISTS refresh_tokens ( id INTEGER PRIMARY KEY, token_hash TEXT NOT NULL, username TEXT NOT NULL, expires_at TEXT NOT NULL ); """ ) conn.commit() return conn def is_initialized() -> bool: """Check if auth has been set up (user exists in DB).""" conn = get_connection() cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user") row = cur.fetchone() conn.close() return bool(row["c"]) def create_initial_user(username: str, password: str, jwt_secret: str) -> None: """ Create the initial admin user in the database. Raises sqlite3.IntegrityError if username already exists. """ from .hashing import hash_password salt, pwd_hash = hash_password(password) conn = get_connection() try: conn.execute( """ INSERT INTO auth_user (username, password_salt, password_hash, jwt_secret) VALUES (?, ?, ?, ?) """, (username, salt, pwd_hash, jwt_secret), ) conn.commit() finally: conn.close() def delete_user(username: str) -> None: """ Delete a user from the database. Used for rollback when setup fails after user creation. """ conn = get_connection() try: conn.execute("DELETE FROM auth_user WHERE username = ?", (username,)) conn.commit() finally: conn.close() def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]: """ Get user's password salt, hash, and JWT secret. Returns (password_salt, password_hash, jwt_secret) or None if user not found. """ conn = get_connection() try: cur = conn.execute( """ SELECT password_salt, password_hash, jwt_secret FROM auth_user WHERE username = ? """, (username,), ) row = cur.fetchone() if not row: return None return row["password_salt"], row["password_hash"], row["jwt_secret"] finally: conn.close() def load_jwt_secret() -> str: """ Load the JWT secret from the database. Raises RuntimeError if auth is not initialized. """ conn = get_connection() try: cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1") row = cur.fetchone() if not row: raise RuntimeError("Auth is not initialized. Please set up a password first.") return row["jwt_secret"] finally: conn.close() def save_setup_token(token: str) -> None: """ Store a hashed setup token, replacing any existing one. """ token_hash = _hash_token(token) conn = get_connection() try: conn.execute("DELETE FROM setup_tokens") conn.execute("INSERT INTO setup_tokens (token_hash) VALUES (?)", (token_hash,)) conn.commit() finally: conn.close() def consume_setup_token(token: str) -> bool: """ Verify a setup token and delete it if valid. Returns True if the token was valid (and is now consumed), False otherwise. """ token_hash = _hash_token(token) conn = get_connection() try: cur = conn.execute( "SELECT id FROM setup_tokens WHERE token_hash = ?", (token_hash,) ) row = cur.fetchone() if row is None: return False conn.execute("DELETE FROM setup_tokens WHERE id = ?", (row["id"],)) conn.commit() return True finally: conn.close() def has_pending_setup_token() -> bool: """Check if a setup token is waiting to be consumed.""" conn = get_connection() try: cur = conn.execute("SELECT COUNT(*) AS c FROM setup_tokens") row = cur.fetchone() return bool(row["c"]) finally: conn.close() def save_refresh_token(token: str, username: str, expires_at: str) -> None: """ Store a hashed refresh token with its associated username and expiry. """ token_hash = _hash_token(token) conn = get_connection() try: conn.execute( """ INSERT INTO refresh_tokens (token_hash, username, expires_at) VALUES (?, ?, ?) """, (token_hash, username, expires_at), ) conn.commit() finally: conn.close() def verify_refresh_token(token: str) -> Optional[str]: """ Verify a refresh token and return the username. Returns the username if valid and not expired, None otherwise. The token is NOT consumed — it stays valid until it expires. """ token_hash = _hash_token(token) conn = get_connection() try: # Clean up any expired tokens while we're here conn.execute( "DELETE FROM refresh_tokens WHERE expires_at < ?", (datetime.now(timezone.utc).isoformat(),), ) conn.commit() cur = conn.execute( """ SELECT id, username, expires_at FROM refresh_tokens WHERE token_hash = ? """, (token_hash,), ) row = cur.fetchone() if row is None: return None # Check expiry expires_at = datetime.fromisoformat(row["expires_at"]) if datetime.now(timezone.utc) > expires_at: conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],)) conn.commit() return None return row["username"] finally: conn.close() def revoke_user_refresh_tokens(username: str) -> None: """Revoke all refresh tokens for a user (e.g. on logout).""" conn = get_connection() try: conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,)) conn.commit() finally: conn.close()