vacuum-wall/lib/db.py

"""Abstract database layer for Vacuum Wall.

Provides query ID constants and an abstract Database baseclass so subsystems
interact with the database through opaque query identifiers, never raw SQL.
Backend implementations (SQLite, PostgreSQL) provide the actual SQL.

Usage:
    from lib.db import Q_INSERT_USER, Database, get_db

    class SQLiteBackend(Database):
        QUERY_MAP = {
            Q_INSERT_USER: "INSERT INTO users (username, password_hash) VALUES (?, ?)",
            ...
        }
"""

from __future__ import annotations

import logging
import os
import secrets
import threading
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, ClassVar

logger = logging.getLogger(__name__)

PROJECT_DIR = Path(__file__).resolve().parent.parent

# ---------------------------------------------------------------------------
# Query ID constants — single source of truth for all database operations
# ---------------------------------------------------------------------------

Q_INIT_TABLES = "init_tables"
Q_INSERT_USER = "insert_user"
Q_SELECT_USER_BY_NAME = "select_user_by_name"
Q_SELECT_USER_BY_ID = "select_user_by_id"
Q_UPDATE_PASSWORD = "update_password"
Q_DELETE_USER = "delete_user"
Q_UPSERT_PERMISSION = "upsert_permission"
Q_SELECT_PERMISSIONS = "select_permissions"
Q_DELETE_PERMISSIONS = "delete_permissions"
Q_DELETE_PERMISSION_SUBSYSTEM = "delete_permission_subsystem"
Q_INSERT_BLACKLIST = "insert_blacklist"
Q_SELECT_BLACKLIST = "select_blacklist_jti"
Q_DELETE_EXPIRED_BLACKLIST = "delete_expired_blacklist"
Q_UPSERT_REFRESH_TOKEN = "upsert_refresh_token"
Q_SELECT_REFRESH_TOKEN = "select_refresh_token"
Q_DELETE_REFRESH_TOKEN = "delete_refresh_token"
Q_INSERT_WEBAUTHN = "insert_webauthn"
Q_SELECT_WEBAUTHN_USER = "select_webauthn_user"
Q_SELECT_WEBAUTHN_ID = "select_webauthn_id"
Q_SELECT_WEBAUTHN_COUNTS = "select_webauthn_counts"
Q_DELETE_WEBAUTHN = "delete_webauthn"
Q_UPDATE_WEBAUTHN_SIGN_COUNT = "update_webauthn_sign_count"
Q_SELECT_ALL_USERS = "select_all_users"
Q_SELECT_USERS_WITH_PERMS = "select_users_with_perms"
Q_SELECT_USER_JWT_SECRET = "select_user_jwt_secret"
Q_UPDATE_JWT_SECRET = "update_jwt_secret"

# ---------------------------------------------------------------------------
# Schema DDL
# ---------------------------------------------------------------------------

INIT_SQL = """
CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT UNIQUE NOT NULL,
    password_hash TEXT NOT NULL,
    jwt_secret TEXT NOT NULL,
    created_at INTEGER NOT NULL DEFAULT (unixepoch())
    );

CREATE TABLE IF NOT EXISTS permissions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT NOT NULL REFERENCES users(username) ON DELETE CASCADE,
    subsystem TEXT NOT NULL,
    level TEXT NOT NULL CHECK (level IN ('read', 'rw')),
    UNIQUE(username, subsystem)
);

CREATE TABLE IF NOT EXISTS token_blacklist (
    jti TEXT PRIMARY KEY,
    token_type TEXT NOT NULL,
    expires INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS refresh_tokens (
    username TEXT UNIQUE NOT NULL,
    jti TEXT NOT NULL,
    issued_at INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS webauthn_creds (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT NOT NULL REFERENCES users(username) ON DELETE CASCADE,
    credential_id TEXT NOT NULL,
    public_key TEXT NOT NULL,
    sign_count INTEGER NOT NULL DEFAULT 0,
    name TEXT NOT NULL DEFAULT '',
    transports TEXT NOT NULL DEFAULT '[]',
    UNIQUE(username, credential_id)
);

CREATE TABLE IF NOT EXISTS init_sequence (
    seq INTEGER PRIMARY KEY
);
"""


