live-hub-py/web/backend/main.py

"""FastAPI 入口"""

import os
from contextlib import asynccontextmanager
from pathlib import Path

import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware

from utils import setup_logger

from .database import init_db
from .routers import (
    account_check,
    accounts,
    audit,
    auth,
    cookies,
    dashboard,
    douyu,
    huya,
    login,
    proxy,
    users,
    yyb,
)
from .schemas import AppInfo
from .version import get_app_version


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 本地 dev.sh 和 Docker 均显式设置 LOG_DIR；直接运行时回退到项目 logs/。
    _log_level = os.getenv("LOG_LEVEL", "INFO")
    _log_dir = Path(
        os.getenv("LOG_DIR", str(Path(__file__).resolve().parents[2] / "logs"))
    )
    setup_logger(level=_log_level, log_dir=str(_log_dir))

    init_db()
    # 进程重启后内存批次丢失，清理历史 pending/running，避免前端被假活跃批次锁死。
    from loguru import logger

    from .database import SessionLocal
    from .services.douyu_service import cleanup_orphan_douyu_tasks
    from .services.huya_service import cleanup_orphan_huya_tasks
    from .services.login_service import (
        cleanup_orphan_login_tasks,
        cleanup_orphan_relogin_tasks,
    )

    db = SessionLocal()
    try:
        cleaned = cleanup_orphan_huya_tasks(db, message="任务已中断（服务重启）")
        if cleaned:
            logger.info(f"启动清理虎牙残留任务: {cleaned} 条")
        cleaned_douyu = cleanup_orphan_douyu_tasks(db, message="任务已中断（服务重启）")
        if cleaned_douyu:
            logger.info(f"启动清理斗鱼残留任务: {cleaned_douyu} 条")
        cleaned_relogin = cleanup_orphan_relogin_tasks(db)
        if cleaned_relogin:
            logger.info(f"启动清理 CK 重登残留任务: {cleaned_relogin} 条")
        cleaned_login = cleanup_orphan_login_tasks(db)
        if cleaned_login:
            logger.info(f"启动清理登录残留任务: {cleaned_login} 条")
        from .services.yyb_service import cleanup_orphan_yyb_tasks

        cleaned_yyb = cleanup_orphan_yyb_tasks(db, message="任务已中断（服务重启）")
        if cleaned_yyb:
            logger.info(f"启动清理应用宝残留任务: {cleaned_yyb} 条")
    finally:
        db.close()
    yield


app = FastAPI(
    title="直播运营后台",
    version=get_app_version(),
    lifespan=lifespan,
)

# CORS（支持通过环境变量配置，逗号分隔；默认开发环境）
_cors_env = os.getenv("CORS_ORIGINS", "")
if _cors_env:
    _cors_origins = [o.strip() for o in _cors_env.split(",") if o.strip()]
else:
    _cors_origins = [
        "http://localhost:5174",
        "http://localhost:5173",
        "http://localhost:3000",
    ]

app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 压缩 JSON 和静态资源响应，避免任务列表轮询反复传输大体积文本。
app.add_middleware(GZipMiddleware, minimum_size=1024)


# 安全响应头中间件
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        if request.url.path.startswith("/assets/"):
            # Vite 产物文件名带 hash，可以长期缓存。
            response.headers["Cache-Control"] = "public, max-age=31536000, immutable"
        return response


app.add_middleware(SecurityHeadersMiddleware)

# 注册路由
app.include_router(auth.router)
app.include_router(users.router)
app.include_router(accounts.router)
app.include_router(account_check.router)
app.include_router(dashboard.router)
app.include_router(login.router)
app.include_router(proxy.router)
app.include_router(cookies.router)
app.include_router(huya.router)
app.include_router(douyu.router)
app.include_router(yyb.router)
app.include_router(audit.router)


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


@app.get("/api/app-info", response_model=AppInfo)
def app_info():
    return AppInfo(version=get_app_version())


# ---- 生产环境：serve 前端静态文件 ----
# Docker 部署时前端构建产物会被复制到 web/frontend/dist
_FRONTEND_DIST = Path(__file__).resolve().parents[2] / "web" / "frontend" / "dist"
_INDEX_HTML = _FRONTEND_DIST / "index.html"

if _INDEX_HTML.exists():
    # 挂载静态资源目录（js/css/图片等）
    _ASSETS_DIR = _FRONTEND_DIST / "assets"
    if _ASSETS_DIR.exists():
        app.mount("/assets", StaticFiles(directory=str(_ASSETS_DIR)), name="assets")

    @app.get("/{full_path:path}")
    async def serve_spa(full_path: str, request: Request):
        """SPA fallback：非 /api 路径返回 index.html"""
        # 排除 API 路径
        if full_path.startswith("api"):
            return {"detail": "Not Found"}
        # 尝试返回静态文件（防护路径遍历）
        file_path = (_FRONTEND_DIST / full_path).resolve()
        try:
            file_path.relative_to(_FRONTEND_DIST.resolve())
        except ValueError:
            # 路径逃逸出 dist 目录，拒绝访问
            return {"detail": "Not Found"}
        if file_path.is_file():
            return FileResponse(str(file_path))
        # SPA fallback 到 index.html
        return FileResponse(str(_INDEX_HTML))


def run():
    _reload = os.getenv("UVICORN_RELOAD", "false").lower() == "true"
    _port = int(os.getenv("BACKEND_PORT", "8800"))
    uvicorn.run("web.backend.main:app", host="0.0.0.0", port=_port, reload=_reload)


if __name__ == "__main__":
    run()