HeXi/hexi/web_auth.py

# -*- coding: utf-8 -*-
"""统一 Web 鉴权模块：OAuth2 Password + SQLite 存储。

- 用户/访问令牌存 SQLite：hexi/data/web_auth.db
- 提供 OAuth2 密码流（tokenUrl 指向 /hub/api/auth/token）
- 供统一 Web /hub 与各插件 API 共用（/api/<plugin>/...）
"""

from __future__ import annotations

import hashlib
import os
import secrets
import sqlite3
import time
from pathlib import Path
from typing import Optional

from fastapi import Depends, HTTPException
from fastapi.security import OAuth2PasswordBearer

basic_path = Path(__file__).resolve().parent
DATA_DIR = basic_path / "data"
DB_PATH = DATA_DIR / "web_auth.db"

TOKEN_TTL_SECONDS = 60 * 60 * 24  # 24 小时
PBKDF2_ITERATIONS = 100_000

oauth2_scheme = OAuth2PasswordBearer(
    tokenUrl="/hub/api/auth/token",
    auto_error=False,
)


def _conn() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def _now() -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S")


def _read_env(key: str, default: str) -> str:
    """取值优先级：os.environ > .env 文件 > default（NoneBot 不一定把 .env 全塞进 env）。"""
    val = os.getenv(key)
    if val:
        return val
    env_path = Path(__file__).resolve().parent.parent / ".env"
    if env_path.exists():
        for line in env_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line.startswith(key + "=") or line.startswith(key + " ="):
                value = line.split("=", 1)[1].strip()
                return value.strip('"').strip("'")
    return default


def _hash_password(password: str) -> str:
    salt = secrets.token_hex(16)
    digest = hashlib.pbkdf2_hmac(
        "sha256", password.encode("utf-8"), salt.encode("utf-8"), PBKDF2_ITERATIONS
    ).hex()
    return f"{salt}${digest}"


def _verify_password(password: str, stored: str) -> bool:
    try:
        salt, digest = stored.split("$", 1)
    except ValueError:
        return False
    calc = hashlib.pbkdf2_hmac(
        "sha256", password.encode("utf-8"), salt.encode("utf-8"), PBKDF2_ITERATIONS
    ).hex()
    return secrets.compare_digest(digest, calc)


def init_auth_db() -> None:
    """只建表（轻量，可每次调用）；用户播种/同步由 sync_admin() 在启动时做一次。"""
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    with _conn() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS web_users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                is_active INTEGER NOT NULL DEFAULT 1,
                created_at TEXT
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS web_tokens (
                token TEXT PRIMARY KEY,
                user_id INTEGER NOT NULL,
                expires_at REAL NOT NULL,
                created_at TEXT
            )
            """
        )
        conn.commit()


def sync_admin() -> None:
    """启动时调用：以 .env 为准创建/同步超级管理员（含改密码场景）。"""
    init_auth_db()
    username = _read_env("hexi_web_username", "admin")
    password = _read_env("hexi_web_password", "admin")
    with _conn() as conn:
        row = conn.execute(
            "SELECT id FROM web_users WHERE username=?", (username,)
        ).fetchone()
        if row is None:
            conn.execute(
                "INSERT INTO web_users(username,password_hash,is_active,created_at) VALUES(?,?,1,?)",
                (username, _hash_password(password), _now()),
            )
        else:
            # 以 .env 为准同步超级管理员密码：解决改 .env 后库里还是旧密码的问题
            conn.execute(
                "UPDATE web_users SET password_hash=?, is_active=1 WHERE id=?",
                (_hash_password(password), row["id"]),
            )
        conn.commit()


def create_user(username: str, password: str) -> int:
    init_auth_db()
    with _conn() as conn:
        cursor = conn.execute(
            "INSERT INTO web_users(username,password_hash,is_active,created_at) VALUES(?,?,1,?)",
            (username, _hash_password(password), _now()),
        )
        conn.commit()
        return int(cursor.lastrowid)


def revoke_user_tokens(user_id: int, except_token: Optional[str] = None) -> None:
    """吊销某用户访问令牌（可保留当前令牌）。改密码后调用，旧登录态全部失效。"""
    init_auth_db()
    with _conn() as conn:
        if except_token:
            conn.execute(
                "DELETE FROM web_tokens WHERE user_id=? AND token<>?",
                (user_id, except_token),
            )
        else:
            conn.execute("DELETE FROM web_tokens WHERE user_id=?", (user_id,))
        conn.commit()


def change_password(
    user_id: int, old_password: str, new_password: str, keep_token: Optional[str] = None
) -> bool:
    """校验旧密码后改新密码，并吊销旧令牌（保留当前令牌可选）。成功返回 True。"""
    if not new_password:
        return False
    init_auth_db()
    with _conn() as conn:
        row = conn.execute(
            "SELECT * FROM web_users WHERE id=? AND is_active=1", (user_id,)
        ).fetchone()
        if not row or not _verify_password(old_password, row["password_hash"]):
            return False
        conn.execute(
            "UPDATE web_users SET password_hash=? WHERE id=?",
            (_hash_password(new_password), user_id),
        )
        conn.commit()
    revoke_user_tokens(user_id, except_token=keep_token)
    return True


def authenticate(username: str, password: str) -> Optional[int]:
    """校验成功返回 user_id，否则 None。"""
    init_auth_db()
    with _conn() as conn:
        row = conn.execute(
            "SELECT * FROM web_users WHERE username=? AND is_active=1", (username,)
        ).fetchone()
        if row and _verify_password(password, row["password_hash"]):
            return int(row["id"])
    return None


def issue_token(user_id: int) -> tuple[str, int]:
    """签发不透明访问令牌（存库可吊销），返回 (token, expires_in)。"""
    init_auth_db()
    token = secrets.token_urlsafe(32)
    expires_at = time.time() + TOKEN_TTL_SECONDS
    with _conn() as conn:
        conn.execute(
            "INSERT INTO web_tokens(token,user_id,expires_at,created_at) VALUES(?,?,?,?)",
            (token, user_id, expires_at, _now()),
        )
        conn.commit()
    return token, TOKEN_TTL_SECONDS


def revoke_token(token: str) -> None:
    init_auth_db()
    with _conn() as conn:
        conn.execute("DELETE FROM web_tokens WHERE token=?", (token,))
        conn.commit()


def get_user_by_token(token: str) -> Optional[dict]:
    init_auth_db()
    with _conn() as conn:
        row = conn.execute(
            """
            SELECT u.id, u.username, u.created_at
            FROM web_tokens t JOIN web_users u ON u.id = t.user_id
            WHERE t.token=? AND t.expires_at > ?
            """,
            (token, time.time()),
        ).fetchone()
        return dict(row) if row else None


async def get_current_user(token: Optional[str] = Depends(oauth2_scheme)) -> dict:
    if not token:
        raise HTTPException(status_code=401, detail="未登录")
    user = get_user_by_token(token)
    if not user:
        raise HTTPException(status_code=401, detail="登录态无效或已过期")
    return user


# 各插件 API 用的统一鉴权依赖
require_admin = Depends(get_current_user)