class Transaction:
    """Context manager for database transactions.

    Provides BEGIN/COMMIT/ROLLBACK semantics. Auto-commit is suppressed
    inside the transaction block.

    Usage:
        with db.in_transaction() as tx:
            tx.run(Q_INSERT_USER, ("user1", "hash"))
            tx.run(Q_UPSERT_PERMISSION, ("user1", "firewall", "rw"))
    """

    def __init__(self, parent: Database) -> None:
        self._parent = parent

    def __enter__(self) -> Transaction:
        self._parent._begin()
        self._parent._suppress_auto_commit()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
        try:
            if exc_type is None:
                self._parent._commit()
            else:
                self._parent._rollback()
        finally:
            self._parent._restore_auto_commit()
        return False

    def query(self, query_id: str, params: tuple = ()) -> list[dict]:
        return self._parent.query(query_id, params)

    def run(self, query_id: str, params: tuple = ()) -> int:
        return self._parent.run(query_id, params)

    def run_one(self, query_id: str, params: tuple = ()) -> int | dict:
        return self._parent.run_one(query_id, params)


class Database(ABC):
    """Abstract database interface.

    All subsystems interact with the database through this interface.
    Queries are identified by string IDs (e.g. Q_INSERT_USER) — never
    raw SQL strings.

    Connection is cached via the ``conn`` property. Prepared statements
    are auto-cached on first use.
    """

    QUERY_MAP: ClassVar[dict[str, str]] = {}

    def __init__(self, connection_string: str) -> None:
        self._connection_string = connection_string
        self._prepared: dict[str, Any] = {}
        self._in_transaction = False
        # Per-thread connections: backend connection objects (e.g. sqlite3)
        # are bound to the thread that created them. The Flask WebUI runs
        # requests in worker threads while the daemon uses a single event-loop
        # thread, so each thread lazily gets its own connection.
        self._local = threading.local()

    @property
    def conn(self) -> Any:
        """Return this thread's cached database connection, creating it lazily."""
        conn = getattr(self._local, "conn", None)
        if conn is None:
            conn = self._connect(self._connection_string)
            self._local.conn = conn
        return conn

    @abstractmethod
    def _connect(self, cs: str) -> Any: ...

    @abstractmethod
    def _prepare(self, sql: str) -> Any: ...

    @abstractmethod
    def _execute(self, stmt: Any, params: tuple) -> tuple[list[dict], int]:
        """Execute a prepared statement. Returns (rows, rowcount)."""
        ...

    @abstractmethod
    def _last_insert_id(self, stmt: Any) -> int | dict: ...

    @abstractmethod
    def _execute_direct(self, sql: str) -> None:
        """Execute raw SQL without prepared statements (for DDL, not implemented by base)."""

    def _begin(self) -> None:  # noqa: B027
        """Begin a transaction (not implemented by base)."""

    def _commit(self) -> None:  # noqa: B027
        """Commit a transaction (not implemented by base)."""

    def _rollback(self) -> None:  # noqa: B027
        """Rollback a transaction (not implemented by base)."""

    def _suppress_auto_commit(self) -> None:  # noqa: B027
        """Suppress auto-commit (not implemented by base)."""

    def _restore_auto_commit(self) -> None:  # noqa: B027
        """Restore auto-commit (not implemented by base)."""

    def query(self, query_id: str, params: tuple = ()) -> list[dict]:
        """Execute a SELECT query. Returns list of row dicts."""
        stmt = self._get_prepared(query_id)
        rows, _ = self._execute(stmt, params)
        return rows

    def run(self, query_id: str, params: tuple = ()) -> int:
        """Execute an INSERT/UPDATE/DELETE. Returns affected row count."""
        stmt = self._get_prepared(query_id)
        _, count = self._execute(stmt, params)
        return count

    def run_one(self, query_id: str, params: tuple = ()) -> int | dict:
        """Execute and return the last insert ID or row dict."""
        stmt = self._get_prepared(query_id)
        _, _ = self._execute(stmt, params)
        return self._last_insert_id(stmt)

    def in_transaction(self) -> Transaction:
        """Return a transaction context manager."""
        return Transaction(self)

    def init_tables(self) -> None:
        """Create schema tables if they don't exist."""
        self._execute_direct(INIT_SQL)

    def _get_prepared(self, query_id: str) -> Any:
        if query_id not in self._prepared:
            if query_id not in self.QUERY_MAP:
                raise KeyError(f"Unknown query ID: {query_id!r}")
            self._prepared[query_id] = self._prepare(self.QUERY_MAP[query_id])
        return self._prepared[query_id]


