""" Database Module for User Authentication ======================================== SQLite-based user storage for Pharma K platform. """ import hashlib import logging import os import secrets import sqlite3 from datetime import datetime from typing import Optional, List, Dict, Any from pathlib import Path logger = logging.getLogger(__name__) #: 模块级标志:管理员未初始化的告警只打印一次,避免每次建连/重渲染刷屏。 _ADMIN_INIT_WARNED = False def _salted_hash(password: str) -> str: """生成加盐密码哈希 ``"$"``(需求 5.4)。 优先复用底座 :mod:`services.auth_service` 的实现,确保全平台哈希格式一致; 在 ``platform/`` 目录未加入导入路径的运行环境下,退回到与之等价的最小 标准库实现(PBKDF2-HMAC-SHA256),保证产出的哈希仍是加盐的、且与 AuthService 的校验逻辑兼容(相同算法 / 迭代轮数 / 盐长度)。 """ try: from services.auth_service import hash_password # type: ignore return hash_password(password) except Exception: # pragma: no cover - 仅在 platform 目录未入路径时触发 salt = secrets.token_bytes(16) dk = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, 200_000) return f"{salt.hex()}${dk.hex()}" # Database file location. # 数据目录可经环境变量 ``PHARMAK_DATA_DIR`` 覆盖(HF Spaces 等部署环境用), # 未设置时回退到仓库内 ``data/`` 目录(本地开发)。 _DEFAULT_DB_DIR = Path(__file__).parent.parent / "data" DB_DIR = Path(os.environ.get("PHARMAK_DATA_DIR") or _DEFAULT_DB_DIR) DB_PATH = DB_DIR / "pharma_k.db" #: 进程级标志:Dataset 拉取每个进程仅执行一次(启动时),避免每次建连重复下载。 _DB_PULLED = False def _maybe_pull_once() -> None: """首次建连时从私有 HF Dataset 拉取数据库(若启用同步)。 HF 免费档文件系统临时、重启即清空,故进程启动后先尝试从 Dataset 恢复最新 DB; 未启用同步(本地开发 / 未配置环境变量)时静默 no-op。 """ global _DB_PULLED if _DB_PULLED: return _DB_PULLED = True # 无论成败都只尝试一次,避免反复触网 try: from utils import hf_storage if hf_storage.is_enabled(): hf_storage.pull_db(DB_PATH) except Exception as exc: # noqa: BLE001 - 同步不可用绝不影响本地存储 logger.warning("启动期数据库同步拉取不可用:%s", exc) def _push_db() -> None: """把本地数据库推回私有 HF Dataset(若启用同步)。写操作后调用。""" try: from utils import hf_storage if hf_storage.is_enabled(): hf_storage.push_db(DB_PATH) except Exception as exc: # noqa: BLE001 logger.warning("数据库同步推送不可用:%s", exc) def get_db_connection() -> sqlite3.Connection: """Get database connection, creating tables if needed.""" # Ensure data directory exists DB_DIR.mkdir(parents=True, exist_ok=True) # 启动期:先从 Dataset 拉取最新 DB(仅一次),再建连建表。 _maybe_pull_once() conn = sqlite3.connect(str(DB_PATH)) conn.row_factory = sqlite3.Row # Create tables if not exist _init_tables(conn) return conn def _init_tables(conn: sqlite3.Connection): """Initialize database tables.""" cursor = conn.cursor() # Users table cursor.execute(''' CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, role TEXT DEFAULT 'user', created_at DATETIME DEFAULT CURRENT_TIMESTAMP, last_login DATETIME ) ''') # Admin config table cursor.execute(''' CREATE TABLE IF NOT EXISTS admin_config ( key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ) ''') # Per-user LLM config table(用户自配模型,密钥加密存储)。 # api_key_enc 存 Fernet 密文(见 utils.crypto);解密失败 / 未加密时按无配置处理。 cursor.execute(''' CREATE TABLE IF NOT EXISTS user_llm_config ( email TEXT PRIMARY KEY, provider TEXT NOT NULL DEFAULT '', model TEXT NOT NULL DEFAULT '', api_key_enc TEXT NOT NULL DEFAULT '', base_url TEXT NOT NULL DEFAULT '', enabled INTEGER NOT NULL DEFAULT 1, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ) ''') conn.commit() # Initialize default admin if not exists _init_default_admin(conn) # Initialize default LLM config _init_default_llm_config(conn) def _init_default_admin(conn: sqlite3.Connection): """初始化管理员账户(需求 5.1)。 管理员初始密码从环境变量 ``ADMIN_INIT_PASSWORD`` 读取,并以**加盐哈希**存储; 未设置该环境变量时**不创建**任何默认账户,仅记录「管理员未初始化」日志 (不含任何明文密码或哈希)。代码与日志中均不出现明文密码。 """ cursor = conn.cursor() cursor.execute("SELECT id FROM users WHERE email = 'admin'") if cursor.fetchone() is not None: return # 管理员已存在,跳过 password = os.environ.get("ADMIN_INIT_PASSWORD") if not password: global _ADMIN_INIT_WARNED if not _ADMIN_INIT_WARNED: logger.warning( "管理员未初始化:未设置 ADMIN_INIT_PASSWORD 环境变量,未创建默认管理员账户。" "如需管理员登录,请设置该环境变量后重启。" ) _ADMIN_INIT_WARNED = True return password_hash = _salted_hash(password) cursor.execute(''' INSERT INTO users (email, password_hash, role) VALUES (?, ?, 'admin') ''', ('admin', password_hash)) conn.commit() logger.info("管理员账户已初始化(初始密码取自环境变量,未记录明文)。") def _init_default_llm_config(conn: sqlite3.Connection): """Initialize default LLM configuration (admin-managed model config).""" cursor = conn.cursor() # 管理员级模型配置默认占位(管理员在后台设置后生效)。 defaults = { "admin_llm_enabled": "0", # 是否启用管理员配置的模型("1"/"0") "admin_llm_provider": "", # 提供商 id(如 deepseek / custom) "admin_llm_model": "", # 模型名(留空则用该提供商默认模型) "admin_llm_api_key": "", # 管理员配置的密钥(存于本地 DB,不入公开库) "admin_llm_base_url": "", # 仅 custom 提供商需要 } for key, value in defaults.items(): cursor.execute( "INSERT OR IGNORE INTO admin_config (key, value) VALUES (?, ?)", (key, value), ) conn.commit() # ============================================================================= # User Operations # ============================================================================= def create_user(email: str, password_hash: str) -> bool: """Create a new user account.""" try: conn = get_db_connection() cursor = conn.cursor() cursor.execute(''' INSERT INTO users (email, password_hash, role) VALUES (?, ?, 'user') ''', (email, password_hash)) conn.commit() conn.close() _push_db() # 注册成功后同步回 Dataset(启用同步时),确保账号持久 return True except sqlite3.IntegrityError: return False # Email already exists def get_user_by_email(email: str) -> Optional[Dict[str, Any]]: """Get user by email.""" conn = get_db_connection() cursor = conn.cursor() cursor.execute("SELECT * FROM users WHERE email = ?", (email,)) row = cursor.fetchone() conn.close() if row: return dict(row) return None def update_last_login(email: str): """Update user's last login time.""" conn = get_db_connection() cursor = conn.cursor() cursor.execute(''' UPDATE users SET last_login = ? WHERE email = ? ''', (datetime.now().isoformat(), email)) conn.commit() conn.close() def get_all_users() -> List[Dict[str, Any]]: """Get all users (for admin).""" conn = get_db_connection() cursor = conn.cursor() cursor.execute("SELECT id, email, role, created_at, last_login FROM users ORDER BY created_at DESC") rows = cursor.fetchall() conn.close() return [dict(row) for row in rows] # ============================================================================= # Admin Config Operations # ============================================================================= def get_admin_config(key: str) -> Optional[str]: """Get admin configuration value.""" conn = get_db_connection() cursor = conn.cursor() cursor.execute("SELECT value FROM admin_config WHERE key = ?", (key,)) row = cursor.fetchone() conn.close() if row: return row['value'] return None def set_admin_config(key: str, value: str): """Set admin configuration value.""" conn = get_db_connection() cursor = conn.cursor() cursor.execute(''' INSERT OR REPLACE INTO admin_config (key, value, updated_at) VALUES (?, ?, ?) ''', (key, value, datetime.now().isoformat())) conn.commit() conn.close() _push_db() # 管理员配置变更后同步回 Dataset(启用同步时) def set_admin_configs(items: Dict[str, str]): """批量写入多个管理员配置项,单次事务 + 单次同步推送。 供 :func:`set_default_llm_config` 等一次性写多键的场景使用,避免逐键多次触网。 """ conn = get_db_connection() cursor = conn.cursor() now = datetime.now().isoformat() for key, value in items.items(): cursor.execute(''' INSERT OR REPLACE INTO admin_config (key, value, updated_at) VALUES (?, ?, ?) ''', (key, value, now)) conn.commit() conn.close() _push_db() # 批量配置变更后单次同步回 Dataset(启用同步时) def get_default_llm_config() -> Dict[str, str]: """返回管理员配置的 LLM 模型设置(admin-managed)。 结构:``{enabled, provider, model, api_key, base_url}``。``enabled`` 为布尔, 其余为字符串。未配置时各项为空 / False。该配置在用户未自行配置时作为兜底 (优先级:用户 > 管理员 > 环境变量)。 """ return { "enabled": (get_admin_config("admin_llm_enabled") or "0") == "1", "provider": get_admin_config("admin_llm_provider") or "", "model": get_admin_config("admin_llm_model") or "", "api_key": get_admin_config("admin_llm_api_key") or "", "base_url": get_admin_config("admin_llm_base_url") or "", } def set_default_llm_config( provider: str, api_key: str, *, model: str = "", base_url: str = "", enabled: bool = True, ): """保存管理员级 LLM 模型设置(写入本地 DB,不入公开仓库)。""" set_admin_configs({ "admin_llm_enabled": "1" if enabled else "0", "admin_llm_provider": provider or "", "admin_llm_model": model or "", "admin_llm_api_key": api_key or "", "admin_llm_base_url": base_url or "", }) def clear_default_llm_config(): """清除 / 停用管理员级模型配置。""" set_admin_config("admin_llm_enabled", "0") # ============================================================================= # Per-User LLM Config Operations(用户自配模型,密钥加密持久化) # ============================================================================= def get_user_llm_config(email: str) -> Optional[Dict[str, Any]]: """读取某用户持久化的自配模型设置(密钥已解密)。 返回 ``{provider, model, api_key, base_url, enabled}``;无记录、加密不可用 或密钥解密失败(secret 轮换 / 密文损坏)时返回 ``None``,调用方按"无持久配置" 处理(退回仅会话存储)。 """ if not email: return None from utils import crypto conn = get_db_connection() cursor = conn.cursor() cursor.execute("SELECT * FROM user_llm_config WHERE email = ?", (email,)) row = cursor.fetchone() conn.close() if not row: return None enc = row["api_key_enc"] or "" api_key = crypto.decrypt(enc) if enc else "" # 有密文但解密失败 → 视为失效配置,按无配置处理。 if enc and not api_key: return None return { "provider": row["provider"] or "", "model": row["model"] or "", "api_key": api_key or "", "base_url": row["base_url"] or "", "enabled": bool(row["enabled"]), } def set_user_llm_config( email: str, provider: str, api_key: str, *, model: str = "", base_url: str = "", enabled: bool = True, ) -> bool: """持久化某用户的自配模型设置(密钥**加密后**落库并同步)。 仅在加密可用(已配置 ``PHARMAK_SECRET_KEY`` 且安装 ``cryptography``)时持久化, 避免明文落库。加密不可用时返回 ``False``,调用方应退回仅会话存储。 """ if not email: return False from utils import crypto if not crypto.is_available(): return False enc = crypto.encrypt(api_key) or "" if api_key else "" conn = get_db_connection() cursor = conn.cursor() cursor.execute(''' INSERT OR REPLACE INTO user_llm_config (email, provider, model, api_key_enc, base_url, enabled, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?) ''', ( email, provider or "", model or "", enc, base_url or "", 1 if enabled else 0, datetime.now().isoformat(), )) conn.commit() conn.close() _push_db() # 用户配置变更后同步回 Dataset(启用同步时) return True def clear_user_llm_config(email: str): """删除某用户的持久化自配模型设置。""" if not email: return conn = get_db_connection() cursor = conn.cursor() cursor.execute("DELETE FROM user_llm_config WHERE email = ?", (email,)) conn.commit() conn.close() _push_db()