auth.py

"""
auth.py — OIDC 核心认证
Authorization Code + PKCE → Authentik → 自动注册 → JWT Session
"""
import os
import base64
import hashlib
import secrets
import time
import logging
import urllib.parse
from typing import Optional

from fastapi import Depends, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse, JSONResponse

from auth_config import load_auth_config
from models_user import UserStore, ROLE_ADMIN, ROLE_MEMBER, ROLE_VIEWER

log = logging.getLogger("auth")

# ── Session 配置 ──────────────────────────────────────
SESSION_COOKIE = "fwb_session"
SESSION_MAX_AGE = 8 * 3600  # 8 小时
STATE_TTL = 600  # 10 分钟

# 内存 state 存储（生产可换 Redis）
_states: dict = {}  # state → {exp, redirect_to}
_sessions: dict = {}  # token → {sub, exp}

# OIDC client（懒初始化）— 保存发现文档 + 客户端凭据
_oidc_client: dict = {}
_oidc_metadata: dict = {}


# ════════════════════════════════════════════════════════
# PKCE 辅助（纯函数，可离线测试）
# ════════════════════════════════════════════════════════
def _gen_pkce_pair() -> tuple:
    """生成 (code_verifier, code_challenge)，S256 方法"""
    verifier = secrets.token_urlsafe(64)[:128]
    digest = hashlib.sha256(verifier.encode("ascii")).digest()
    challenge = base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
    return verifier, challenge


def _build_authorization_url(
    authorization_endpoint: str,
    client_id: str,
    redirect_uri: str,
    scope: str,
    state: str,
    nonce: str,
    code_challenge: str,
) -> str:
    """构建 Authorization Code + PKCE 授权跳转 URL"""
    params = {
        "response_type": "code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "scope": scope,
        "state": state,
        "nonce": nonce,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
    }
    sep = "&" if "?" in authorization_endpoint else "?"
    return f"{authorization_endpoint}{sep}{urllib.parse.urlencode(params)}"


# ════════════════════════════════════════════════════════
# 初始化
# ════════════════════════════════════════════════════════
def init_oidc(oidc_cfg: dict):
    """初始化 OIDC：拉取发现文档，缓存端点与凭据（不依赖 SessionMiddleware）"""
    import requests

    issuer = oidc_cfg["issuer"].rstrip("/")
    discovery = requests.get(
        f"{issuer}/.well-known/openid-configuration", timeout=5
    ).json()

    global _oidc_client, _oidc_metadata
    _oidc_metadata = discovery
    _oidc_client = {
        "client_id": oidc_cfg["client_id"],
        "client_secret": oidc_cfg.get("client_secret", ""),
        "scope": oidc_cfg.get("scope", "openid email profile groups"),
        "authorization_endpoint": discovery.get("authorization_endpoint", ""),
        "token_endpoint": discovery.get("token_endpoint", ""),
        "userinfo_endpoint": discovery.get("userinfo_endpoint", ""),
    }
    log.info(f"OIDC discovery OK — issuer={issuer}")
    return _oidc_client


# ════════════════════════════════════════════════════════
# 登录 / 回调 / 登出
# ════════════════════════════════════════════════════════
async def login_page(request: Request):
    """GET /login → 跳转 Authentik"""
    app = request.app
    auth_cfg = app.state.auth_config
    if auth_cfg.get("mode") != "oidc" or not auth_cfg.get("oidc", {}).get("enabled"):
        return JSONResponse({"detail": "OIDC not enabled"}, status_code=503)

    # 生成 state + nonce + PKCE
    state = secrets.token_urlsafe(32)
    nonce = secrets.token_urlsafe(16)
    verifier, challenge = _gen_pkce_pair()
    redirect_to = request.query_params.get("next", "/")

    _states[state] = {
        "exp": time.time() + STATE_TTL,
        "redirect_to": redirect_to,
        "code_verifier": verifier,
        "nonce": nonce,
    }

    # 清理过期 state
    now = time.time()
    expired = [s for s, v in _states.items() if v["exp"] < now]
    for s in expired:
        del _states[s]

    # 构建授权 URL
    oidc = app.state.oidc_client
    redirect_uri = auth_cfg["oidc"]["redirect_uri"]
    auth_url = _build_authorization_url(
        authorization_endpoint=oidc["authorization_endpoint"],
        client_id=oidc["client_id"],
        redirect_uri=redirect_uri,
        scope=oidc["scope"],
        state=state,
        nonce=nonce,
        code_challenge=challenge,
    )
    return RedirectResponse(auth_url, status_code=302)


async def _exchange_code_for_userinfo(
    oidc: dict, code: str, code_verifier: str, redirect_uri: str
) -> dict:
    """用授权码换取 token，再拉取 userinfo。失败抛 HTTPException。"""
    import httpx

    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": redirect_uri,
        "client_id": oidc["client_id"],
        "code_verifier": code_verifier,
    }
    # Confidential client：带上 client_secret
    if oidc.get("client_secret"):
        data["client_secret"] = oidc["client_secret"]

    async with httpx.AsyncClient(timeout=10) as client:
        tok_resp = await client.post(oidc["token_endpoint"], data=data)
        if tok_resp.status_code != 200:
            log.error(f"Token exchange failed: {tok_resp.status_code} {tok_resp.text[:300]}")
            raise HTTPException(400, "OIDC token exchange failed")
        token = tok_resp.json()

        access_token = token.get("access_token")
        if not access_token:
            raise HTTPException(400, "OIDC: no access_token in response")

        ui_resp = await client.get(
            oidc["userinfo_endpoint"],
            headers={"Authorization": f"Bearer {access_token}"},
        )
        if ui_resp.status_code != 200:
            log.error(f"Userinfo failed: {ui_resp.status_code} {ui_resp.text[:300]}")
            raise HTTPException(400, "OIDC userinfo fetch failed")
        return ui_resp.json()


