"""PostgreSQL database models, initialization, and user-scoped helpers. Multi-user design (unchanged from the SQLite era): - The `users` table stores every account. Root (ID 0) is the administrator and is created automatically on first boot (fresh DB) or when migrating an existing single-user database (all pre-existing data becomes Root's data). - Every per-user table carries a `user_id INTEGER NOT NULL DEFAULT 0` column, so legacy single-user rows are automatically attributed to Root on import and any NEW column added in the future gets a DEFAULT (the migration fills missing columns with their SCHEMA defaults). - `auth_config` is keyed by `user_id` (legacy `id=1` row maps to user 0). The data-access layer talks PostgreSQL through ``backend.pg_compat``, a thin sqlite3-compatible shim (``?`` placeholders, ``INSERT OR IGNORE``, ``datetime('now')``, ``cursor.lastrowid``, ``conn.total_changes`` are all translated transparently). """ import re import sys import time from backend.pg_compat import connect as _pg_connect # ─── PostgreSQL schema ─────────────────────── # All new columns must carry a DEFAULT (migration contract: missing columns are # filled with their SCHEMA default when the DB is upgraded). SCHEMA = """ -- Global schema version (single row). CREATE TABLE IF NOT EXISTS meta ( id INTEGER PRIMARY KEY DEFAULT 1, version INTEGER NOT NULL DEFAULT 1 ); -- Auth tokens (dual-token auth): access (short-lived, per-request) and -- refresh (long-lived, stored in an httpOnly cookie). Only hashes are stored. CREATE TABLE IF NOT EXISTS auth_tokens ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, token_hash TEXT NOT NULL, kind TEXT NOT NULL DEFAULT 'access' CHECK(kind IN ('access','refresh')), view TEXT NOT NULL DEFAULT 'display' CHECK(view IN ('control','display')), temp INTEGER NOT NULL DEFAULT 0, session_version INTEGER NOT NULL DEFAULT 0, reauth_at INTEGER NOT NULL DEFAULT 0, created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')), expires_at INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS idx_auth_tokens_hash ON auth_tokens(token_hash); CREATE INDEX IF NOT EXISTS idx_auth_tokens_user ON auth_tokens(user_id); CREATE TABLE IF NOT EXISTS users ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, remark TEXT NOT NULL DEFAULT '', role TEXT NOT NULL DEFAULT 'user' CHECK(role IN ('admin','user')), is_temp INTEGER NOT NULL DEFAULT 0, temp_totp_secret TEXT DEFAULT '', status TEXT NOT NULL DEFAULT 'active' CHECK(status IN ('active','pending')), session_version INTEGER NOT NULL DEFAULT 0, theme TEXT NOT NULL DEFAULT 'dark' CHECK(theme IN ('dark','light','auto')), created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')) ); CREATE TABLE IF NOT EXISTS auth_config ( user_id INTEGER PRIMARY KEY DEFAULT 0, passkey_credential_id TEXT DEFAULT '', passkey_public_key TEXT DEFAULT '', totp_secret TEXT DEFAULT '', totp_enabled INTEGER DEFAULT 0, recovery_codes TEXT DEFAULT '', temp_password TEXT DEFAULT '', must_configure_mfa INTEGER DEFAULT 1 ); CREATE TABLE IF NOT EXISTS calendar_groups ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, name TEXT NOT NULL, color TEXT DEFAULT '#4A90D9', enabled INTEGER DEFAULT 1, created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')) ); CREATE TABLE IF NOT EXISTS calendar_sources ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, group_id INTEGER NOT NULL, type TEXT NOT NULL CHECK(type IN ('upload', 'subscription', 'manual')), name TEXT NOT NULL, url TEXT, ics_data TEXT, last_sync TEXT, FOREIGN KEY (group_id) REFERENCES calendar_groups(id) ); CREATE TABLE IF NOT EXISTS calendar_events ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, source_id INTEGER NOT NULL, uid TEXT, summary TEXT, description TEXT, location TEXT, dtstart TEXT, dtend TEXT, is_all_day INTEGER DEFAULT 0, FOREIGN KEY (source_id) REFERENCES calendar_sources(id) ); CREATE TABLE IF NOT EXISTS reminders ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, calendar_group_id INTEGER NOT NULL, minutes_before INTEGER DEFAULT 15, enabled INTEGER DEFAULT 1, FOREIGN KEY (calendar_group_id) REFERENCES calendar_groups(id) ); CREATE TABLE IF NOT EXISTS weather_config ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, api_key TEXT DEFAULT '', city TEXT DEFAULT 'Beijing', lat REAL, lon REAL, units TEXT DEFAULT 'metric', update_interval INTEGER DEFAULT 1800, jwt_private_key TEXT DEFAULT '', jwt_public_key TEXT DEFAULT '', credential_id TEXT DEFAULT '', project_id TEXT DEFAULT '', qweather_host TEXT DEFAULT '' ); CREATE TABLE IF NOT EXISTS weather_cache ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, data TEXT, updated_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')) ); CREATE TABLE IF NOT EXISTS weather_hourly_cache ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, data TEXT, fetched_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')) ); CREATE TABLE IF NOT EXISTS weather_minutely_cache ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, data TEXT, fetched_at TEXT ); CREATE TABLE IF NOT EXISTS weather_warning_cache ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, data TEXT, fetched_at TEXT ); CREATE TABLE IF NOT EXISTS clock_config ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, city TEXT DEFAULT 'Beijing', timezone TEXT DEFAULT 'Asia/Shanghai', is_primary INTEGER DEFAULT 1 ); CREATE TABLE IF NOT EXISTS alarms ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, label TEXT DEFAULT '', time TEXT NOT NULL, days TEXT DEFAULT '1,2,3,4,5,6,7', enabled INTEGER DEFAULT 1, sound TEXT DEFAULT 'default', once INTEGER DEFAULT 0 ); CREATE TABLE IF NOT EXISTS countdowns ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, label TEXT DEFAULT '', duration_seconds INTEGER DEFAULT 300, end_time TEXT DEFAULT '', enabled INTEGER DEFAULT 1, sound TEXT DEFAULT 'default' ); CREATE TABLE IF NOT EXISTS display_config ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, mode TEXT NOT NULL DEFAULT 'default' CHECK(mode IN ('default', 'custom')), custom_html TEXT DEFAULT '', custom_html_1 TEXT DEFAULT '', custom_html_2 TEXT DEFAULT '', custom_html_3 TEXT DEFAULT '' ); CREATE TABLE IF NOT EXISTS pomodoro_state ( id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, user_id INTEGER NOT NULL DEFAULT 0, running INTEGER DEFAULT 0, phase TEXT DEFAULT 'work' CHECK(phase IN ('work', 'break')), work_minutes INTEGER DEFAULT 25, break_minutes INTEGER DEFAULT 5, long_break_minutes INTEGER DEFAULT 15, cycles_before_long_break INTEGER DEFAULT 4, end_time TEXT, remaining_seconds INTEGER DEFAULT 1500, cycles INTEGER DEFAULT 0, enabled INTEGER DEFAULT 0, display_controls_enabled INTEGER DEFAULT 1, cycles_date TEXT DEFAULT '', work_label TEXT DEFAULT '工作中', break_label TEXT DEFAULT '休息中' ); CREATE UNIQUE INDEX IF NOT EXISTS idx_pomodoro_state_user ON pomodoro_state(user_id); """ # Ordered list of all user-data tables (children first) for FK-safe deletion. _DELETE_ORDER = [ "calendar_events", "calendar_sources", "reminders", "calendar_groups", "weather_cache", "weather_hourly_cache", "weather_minutely_cache", "weather_warning_cache", "weather_config", "clock_config", "alarms", "countdowns", "display_config", "pomodoro_state", "auth_config", "auth_tokens", ] def get_db(db_path: str = "data.db"): """Open a PostgreSQL connection (sqlite3-compatible wrapper). `db_path` is a PostgreSQL connection string (kept as the historical parameter name to minimise churn across the codebase). """ return _pg_connect(db_path) def _db_path(preferred: str = None) -> str: """Get the actual database connection string, preferring current_app config.""" if preferred: return preferred try: from flask import current_app return current_app.config.get("DB_PATH", "data.db") except Exception: return "data.db" # ─── Users ─────────────────────────────────── def get_users(db_path: str = None) -> list: """Return all users ordered by id.""" conn = get_db(_db_path(db_path)) rows = [dict(r) for r in conn.execute("SELECT * FROM users ORDER BY id").fetchall()] conn.close() return rows def get_active_user_ids(db_path: str = None) -> list: """Return ids of all active (non-pending) users. Used by the scheduler.""" conn = get_db(_db_path(db_path)) rows = conn.execute("SELECT id FROM users WHERE status='active' ORDER BY id").fetchall() conn.close() return [r["id"] for r in rows] def get_user(user_id: int, db_path: str = None) -> dict | None: conn = get_db(_db_path(db_path)) row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return dict(row) if row else None def user_exists(user_id: int, db_path: str = None) -> bool: conn = get_db(_db_path(db_path)) row = conn.execute("SELECT 1 FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return row is not None def is_admin_user(user_id: int, db_path: str = None) -> bool: """Root is always admin; otherwise check the role column.""" if user_id == 0: return True conn = get_db(_db_path(db_path)) row = conn.execute("SELECT role FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return bool(row and row["role"] == "admin") def create_user(remark: str = "", role: str = "user", is_temp: int = 0, temp_totp_secret: str = "", status: str = "active", db_path: str = None) -> int: """Create a user. Returns the new user id. The remark defaults to the user's own id (the only place an id is shown). """ conn = get_db(_db_path(db_path)) cur = conn.execute( "INSERT INTO users (remark, role, is_temp, temp_totp_secret, status) VALUES (?, ?, ?, ?, ?)", (remark, role, 1 if is_temp else 0, temp_totp_secret, status) ) uid = cur.lastrowid if not remark: conn.execute("UPDATE users SET remark=? WHERE id=?", (str(uid), uid)) conn.commit() conn.close() return uid def update_user(user_id: int, remark: str = None, role: str = None, is_temp: int = None, temp_totp_secret: str = None, status: str = None, theme: str = None, db_path: str = None): conn = get_db(_db_path(db_path)) fields = [] vals = [] for key, val in [("remark", remark), ("role", role), ("is_temp", is_temp), ("temp_totp_secret", temp_totp_secret), ("status", status), ("theme", theme)]: if val is not None: fields.append(f"{key}=?") vals.append(val) if fields: vals.append(user_id) conn.execute(f"UPDATE users SET {', '.join(fields)} WHERE id=?", vals) conn.commit() conn.close() def delete_user(user_id: int, db_path: str = None): """Delete a user and ALL of their data. Root (id 0) cannot be deleted. Deletion order respects foreign keys (children before parents). """ if user_id == 0: return conn = get_db(_db_path(db_path)) try: for table in _DELETE_ORDER: conn.execute(f"DELETE FROM {table} WHERE user_id=?", (user_id,)) conn.execute("DELETE FROM users WHERE id=?", (user_id,)) conn.commit() finally: conn.close() def get_user_session_version(user_id: int, db_path: str = None) -> int: conn = get_db(_db_path(db_path)) row = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return row["session_version"] if row else -1 def bump_user_session_version(user_id: int, db_path: str = None) -> int: """Invalidate all sessions of one user by incrementing their version.""" conn = get_db(_db_path(db_path)) conn.execute("UPDATE users SET session_version = session_version + 1 WHERE id=?", (user_id,)) conn.commit() row = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return row["session_version"] if row else -1 def bump_all_session_versions(db_path: str = None) -> int: """Invalidate every session (admin 'kick all' action).""" conn = get_db(_db_path(db_path)) conn.execute("UPDATE users SET session_version = session_version + 1") conn.commit() conn.close() return 0 def is_temp_account(user_id: int, db_path: str = None) -> bool: """True if the account is a temp account (awaiting first login-method setup).""" conn = get_db(_db_path(db_path)) row = conn.execute("SELECT is_temp FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return bool(row and row["is_temp"]) def get_temp_totp_secret(user_id: int, db_path: str = None) -> str: conn = get_db(_db_path(db_path)) row = conn.execute("SELECT temp_totp_secret FROM users WHERE id=?", (user_id,)).fetchone() conn.close() return (row["temp_totp_secret"] or "") if row else "" def finalize_temp_account(user_id: int, db_path: str = None): """Convert a temp account into a formal account: invalidates the temp TOTP, clears the temp flag, and revokes every temp=1 token (a temp refresh must never mint formal tokens afterwards). """ conn = get_db(_db_path(db_path)) conn.execute( "UPDATE users SET is_temp=0, temp_totp_secret='' WHERE id=?", (user_id,) ) conn.execute( "UPDATE auth_config SET must_configure_mfa=0 WHERE user_id=?", (user_id,) ) conn.execute( "DELETE FROM auth_tokens WHERE user_id=? AND temp=1", (user_id,) ) conn.commit() conn.close() # ─── Display config (per user) ────────────── def _ensure_display_row(conn, user_id: int): row = conn.execute("SELECT id FROM display_config WHERE user_id=?", (user_id,)).fetchone() if row is None: conn.execute( "INSERT INTO display_config (user_id, mode) VALUES (?, 'default')", (user_id,) ) def get_display_config(db_path: str = None, user_id: int = 0) -> dict: db = get_db(_db_path(db_path)) row = db.execute( "SELECT mode, custom_html, custom_html_1, custom_html_2, custom_html_3 " "FROM display_config WHERE user_id=?", (user_id,) ).fetchone() if row is None: db.execute( "INSERT INTO display_config (user_id, mode, custom_html, custom_html_1, custom_html_2, custom_html_3) " "VALUES (?, 'default', '', '', '', '')", (user_id,) ) db.commit() cfg = {"mode": "default", "custom_html": "", "custom_html_1": "", "custom_html_2": "", "custom_html_3": ""} else: cfg = { "mode": row["mode"], "custom_html": row["custom_html"] or "", "custom_html_1": row["custom_html_1"] or "", "custom_html_2": row["custom_html_2"] or "", "custom_html_3": row["custom_html_3"] or "", } db.close() return cfg def set_display_mode(mode: str, db_path: str = None, user_id: int = 0) -> bool: if mode not in ("default", "custom"): return False db = get_db(_db_path(db_path)) _ensure_display_row(db, user_id) db.execute("UPDATE display_config SET mode=? WHERE user_id=?", (mode, user_id)) db.commit() db.close() return True def set_custom_html(html: str, db_path: str = None, user_id: int = 0) -> bool: db = get_db(_db_path(db_path)) _ensure_display_row(db, user_id) db.execute("UPDATE display_config SET custom_html=? WHERE user_id=?", (html, user_id)) db.commit() db.close() return True def set_custom_html_template(index: int, html: str, db_path: str = None, user_id: int = 0) -> bool: if index not in (1, 2, 3): return False db = get_db(_db_path(db_path)) _ensure_display_row(db, user_id) db.execute( f"UPDATE display_config SET custom_html_{index}=? WHERE user_id=?", (html, user_id) ) db.commit() db.close() return True def apply_custom_html(index: int, db_path: str = None, user_id: int = 0) -> bool: """Copy template N (0=main custom_html, 1-3=custom_html_1-3) into custom_html and set mode to custom.""" db = get_db(_db_path(db_path)) row = db.execute( "SELECT custom_html, custom_html_1, custom_html_2, custom_html_3 " "FROM display_config WHERE user_id=?", (user_id,) ).fetchone() if row is None: db.close() return False if index == 0: html = row["custom_html"] or "" elif index == 1: html = row["custom_html_1"] or "" elif index == 2: html = row["custom_html_2"] or "" elif index == 3: html = row["custom_html_3"] or "" else: db.close() return False db.execute("UPDATE display_config SET mode='custom', custom_html=? WHERE user_id=?", (html, user_id)) db.commit() db.close() return True # ─── Auth tokens (dual-token auth) ─────────── # Access tokens are short-lived and validated on every request; refresh tokens # are long-lived and stored in an httpOnly cookie. Only SHA-256 hashes are # stored in the DB. Kicking a user/device = delete their tokens + bump the # user's session_version (which invalidates every outstanding access token). ACCESS_TOKEN_TTL = 15 * 60 # 15 minutes REFRESH_TOKEN_TTL = 30 * 24 * 3600 # 30 days (sliding, renewed on each refresh) def _hash_token(raw: str) -> str: import hashlib return hashlib.sha256(raw.encode("utf-8")).hexdigest() def _prune_expired_tokens(conn): conn.execute("DELETE FROM auth_tokens WHERE expires_at < ?", (int(time.time()),)) def create_auth_token(user_id: int, kind: str, view: str = "display", temp: bool = False, ttl: int = None, db_path: str = None) -> str: """Create an auth token row and return the raw (unhashed) token. Only the SHA-256 hash is persisted.""" import secrets as _secrets if ttl is None: ttl = ACCESS_TOKEN_TTL if kind == "access" else REFRESH_TOKEN_TTL raw = _secrets.token_urlsafe(48 if kind == "refresh" else 32) conn = get_db(_db_path(db_path)) try: _prune_expired_tokens(conn) ver_row = conn.execute( "SELECT session_version FROM users WHERE id=?", (user_id,) ).fetchone() ver = ver_row["session_version"] if ver_row else 0 conn.execute( "INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, " "session_version, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)", (user_id, _hash_token(raw), kind, view, 1 if temp else 0, ver, int(time.time()) + ttl) ) conn.commit() return raw finally: conn.close() def _token_lookup(raw: str, kind: str, db_path: str = None) -> dict | None: """Look up a token by its raw value. For access tokens the row is only valid while it matches the user's current session_version (kick support). Returns None when invalid/expired/user inactive/version bumped.""" if not raw: return None conn = get_db(_db_path(db_path)) try: h = _hash_token(raw) row = conn.execute( "SELECT t.*, u.status AS user_status FROM auth_tokens t " "JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind=?", (h, kind) ).fetchone() if not row: return None if row["expires_at"] <= int(time.time()): return None if row["user_status"] != "active": return None if kind == "access" and row["session_version"] != _cur_session_version(row["user_id"], conn): return None return dict(row) finally: conn.close() def _cur_session_version(user_id: int, conn) -> int: r = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone() return r["session_version"] if r else -1 def find_access_token(raw: str, db_path: str = None) -> dict | None: """Validate an access token. Returns the row dict or None.""" return _token_lookup(raw, "access", db_path) def rotate_refresh_token(raw: str, db_path: str = None) -> dict | None: """Atomically validate and rotate a refresh token (concurrency-safe). Returns a fresh {access_token, refresh_token} pair or None. The old refresh token is consumed inside a single transaction, so a replayed/concurrent duplicate can never mint two valid pairs.""" if not raw: return None import secrets as _secrets conn = get_db(_db_path(db_path)) try: conn.execute("BEGIN IMMEDIATE") row = conn.execute( "SELECT t.*, u.status AS user_status FROM auth_tokens t " "JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind='refresh'", (_hash_token(raw),) ).fetchone() if (not row or row["expires_at"] <= int(time.time()) or row["user_status"] != "active"): conn.rollback() return None cur = conn.execute("DELETE FROM auth_tokens WHERE id=?", (row["id"],)) if cur.rowcount == 0: conn.rollback() return None now = int(time.time()) acc_raw = _secrets.token_urlsafe(32) ref_raw = _secrets.token_urlsafe(48) ver = row["session_version"] conn.execute( "INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, " "session_version, expires_at) VALUES (?, ?, 'access', ?, ?, ?, ?)", (row["user_id"], _hash_token(acc_raw), row["view"], row["temp"], ver, now + ACCESS_TOKEN_TTL) ) conn.execute( "INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, " "session_version, expires_at) VALUES (?, ?, 'refresh', ?, ?, ?, ?)", (row["user_id"], _hash_token(ref_raw), row["view"], row["temp"], ver, now + REFRESH_TOKEN_TTL) ) conn.commit() return {"access_token": acc_raw, "refresh_token": ref_raw} finally: conn.close() def find_refresh_token(raw: str, db_path: str = None) -> dict | None: """Validate a refresh token (status check only; version is not compared here because a refresh token belongs to one device and is deleted on kick, not merely invalidated).""" if not raw: return None conn = get_db(_db_path(db_path)) try: h = _hash_token(raw) row = conn.execute( "SELECT t.*, u.status AS user_status FROM auth_tokens t " "JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind='refresh'", (h,) ).fetchone() if not row: return None if row["expires_at"] <= int(time.time()): return None if row["user_status"] != "active": return None return dict(row) finally: conn.close() def mark_token_reauth(token_row_id: int, db_path: str = None): """Stamp the current access token row as re-authenticated (per-device).""" if not token_row_id: return conn = get_db(_db_path(db_path)) try: conn.execute("UPDATE auth_tokens SET reauth_at=? WHERE id=?", (int(time.time()), token_row_id)) conn.commit() finally: conn.close() def delete_auth_token(raw: str, kind: str = None, db_path: str = None): """Delete a single token row by raw value (logout of one device).""" if not raw: return conn = get_db(_db_path(db_path)) try: h = _hash_token(raw) if kind: conn.execute("DELETE FROM auth_tokens WHERE token_hash=? AND kind=?", (h, kind)) else: conn.execute("DELETE FROM auth_tokens WHERE token_hash=?", (h,)) conn.commit() finally: conn.close() def revoke_user_tokens(user_id: int, db_path: str = None): """Delete every token of a user (kick all devices of one account). Callers should also bump_user_session_version() to invalidate access tokens.""" conn = get_db(_db_path(db_path)) try: conn.execute("DELETE FROM auth_tokens WHERE user_id=?", (user_id,)) conn.commit() finally: conn.close() def revoke_all_tokens(db_path: str = None): """Delete every token of every user (Root: kick all users).""" conn = get_db(_db_path(db_path)) try: conn.execute("DELETE FROM auth_tokens") conn.commit() finally: conn.close() # ─── Auth / MFA helpers (per user) ────────── def get_auth_config(db_path: str = None, user_id: int = 0) -> dict: """Return the full auth_config row for a user, with defaults for missing fields.""" conn = get_db(_db_path(db_path)) row = conn.execute("SELECT * FROM auth_config WHERE user_id = ?", (user_id,)).fetchone() if row is None: conn.execute( """INSERT INTO auth_config (user_id, passkey_credential_id, passkey_public_key, totp_secret, totp_enabled, recovery_codes, temp_password, must_configure_mfa) VALUES (?, '', '', '', 0, '', '', 1)""", (user_id,) ) conn.commit() cfg = { "user_id": user_id, "passkey_credential_id": "", "passkey_public_key": "", "totp_secret": "", "totp_enabled": 0, "recovery_codes": "", "temp_password": "", "must_configure_mfa": 1, } else: cfg = { "user_id": row["user_id"], "passkey_credential_id": row["passkey_credential_id"] or "", "passkey_public_key": row["passkey_public_key"] or "", "totp_secret": row["totp_secret"] or "", "totp_enabled": bool(row["totp_enabled"]), "recovery_codes": row["recovery_codes"] or "", "temp_password": row["temp_password"] or "", "must_configure_mfa": bool(row["must_configure_mfa"]), } conn.close() return cfg def update_auth_config(user_id: int = 0, **kwargs): """Update auth_config fields. Keys match DB column names.""" db_path = _db_path(kwargs.pop("_db_path", None)) conn = get_db(db_path) row = conn.execute("SELECT user_id FROM auth_config WHERE user_id = ?", (user_id,)).fetchone() fields = [] vals = [] for key, val in kwargs.items(): if key in ("passkey_credential_id", "passkey_public_key", "totp_secret", "recovery_codes", "temp_password", "must_configure_mfa", "totp_enabled"): fields.append(f"{key}=?") vals.append(val) if not fields: conn.close() return if row is None: cols = ["user_id"] + [f.split("=")[0] for f in fields] col_str = ", ".join(cols) ph = ", ".join("?" * len(cols)) conn.execute(f"INSERT INTO auth_config ({col_str}) VALUES ({ph})", (user_id, *vals)) else: vals.append(user_id) conn.execute(f"UPDATE auth_config SET {', '.join(fields)} WHERE user_id=?", vals) conn.commit() conn.close() def get_temp_password(db_path: str = None, user_id: int = 0) -> str: cfg = get_auth_config(db_path, user_id) return cfg["temp_password"] def set_temp_password(password: str, db_path: str = None, user_id: int = 0): update_auth_config(user_id, temp_password=password, _db_path=db_path) def clear_temp_password(db_path: str = None, user_id: int = 0): update_auth_config(user_id, temp_password="", must_configure_mfa=0, _db_path=db_path) def get_recovery_codes(db_path: str = None, user_id: int = 0) -> list: cfg = get_auth_config(db_path, user_id) raw = cfg["recovery_codes"] if not raw: return [] return raw.split(",") def set_recovery_codes(codes: list, db_path: str = None, user_id: int = 0): update_auth_config(user_id, recovery_codes=",".join(codes), _db_path=db_path) def has_mfa_configured(db_path: str = None, user_id: int = 0) -> bool: cfg = get_auth_config(db_path, user_id) return bool(cfg["passkey_credential_id"]) or bool(cfg["totp_enabled"]) def find_user_by_credential_id(credential_id: str, db_path: str = None) -> int | None: """Find the user who owns a given WebAuthn credential id (for login).""" conn = get_db(_db_path(db_path)) row = conn.execute( "SELECT a.user_id FROM auth_config a " "JOIN users u ON u.id = a.user_id " "WHERE a.passkey_credential_id = ? AND u.status = 'active'", (credential_id,) ).fetchone() conn.close() return row["user_id"] if row else None # ─── DB init / migration ───────────────────── # Stores passwords that were just generated, for the startup script to print. _EMITTED_PASSWORDS: list = [] # Current schema version supported by this program. Bump it whenever the # SCHEMA changes so existing databases get upgraded on next start. DB_VERSION = 2 def _ensure_defaults(conn): """Insert default rows for config tables and the Root user when missing. Generates a random 6-digit temp password when Root's auth_config is created (or when a migrated DB has no MFA configured yet). """ import secrets import string # Schema-version row (single row, id=1). conn.execute("INSERT OR IGNORE INTO meta (id, version) VALUES (1, ?)", (DB_VERSION,)) # Root user (id 0) root = conn.execute("SELECT id FROM users WHERE id=0").fetchone() if root is None: conn.execute( "INSERT OR IGNORE INTO users (id, remark, role, status) VALUES (0, 'Root', 'admin', 'active')" ) # Keep the users identity sequence ahead of explicit id=0 insertion. conn.execute( "SELECT setval(pg_get_serial_sequence('users', 'id'), " "GREATEST((SELECT COALESCE(MAX(id), 1) FROM users), 1))" ) # weather_config for Root count = conn.execute("SELECT COUNT(*) FROM weather_config WHERE user_id=0").fetchone()[0] if count == 0: conn.execute("INSERT INTO weather_config (user_id, city, lat, lon) " "VALUES (0, 'Beijing', 39.9042, 116.4074)") # clock_config for Root count = conn.execute("SELECT COUNT(*) FROM clock_config WHERE user_id=0").fetchone()[0] if count == 0: conn.execute("INSERT INTO clock_config (user_id, city, timezone, is_primary) " "VALUES (0, 'Beijing', 'Asia/Shanghai', 1)") # auth_config for Root row = conn.execute("SELECT * FROM auth_config WHERE user_id=0").fetchone() mfa_configured = bool(row and (row["passkey_credential_id"] or row["totp_enabled"])) if row is None: temp_pw = "".join(secrets.choice(string.digits) for _ in range(6)) conn.execute( "INSERT INTO auth_config (user_id, passkey_credential_id, passkey_public_key, " "totp_secret, totp_enabled, recovery_codes, temp_password, must_configure_mfa) " "VALUES (0, '', '', '', 0, '', ?, 1)", (temp_pw,) ) _EMITTED_PASSWORDS.append(temp_pw) print(f"\n{'='*50}", file=sys.stderr) print(f" 🔐 Root 初始临时密码: {temp_pw}", file=sys.stderr) print(f" ⚠️ 请使用此密码登录 Root 并配置双因素认证", file=sys.stderr) print(f"{'='*50}\n", file=sys.stderr) else: # Migrated DB. If Root already has MFA, no temp password is needed. if mfa_configured: if row["temp_password"] or row["must_configure_mfa"]: conn.execute( "UPDATE auth_config SET temp_password='', must_configure_mfa=0 WHERE user_id=0" ) elif (not row["temp_password"]) and row["must_configure_mfa"]: # Migrated DB without MFA — generate a fresh temp password. temp_pw = "".join(secrets.choice(string.digits) for _ in range(6)) conn.execute( "UPDATE auth_config SET temp_password=? WHERE user_id=0", (temp_pw,) ) _EMITTED_PASSWORDS.append(temp_pw) print(f"\n{'='*50}", file=sys.stderr) print(f" 🔐 迁移已存在数据库 — Root 临时密码: {temp_pw}", file=sys.stderr) print(f" ⚠️ 请使用此密码登录 Root 并配置双因素认证", file=sys.stderr) print(f"{'='*50}\n", file=sys.stderr) conn.commit() def _table_columns(conn, table_name: str) -> list: """Return the column names of a table (information_schema).""" rows = conn.execute( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = current_schema() AND table_name = ? " "ORDER BY ordinal_position", (table_name,) ).fetchall() return [r["column_name"] for r in rows] def _read_db_version(conn) -> int | None: """Read the schema version from the `meta` table. Returns None when the DB is legacy (no meta table / no version row). """ try: row = conn.execute("SELECT version FROM meta WHERE id=1").fetchone() return int(row["version"]) if row else None except Exception: conn._safe_rollback() return None # ─── SCHEMA parsing for column-level migration ────── _CREATE_TABLE_RE = re.compile( r"CREATE TABLE IF NOT EXISTS\s+(\w+)\s*\((.*?)\)\s*;", re.IGNORECASE | re.DOTALL, ) def _split_top_level(text: str) -> list: """Split `text` on top-level commas (ignoring commas inside parentheses AND inside single-quoted strings, e.g. DEFAULT '1,2,3,4,5,6,7').""" parts = [] depth = 0 cur = [] in_str = False i = 0 while i < len(text): ch = text[i] if ch == "'": # SQL 单引号转义用两个单引号表示 if in_str and i + 1 < len(text) and text[i + 1] == "'": cur.append("''") i += 2 continue in_str = not in_str cur.append(ch) elif not in_str: if ch == "(": depth += 1 elif ch == ")": depth -= 1 if ch == "," and depth == 0: parts.append("".join(cur)) cur = [] i += 1 continue cur.append(ch) else: cur.append(ch) i += 1 if cur: parts.append("".join(cur)) return parts def _strip_check(s: str) -> str: """Remove every CHECK(...) clause (parenthesis-balanced).""" out = s while True: idx = out.upper().find("CHECK(") if idx == -1: break depth = 0 end = -1 for j in range(idx, len(out)): if out[j] == "(": depth += 1 elif out[j] == ")": depth -= 1 if depth == 0: end = j break if end == -1: break out = out[:idx] + out[end + 1:] return out def _migrate_missing_columns(conn): """Add columns that exist in SCHEMA but are missing from existing tables. New columns are created with their SCHEMA DEFAULT (migration contract). Best-effort: a failing column does not block startup. """ for m in _CREATE_TABLE_RE.finditer(SCHEMA): table = m.group(1) try: existing = set(_table_columns(conn, table)) except Exception: conn._safe_rollback() continue for col_def in _split_top_level(m.group(2)): col_def = col_def.strip() if not col_def: continue upper = col_def.upper() if upper.startswith(("PRIMARY KEY", "UNIQUE", "CHECK", "FOREIGN KEY", "CONSTRAINT")): continue name = col_def.split()[0].strip('"') if not name or name in existing: continue # Make the ADD COLUMN safe for rows that already exist: # drop NOT NULL / PK / UNIQUE / IDENTITY; keep TYPE + DEFAULT. safe = re.sub(r"\s+NOT\s+NULL\b", "", col_def, flags=re.IGNORECASE) safe = re.sub(r"\s+PRIMARY\s+KEY\b", "", safe, flags=re.IGNORECASE) safe = re.sub(r"\s+UNIQUE\b", "", safe, flags=re.IGNORECASE) safe = re.sub(r"\s+GENERATED\s+BY\s+DEFAULT\s+AS\s+IDENTITY\b", "", safe, flags=re.IGNORECASE) safe = _strip_check(safe).strip().rstrip(",") try: conn.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {safe}") except Exception: # rollback done by the shim; the next column starts a fresh tx pass def _fix_identity_sequences(conn): """Sync every IDENTITY sequence to max(id) (safe after explicit id inserts).""" try: rows = conn.execute( "SELECT table_name, column_name FROM information_schema.columns " "WHERE table_schema = current_schema() AND is_identity = 'YES'" ).fetchall() except Exception: conn._safe_rollback() return for r in rows: t, c = r["table_name"], r["column_name"] try: conn.execute( f"SELECT setval(pg_get_serial_sequence('{t}', '{c}'), " f"GREATEST((SELECT COALESCE(MAX({c}), 1) FROM {t}), 1))" ) except Exception: conn._safe_rollback() # Timestamp columns whose DEFAULT was generated by early builds with the wrong # to_char pattern ('YYYY-MM-DD HH:MM:SS' — HH is 12-hour, MM is the month). _TIMESTAMP_DEFAULT_FIX = [ ("auth_tokens", "created_at"), ("users", "created_at"), ("calendar_groups", "created_at"), ("weather_cache", "updated_at"), ("weather_hourly_cache", "fetched_at"), ] def _fix_timestamp_defaults(conn): """Repair legacy DEFAULT expressions (idempotent, best-effort).""" for t, c in _TIMESTAMP_DEFAULT_FIX: try: conn.execute( f"ALTER TABLE {t} ALTER COLUMN {c} " "SET DEFAULT to_char(now(), 'YYYY-MM-DD HH24:MI:SS')" ) except Exception: conn._safe_rollback() def init_db(db_path: str = "data.db"): """Initialize the database with schema-version aware migration (PostgreSQL). - Fresh database (no `meta` row) -> create the latest SCHEMA + defaults. - Existing database: * DB version == DB_VERSION -> use as-is (fast path, defaults ensured) * no version / DB version < DB_VERSION -> create missing tables/columns; missing columns are filled with their SCHEMA DEFAULT; stamp the new version * DB version > DB_VERSION -> abort with an error (DB is newer than this program) """ conn = get_db(db_path) try: old_version = _read_db_version(conn) if old_version is not None and old_version > DB_VERSION: raise SystemExit( f"[FATAL] 数据库 schema 版本 (v{old_version}) 高于当前程序支持的版本 " f"(v{DB_VERSION})。请先升级程序,或删除/降级数据库后再启动。" ) if old_version is not None and old_version == DB_VERSION: # Same version — no upgrade needed. Just ensure defaults exist. _ensure_defaults(conn) _fix_timestamp_defaults(conn) conn.commit() return # Create any missing tables (CREATE TABLE IF NOT EXISTS is idempotent), # then add any columns missing from older databases. The DDL is # committed first so a later best-effort column migration (which may # roll back its own transaction) can never undo the schema creation. conn.executescript(SCHEMA) conn.commit() _migrate_missing_columns(conn) # Bootstrap defaults (Root user, meta row, config rows). _ensure_defaults(conn) # Repair legacy timestamp DEFAULT expressions (old to_char HH:MM pattern). _fix_timestamp_defaults(conn) # Stamp the database with the program's schema version. conn.execute("UPDATE meta SET version=? WHERE id=1", (DB_VERSION,)) conn.execute("INSERT OR IGNORE INTO meta (id, version) VALUES (1, ?)", (DB_VERSION,)) _fix_identity_sequences(conn) conn.commit() finally: conn.close()