zTools2/app/main.py

"""FastAPI 应用工厂。

路由概览：
    GET  /                       -> 仅返回版本号
    GET  /docs                   -> Swagger UI（Basic Auth）
    GET  /redoc                  -> ReDoc       （Basic Auth）
    GET  /openapi.json           -> OpenAPI 文档（Basic Auth）
    GET  /health                 -> 存活探针（公开）
    GET  /upload                 -> 上传页面（公开 HTML）
    GET  /files                  -> 文件浏览页（Basic Auth，同 docs）
    GET  /whiteboard/{id}        -> 白板页面（公开，不存在则新建）
    GET  /whiteboard-admin       -> 白板管理页（Basic Auth，同 docs）
    GET  /api/...                -> 业务接口
    WS   /ws/whiteboard/{id}     -> 白板实时同步（公开）
    /static/...                  -> 前端静态资源（JS/CSS）
"""

from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import Depends, FastAPI
from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles

from .controllers import (
    chunk_upload_router,
    file_admin_router,
    file_router,
    system_router,
    tunnel_router,
    whiteboard_router,
)
from .database import init_db_schema
from .security import require_docs_auth
from .services.chunk_upload_service import ChunkUploadService
from .services.whiteboard_hub import get_hub
from .views.upload_html import render as render_upload_html

logger = logging.getLogger("zikai")

# 后台 reaper 的扫描间隔（秒）。不依赖 start.sh，进程存活期间持续清理过期会话。
_REAPER_INTERVAL_SECONDS = 60

_STATIC_DIR = Path(__file__).resolve().parent.parent / "static"


def _reap_once() -> None:
    """同步执行一次过期会话清理（在 worker 线程里跑，避免阻塞事件循环）。"""
    try:
        from .database import get_session_local
        from .dao.upload_session_dao import UploadSessionDAO
        from .dao.uploaded_file_dao import UploadedFileDAO

        db = get_session_local()()
        try:
            n = ChunkUploadService(UploadSessionDAO(db), UploadedFileDAO(db)).reap_stale_sessions()
            if n:
                logger.info("reaper 清理了 %d 个过期分片会话", n)
        finally:
            db.close()
    except Exception as exc:  # pragma: no cover
        logger.warning("reaper 执行失败：%s", exc)


async def _reaper_loop(stop: asyncio.Event) -> None:
    """周期性清理被放弃的分片上传会话。"""
    while not stop.is_set():
        try:
            await asyncio.to_thread(_reap_once)
        except Exception as exc:  # pragma: no cover
            logger.warning("reaper 循环异常：%s", exc)
        # 用 wait_for 实现「可被 stop 提前唤醒的 sleep」
        try:
            await asyncio.wait_for(stop.wait(), timeout=_REAPER_INTERVAL_SECONDS)
        except asyncio.TimeoutError:
            pass


def _tunnel_reap_once() -> None:
    """启动时清理 DB 里残留的 active 隧道会话（进程异常重启后的孤儿记录）。"""
    try:
        from .database import get_session_local
        from .dao.tunnel_session_dao import TunnelSessionDAO
        from .services.tunnel_service import TunnelService

        db = get_session_local()()
        try:
            n = TunnelService(TunnelSessionDAO(db)).reap_orphans()
            if n:
                logger.info("tunnel reaper 清理了 %d 个孤儿隧道会话", n)
        finally:
            db.close()
    except Exception as exc:  # pragma: no cover
        logger.warning("tunnel reaper 执行失败：%s", exc)


@asynccontextmanager
async def lifespan(app: FastAPI):
    try:
        init_db_schema()
        logger.info("数据库表已就绪。")
    except Exception as exc:  # pragma: no cover
        logger.error("初始化数据库失败：%s", exc)
    # 启动时清理残留的 active 隧道会话（进程重启后的孤儿记录）
    await asyncio.to_thread(_tunnel_reap_once)
    # 启动后台 reaper，清理被放弃的分片上传会话与临时文件
    stop = asyncio.Event()
    reaper = asyncio.create_task(_reaper_loop(stop))
    logger.info("分片会话 reaper 已启动（间隔 %ds）。", _REAPER_INTERVAL_SECONDS)
    # 启动白板心跳 reaper，清理失活的 WebSocket 连接
    wb_reaper = asyncio.create_task(get_hub().reap_loop(stop))
    logger.info("白板心跳 reaper 已启动。")
    try:
        yield
    finally:
        stop.set()
        for task in (reaper, wb_reaper):
            try:
                await asyncio.wait_for(task, timeout=5)
            except (asyncio.TimeoutError, asyncio.CancelledError):  # pragma: no cover
                task.cancel()