# ---------------------------------------------------------------------------
# Singleton accessor
# ---------------------------------------------------------------------------

_db_instance: Database | None = None


def _get_backend_name() -> str:
    return os.environ.get("VACUUM_WALL_DB_BACKEND", "sqlite")


def _get_db_path() -> str:
    return os.environ.get("VACUUM_WALL_DB_PATH", str(PROJECT_DIR / "data" / "auth.db"))


def get_db() -> Database:
    """Return the singleton Database instance.

    Creates the instance on first call using the backend specified by
    ``VACUUM_WALL_DB_BACKEND`` env var (default: sqlite).

    Call this at application startup to ensure the DB is initialized.
    """
    global _db_instance
    if _db_instance is None:
        backend = _get_backend_name()
        if backend == "sqlite":
            from lib.db_sqlite import SQLiteBackend

            path = _get_db_path()
            _db_instance = SQLiteBackend(path)
            Path(path).parent.mkdir(parents=True, exist_ok=True)
            _db_instance.init_tables()
            _seed_builtin_admin(_db_instance)
        else:
            raise ValueError(f"Unknown database backend: {backend!r}")
    return _db_instance


def _seed_builtin_admin(db: Database) -> None:
    """Create the builtin admin user if they don't exist.

    The builtin admin has full (rw) access to all subsystems and cannot
    be deleted or have permissions modified through the normal API.
    """
    from lib.auth_users import ALL_SUBSYSTEMS, BUILTIN_ADMIN_USERNAME
    from lib.password import hash_password

    # Check if admin already exists
    rows = db.query(Q_SELECT_USER_BY_NAME, (BUILTIN_ADMIN_USERNAME,))
    if rows:
        return

    # Generate a random password — this fallback should only fire if
    # bootstrap_auth.py was skipped. Log the password prominently.
    random_password = secrets.token_urlsafe(24)
    placeholder_hash = hash_password(random_password)
    jwt_secret = secrets.token_urlsafe(32)

    try:
        with db.in_transaction() as tx:
            tx.run_one(
                Q_INSERT_USER, (BUILTIN_ADMIN_USERNAME, placeholder_hash, jwt_secret)
            )
            for subsystem in ALL_SUBSYSTEMS:
                tx.run(Q_UPSERT_PERMISSION, (BUILTIN_ADMIN_USERNAME, subsystem, "rw"))
    except Exception as exc:
        rows = db.query(Q_SELECT_USER_BY_NAME, (BUILTIN_ADMIN_USERNAME,))
        if not rows:
            raise
        logger.warning(
            "Concurrent builtin admin seed detected (%s); proceeding with existing user",
            exc,
        )
        return

    auth_log = Path("/var/log/vacuum-wall/auth.log")
    auth_log_written = False
    try:
        auth_log.write_text(
            f"Builtin admin password: {random_password}\n", encoding="utf-8"
        )
        import os as _os

        _os.chmod(str(auth_log), 0o600)
        auth_log_written = True
    except OSError:
        logger.error("Could not write admin password to %s", auth_log)
    logger.warning(
        "Builtin admin user created. THIS IS A FALLBACK — bootstrap_auth.py "
        "should have run during install. Admin password: %s... (%s)",
        random_password[:6],
        "see /var/log/vacuum-wall/auth.log for the full password"
        if auth_log_written
        else "full password NOT written — check the error above",
    )


def reset_db_for_test() -> None:
    """Reset the singleton — only for tests."""
    global _db_instance
    if _db_instance is not None:
        _db_instance = None