src/auth.py

"""认证与 Session 管理模块。"""

from __future__ import annotations

import logging
import os
import secrets
import uuid
from typing import Optional

import bcrypt
from fastapi import HTTPException, Request
from itsdangerous import BadSignature, SignatureExpired, TimestampSigner

# 保持与原实现一致的常量和日志名称
logger = logging.getLogger("gateway")


# ── Session 配置 ───────────────────────────────────────────────────────────────
SESSION_COOKIE = "gw_session"
SESSION_MAX_AGE = 86400  # 24 hours

SECRET_KEY = os.environ.get("SESSION_SECRET") or secrets.token_hex(32)
signer = TimestampSigner(SECRET_KEY)

INTERNAL_KEY_SALT = (os.environ.get("INTERNAL_KEY_SALT") or SECRET_KEY).strip()


# ── 用户加载与认证 ────────────────────────────────────────────────────────────
def _load_users() -> dict[str, str]:
    """从 BASIC_AUTH_USERS 加载用户名密码。"""
    raw = os.environ.get("BASIC_AUTH_USERS", "").replace("\\n", "\n")
    users: dict[str, str] = {}
    for line in raw.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if ":" not in line:
            logger.warning("Skipping invalid BASIC_AUTH_USERS line (no colon)")
            continue
        username, password = line.split(":", 1)
        username = username.strip()
        password = password.strip()
        if username and password:
            users[username] = password
    if not users:
        logger.error("No valid users found — authentication will always fail")
    return users


USERS = _load_users()
INTERNAL_KEY_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_DNS, INTERNAL_KEY_SALT)


def _make_internal_api_key(username: str) -> str:
    """基于用户名生成稳定内部 Key（仅服务端使用）。"""
    value = uuid.uuid5(INTERNAL_KEY_NAMESPACE, username)
    return f"sk-{value}"


INTERNAL_KEY_TO_USER = {
    _make_internal_api_key(username): username for username in USERS.keys()
}


def _verify_credentials(username: str, password: str) -> bool:
    """验证用户名密码，支持明文与 bcrypt。"""
    stored = USERS.get(username)
    if stored is None:
        return False
    if stored.startswith("$2"):
        return bcrypt.checkpw(password.encode(), stored.encode())
    return secrets.compare_digest(stored, password)


# ── Session ──────────────────────────────────────────────────────────────────
def _make_session(username: str) -> str:
    """生成签名的 Session token。"""
    return signer.sign(username).decode()


def _verify_session(token: str) -> Optional[str]:
    """验证 Session token，返回用户名或 None。"""
    try:
        username = signer.unsign(token, max_age=SESSION_MAX_AGE).decode()
    except (BadSignature, SignatureExpired):
        return None
    if username not in USERS:
        return None
    return username


def _get_session_user(request: Request) -> Optional[str]:
    """从请求 Cookie 中解析当前用户。"""
    token = request.cookies.get(SESSION_COOKIE)
    return _verify_session(token) if token else None


def _require_user(request: Request) -> str:
    """FastAPI 依赖：要求用户已登录，否则抛出 401。"""
    username = _get_session_user(request)
    if not username:
        raise HTTPException(status_code=401, detail="Unauthorized")
    return username