async def auth_callback(request: Request):
    """GET /auth/callback → Authentik 回调"""
    app = request.app
    auth_cfg = app.state.auth_config
    user_store: UserStore = app.state.user_store
    oidc = app.state.oidc_client

    # 错误回参（如用户取消授权）
    if request.query_params.get("error"):
        raise HTTPException(
            400,
            f"OIDC error: {request.query_params.get('error')} "
            f"{request.query_params.get('error_description', '')}".strip(),
        )

    # 校验 state
    state = request.query_params.get("state", "")
    if state not in _states or _states[state]["exp"] < time.time():
        raise HTTPException(400, "Invalid or expired state")
    st = _states.pop(state)
    redirect_to = st["redirect_to"]
    code_verifier = st["code_verifier"]

    code = request.query_params.get("code", "")
    if not code:
        raise HTTPException(400, "Missing authorization code")

    # 用授权码 + PKCE code_verifier 换 token
    redirect_uri = auth_cfg["oidc"]["redirect_uri"]
    userinfo = await _exchange_code_for_userinfo(oidc, code, code_verifier, redirect_uri)

    # 解析用户
    sub = userinfo.get("sub") or userinfo.get("preferred_username") or userinfo.get("email")
    email = userinfo.get("email", "")
    name = userinfo.get("name") or userinfo.get("preferred_username", "")
    groups = userinfo.get("groups", []) or []

    if not sub:
        raise HTTPException(400, "OIDC: missing subject")

    # JIT 自动注册 / 更新
    role_mapping = auth_cfg.get("role_mapping", {})
    user = user_store.get_or_create(
        sub=sub,
        email=email,
        name=name,
        groups=groups,
        role_mapping=role_mapping,
    )

    if not user.is_active:
        raise HTTPException(403, "User disabled")

    # 创建本地 session
    session_token = secrets.token_urlsafe(48)
    _sessions[session_token] = {
        "sub": sub,
        "exp": time.time() + SESSION_MAX_AGE,
    }
    # 清理过期 session
    now = time.time()
    expired = [t for t, v in _sessions.items() if v["exp"] < now]
    for t in expired:
        del _sessions[t]

    resp = RedirectResponse(redirect_to, status_code=302)
    resp.set_cookie(
        SESSION_COOKIE,
        session_token,
        max_age=SESSION_MAX_AGE,
        httponly=True,
        samesite="lax",
    )
    log.info(f"User logged in: {email} (role={user.role})")
    return resp


async def logout_page(request: Request):
    """GET /logout → 清本地 session + 跳 Authentik 登出"""
    app = request.app
    resp = RedirectResponse("/", status_code=302)
    token = request.cookies.get(SESSION_COOKIE)
    if token and token in _sessions:
        del _sessions[token]
    resp.delete_cookie(SESSION_COOKIE)

    # Authentik 全局登出（如果配了 end_session_endpoint）
    auth_cfg = app.state.auth_config
    if auth_cfg.get("mode") == "oidc":
        end_session = _oidc_metadata.get("end_session_endpoint")
        if end_session:
            return RedirectResponse(end_session, status_code=302)
    return resp


# ════════════════════════════════════════════════════════
# FastAPI 依赖注入
# ════════════════════════════════════════════════════════
async def get_current_user(request: Request) -> dict:
    """
    Depends：解析当前登录用户
    - OIDC 模式：从 cookie session 取
    - None 模式：返回虚拟 admin（向后兼容）
    """
    app = request.app
    auth_cfg = app.state.auth_config
    mode = auth_cfg.get("mode", "none")

    # ── None 模式：单用户，虚拟 admin ──
    if mode == "none":
        return {
            "sub": "local-admin",
            "email": "admin@local",
            "name": "Local Admin",
            "role": "admin",
            "is_active": True,
        }

    # ── OIDC 模式 ──
    token = request.cookies.get(SESSION_COOKIE)
    if not token or token not in _sessions:
        # 检查 Bearer token（API 调用场景）
        auth_header = request.headers.get("authorization", "")
        if auth_header.startswith("Bearer "):
            # 这里可以验 JWT（Authentik 直发），暂时简化为跳过
            pass
        raise HTTPException(status_code=401, detail="Not authenticated")

    sess = _sessions[token]
    if sess["exp"] < time.time():
        del _sessions[token]
        raise HTTPException(status_code=401, detail="Session expired")

    user_store: UserStore = app.state.user_store
    user = user_store.find_by_sub(sess["sub"])
    if not user or not user.is_active:
        raise HTTPException(status_code=403, detail="User not found or disabled")

    return user.to_public_dict()


async def require_admin(request: Request) -> dict:
    """Depends：要求 admin 角色"""
    user = await get_current_user(request)
    if user.get("role") != "admin":
        raise HTTPException(status_code=403, detail="Admin only")
    return user


# ════════════════════════════════════════════════════════
# Webhook Token 旁路（给 CI / n8n 用）
# ════════════════════════════════════════════════════════
async def check_webhook_auth(request: Request) -> bool:
    """
    Webhook 鉴权：
    - 没配 token → 放行（不安全但向后兼容）
    - 配了 token → 必须匹配（header 或 query 参数）
    """
    app = request.app
    auth_cfg = app.state.auth_config
    webhook_cfg = auth_cfg.get("webhook", {})
    expected = webhook_cfg.get("token", "").strip()

    if not expected:
        return True  # 未配置 token，放行

    # 从 header 取
    header_name = webhook_cfg.get("header_name", "X-Webhook-Token")
    provided = request.headers.get(header_name, "")
    if not provided:
        provided = request.query_params.get("token", "")

    return secrets.compare_digest(provided, expected)