gaokao-volunteer-system/admin/db.py

"""SQLite 连接管理 (T6.1).

提供:
- get_connection(db_path) : 打开 SQLite（启用外键 + WAL）
- ensure_schema(db_path)  : 幂等创建 admin_users 表
- bootstrap_admin(...)    : 启动时若无账户则自动创建默认管理员
- AdminUser / AdminUserRepo : dataclass + 简单 DAO
"""

from __future__ import annotations

import sqlite3
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from admin.config import Settings
from admin.password import hash_password, verify_password


def utc_now_iso() -> str:
    """当前 UTC 时间（ISO8601 秒精度）。"""
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def get_connection(db_path: str) -> sqlite3.Connection:
    """打开 SQLite 连接。

    启用外键约束；row_factory = Row（便于列名访问）。
    默认路径会创建父目录。
    """
    path = Path(db_path)
    if path != Path(":memory:") and str(db_path) != ":memory:":
        path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path, isolation_level=None)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    # WAL 仅对真实文件有效（:memory: 会忽略）
    if str(db_path) != ":memory:":
        conn.execute("PRAGMA journal_mode = WAL")
    return conn


ADMIN_USERS_DDL = """
CREATE TABLE IF NOT EXISTS admin_users (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    username        TEXT UNIQUE NOT NULL,
    password_hash   TEXT NOT NULL,
    role            TEXT NOT NULL DEFAULT 'admin',
    is_active       INTEGER NOT NULL DEFAULT 1,
    created_at      TEXT NOT NULL,
    last_login_at   TEXT
);
CREATE INDEX IF NOT EXISTS idx_admin_users_username ON admin_users(username);
"""


def ensure_schema(db_path: str) -> None:
    """幂等创建 admin_users 表。"""
    with get_connection(db_path) as conn:
        conn.executescript(ADMIN_USERS_DDL)


@dataclass(frozen=True)
class AdminUser:
    """管理员用户（不可变视图）。"""

    id: int
    username: str
    role: str
    is_active: bool
    created_at: str
    last_login_at: Optional[str] = None

    @classmethod
    def from_row(cls, row: sqlite3.Row) -> "AdminUser":
        return cls(
            id=row["id"],
            username=row["username"],
            role=row["role"],
            is_active=bool(row["is_active"]),
            created_at=row["created_at"],
            last_login_at=row["last_login_at"],
        )

    def to_public_dict(self) -> dict:
        """公开字段（不含密码）。"""
        return {
            "id": self.id,
            "username": self.username,
            "role": self.role,
            "is_active": self.is_active,
            "created_at": self.created_at,
            "last_login_at": self.last_login_at,
        }


class AdminUserRepo:
    """AdminUser 仓库（基于 sqlite3）。"""

    def __init__(self, db_path: str):
        self.db_path = db_path

    def _conn(self) -> sqlite3.Connection:
        return get_connection(self.db_path)

    def get_by_username(self, username: str) -> Optional[tuple[AdminUser, str]]:
        """按用户名查找，返回 (AdminUser, password_hash) 或 None。"""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT id, username, password_hash, role, is_active, "
                "created_at, last_login_at FROM admin_users WHERE username = ?",
                (username,),
            ).fetchone()
        if row is None:
            return None
        return AdminUser.from_row(row), row["password_hash"]

    def get_by_id(self, user_id: int) -> Optional[AdminUser]:
        """按 ID 查找（不含密码）。"""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT id, username, role, is_active, created_at, last_login_at "
                "FROM admin_users WHERE id = ?",
                (user_id,),
            ).fetchone()
        return AdminUser.from_row(row) if row else None

    def create(self, username: str, password: str, role: str = "admin") -> AdminUser:
        """创建管理员。重复 username 抛 sqlite3.IntegrityError。"""
        now = utc_now_iso()
        password_hash = hash_password(password)
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO admin_users(username, password_hash, role, "
                "is_active, created_at) VALUES(?, ?, ?, 1, ?)",
                (username, password_hash, role, now),
            )
            user_id = cur.lastrowid or 0
        if user_id <= 0:  # pragma: no cover - 数据库异常兜底
            raise RuntimeError("failed to obtain lastrowid for new admin user")
        return AdminUser(
            id=user_id,
            username=username,
            role=role,
            is_active=True,
            created_at=now,
        )

    def count(self) -> int:
        """当前账户总数。"""
        with self._conn() as conn:
            row = conn.execute("SELECT COUNT(*) AS n FROM admin_users").fetchone()
        return int(row["n"])

    def update_last_login(self, user_id: int) -> None:
        """更新最后登录时间。"""
        with self._conn() as conn:
            conn.execute(
                "UPDATE admin_users SET last_login_at = ? WHERE id = ?",
                (utc_now_iso(), user_id),
            )


def bootstrap_admin(settings: Settings) -> tuple[bool, str]:
    """若 admin_users 为空,自动创建默认管理员。

    Returns:
        (created, message): created=True 表示已创建, message 为提示文本
    """
    ensure_schema(settings.db_path)
    repo = AdminUserRepo(settings.db_path)
    if repo.count() > 0:
        return False, "已存在管理员账户,跳过 bootstrap"
    user = repo.create(
        username=settings.default_admin_username,
        password=settings.default_admin_password,
    )
    return True, (
        f"已创建默认管理员: username={user.username} id={user.id} "
        f"(请尽快通过 GAOKAO_ADMIN_PASS 覆盖默认密码!)"
    )


def authenticate(
    repo: AdminUserRepo, username: str, password: str
) -> Optional[AdminUser]:
    """校验用户名+密码,成功返回 AdminUser 并更新 last_login_at。"""
    result = repo.get_by_username(username)
    if result is None:
        return None
    user, password_hash = result
    if not user.is_active:
        return None
    if not verify_password(password, password_hash):
        return None
    repo.update_last_login(user.id)
    return user