PlantAI-rag-new/sdk/app/services/auth.py

"""services/auth.py —— 管理员鉴权:密码哈希、会话凭据、用户 CRUD、种子管理员。

密码用 stdlib pbkdf2_hmac + 每用户随机盐(不新增依赖);
会话凭据用不透明 token 存 auth_tokens,逐次查库(吊销/过期即时生效)。
所有 DB 访问经 `db` 参数注入(与 services/history.py 一致,避免依赖全局单例)。
"""
from __future__ import annotations

import hashlib
import logging
import secrets
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any

from ..config import get_settings
from ..db import Database, dumps_json, loads_json
from ..logwriter import log_event

logger = logging.getLogger(__name__)

# 授权清单功能单元(9 项,与管理后台界面一致;后续可调)
AUTHORIZATION_UNITS = (
    "qa", "kg", "extraction", "docgen", "figures",
    "private_docs", "bioflow", "paper", "legacy",
)

_PBKDF2_ITERATIONS = 200_000


def _now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def hash_password(password: str, salt: str | None = None) -> tuple[str, str]:
    salt = salt or secrets.token_hex(16)
    digest = hashlib.pbkdf2_hmac(
        "sha256", password.encode("utf-8"), salt.encode("utf-8"), _PBKDF2_ITERATIONS
    )
    return digest.hex(), salt


def verify_password(password: str, salt: str, expected_hex: str) -> bool:
    digest, _ = hash_password(password, salt)
    return secrets.compare_digest(digest, expected_hex)


def _row_to_user(row: Any) -> dict[str, Any]:
    return {
        "id": row["id"],
        "username": row["username"],
        "role": row["role"],
        "authorizations": loads_json(row["authorizations"], default=[]),
        "created_at": row["created_at"],
    }


# ---- 登录校验 ----

async def verify_login(db: Database, username: str, password: str) -> dict[str, Any] | None:
    row = await db.fetchone(
        "SELECT id, username, password_hash, salt, role, authorizations, created_at "
        "FROM users WHERE username = ?",
        (username,),
    )
    if row is None:
        return None
    if not verify_password(password, row["salt"], row["password_hash"]):
        return None
    return _row_to_user(row)


# ---- 会话凭据 ----

async def issue_token(db: Database, user_id: str, ttl_seconds: int | None = None) -> str:
    settings = get_settings()
    ttl = ttl_seconds if ttl_seconds is not None else settings.admin_token_ttl_seconds
    token = secrets.token_urlsafe(32)
    now = _now_iso()
    expires = datetime.now(timezone.utc) + timedelta(seconds=ttl)
    await db.execute(
        "INSERT INTO auth_tokens (token, user_id, expires_at, created_at, revoked) "
        "VALUES (?, ?, ?, ?, 0)",
        (token, user_id, expires.strftime("%Y-%m-%dT%H:%M:%SZ"), now),
    )
    return token


async def resolve_token(db: Database, token: str) -> dict[str, Any] | None:
    row = await db.fetchone(
        "SELECT t.expires_at, t.revoked, u.id, u.username, u.role, u.authorizations, u.created_at "
        "FROM auth_tokens t JOIN users u ON u.id = t.user_id WHERE t.token = ?",
        (token,),
    )
    if row is None or row["revoked"]:
        return None
    try:
        expires = datetime.strptime(row["expires_at"], "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
    except ValueError:
        return None
    if datetime.now(timezone.utc) >= expires:
        return None
    return _row_to_user(row)


async def revoke_token(db: Database, token: str) -> None:
    await db.execute("UPDATE auth_tokens SET revoked = 1 WHERE token = ?", (token,))


# ---- 用户 CRUD ----

async def list_users(db: Database) -> list[dict[str, Any]]:
    rows = await db.fetchall(
        "SELECT id, username, role, authorizations, created_at FROM users ORDER BY created_at ASC"
    )
    return [_row_to_user(r) for r in rows]


async def get_user(db: Database, user_id: str) -> dict[str, Any] | None:
    row = await db.fetchone(
        "SELECT id, username, role, authorizations, created_at FROM users WHERE id = ?", (user_id,)
    )
    return _row_to_user(row) if row else None


async def get_user_by_username(db: Database, username: str) -> dict[str, Any] | None:
    row = await db.fetchone(
        "SELECT id, username, role, authorizations, created_at FROM users WHERE username = ?", (username,)
    )
    return _row_to_user(row) if row else None


async def create_user(
    db: Database, username: str, password: str, role: str, authorizations: list[str]
) -> dict[str, Any]:
    user_id = str(uuid.uuid4())
    ph, salt = hash_password(password)
    await db.execute(
        "INSERT INTO users (id, username, password_hash, salt, role, authorizations, created_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?)",
        (user_id, username, ph, salt, role, dumps_json(authorizations), _now_iso()),
    )
    log_event("user_created", detail={"target": username, "role": role})
    return {"id": user_id, "username": username, "role": role, "authorizations": authorizations, "created_at": _now_iso()}


async def update_user(
    db: Database,
    user_id: str,
    *,
    role: str | None = None,
    authorizations: list[str] | None = None,
    password: str | None = None,
) -> dict[str, Any] | None:
    fields: list[str] = []
    params: list[Any] = []
    if role is not None:
        fields.append("role = ?")
        params.append(role)
    if authorizations is not None:
        fields.append("authorizations = ?")
        params.append(dumps_json(authorizations))
    if password is not None:
        ph, salt = hash_password(password)
        fields.append("password_hash = ?")
        params.append(ph)
        fields.append("salt = ?")
        params.append(salt)
    if fields:
        params.append(user_id)
        await db.execute(f"UPDATE users SET {', '.join(fields)} WHERE id = ?", params)
        log_event("user_updated", detail={"target_id": user_id, "role": role, "password_reset": password is not None})
    return await get_user(db, user_id)


async def delete_user(db: Database, user_id: str) -> None:
    await db.execute("DELETE FROM users WHERE id = ?", (user_id,))
    log_event("user_deleted", detail={"target_id": user_id})


# ---- 种子管理员 ----

async def seed_admin(db: Database) -> None:
    """启动时确保存在至少一个 admin 用户(默认 admin/admin)。"""
    row = await db.fetchone("SELECT id FROM users WHERE role = 'admin' LIMIT 1")
    if row is not None:
        return
    settings = get_settings()
    await create_user(db, settings.admin_seed_username, settings.admin_seed_password, "admin", list(AUTHORIZATION_UNITS))
    logger.info("已创建种子管理员 %s", settings.admin_seed_username)