def create_app() -> FastAPI:
    app = FastAPI(
        title="zikai file service",
        description=(
            "f.zikai.wang 的文件上传与主机监控服务。\n\n"
            "- `GET /api/system/status` — CPU / 内存 / 磁盘（HTML 或 JSON）\n"
            "- `POST /api/files/upload` — 大文件流式上传\n"
            "- 内置 SFTP 服务（详见 README）\n\n"
            "/docs 等接口需 Basic Auth，凭据见 config.yaml 的 docs 段。"
        ),
        version="1.0.0",
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
        lifespan=lifespan,
    )

    app.include_router(system_router)
    app.include_router(file_router)
    app.include_router(file_admin_router)
    app.include_router(chunk_upload_router)
    app.include_router(tunnel_router)
    app.include_router(whiteboard_router)

    # 前端静态资源（JS/CSS）；HTML 壳由下面的具名路由返回，便于各自挂 Basic Auth
    if _STATIC_DIR.is_dir():
        app.mount("/static", StaticFiles(directory=str(_STATIC_DIR)), name="static")

    # 受 Basic Auth 保护的文档接口
    @app.get("/openapi.json")
    def protected_openapi(_: str = Depends(require_docs_auth)) -> JSONResponse:
        return JSONResponse(app.openapi())

    @app.get("/docs")
    def protected_docs(_: str = Depends(require_docs_auth)):
        return get_swagger_ui_html(
            openapi_url="/openapi.json", title="zikai docs", swagger_favicon_url=""
        )

    @app.get("/redoc")
    def protected_redoc(_: str = Depends(require_docs_auth)):
        return get_redoc_html(
            openapi_url="/openapi.json", title="zikai docs", redoc_favicon_url=""
        )

    @app.get("/", response_class=PlainTextResponse)
    def root() -> PlainTextResponse:
        return PlainTextResponse(f"zikai {app.version}\n")

    @app.get("/health")
    def health() -> dict:
        return {"status": "ok"}

    @app.get("/upload", response_class=HTMLResponse)
    def upload_page() -> HTMLResponse:
        """拖拽 / 多文件 / 分片上传页面（公开，对齐 /api/files/upload）。"""
        return HTMLResponse(render_upload_html())

    @app.get("/files", response_class=HTMLResponse)
    def files_page(_: str = Depends(require_docs_auth)) -> HTMLResponse:
        """文件浏览页（Basic Auth，同 docs）：列出/下载/删除已上传文件。"""
        return _serve_static_html("file_browser.html")

    @app.get("/whiteboard-admin", response_class=HTMLResponse)
    def whiteboard_admin_page(_: str = Depends(require_docs_auth)) -> HTMLResponse:
        """白板管理页（Basic Auth，同 docs）：查看/删除白板。"""
        return _serve_static_html("whiteboard_admin.html")

    @app.get("/whiteboard/{board_id}", response_class=HTMLResponse)
    def whiteboard_page(board_id: str) -> HTMLResponse:
        """白板页面（公开）：访问即协作，不存在则前端拉取时自动新建。"""
        return _serve_static_html("whiteboard.html")

    return app


def _serve_static_html(filename: str) -> HTMLResponse:
    """读取 static/ 下的 HTML 文件并返回；缺失时返回 404 文本。"""
    path = _STATIC_DIR / filename
    if not path.is_file():
        return HTMLResponse(f"<h1>未找到 {filename}</h1>", status_code=404)
    return HTMLResponse(path.read_text(encoding="utf-8"))


app = create_app()