desungongpai/app/main.py

from pathlib import Path

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from contextlib import asynccontextmanager
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.middleware import SlowAPIMiddleware
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

from app.database import init_db, async_session, engine
from app.tcp_server import tcp_manager
from app.config import settings
from app.routers import devices, locations, alarms, attendance, commands, bluetooth, beacons, heartbeats
from app.dependencies import verify_api_key

import asyncio
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Rate limiter (shared instance via app.extensions for router access)
from app.extensions import limiter


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    logger.info("Initializing database...")
    await init_db()
    # Reset all devices to offline on startup (stale state from previous run)
    try:
        from sqlalchemy import update
        from app.models import Device
        async with async_session() as session:
            async with session.begin():
                await session.execute(update(Device).values(status="offline"))
        logger.info("All devices reset to offline on startup")
    except Exception:
        logger.exception("Failed to reset device statuses on startup")
    logger.info("Starting TCP server on %s:%d", settings.TCP_HOST, settings.TCP_PORT)
    tcp_task = asyncio.create_task(tcp_manager.start(settings.TCP_HOST, settings.TCP_PORT))
    yield
    # Shutdown
    logger.info("Shutting down TCP server...")
    await tcp_manager.stop()
    tcp_task.cancel()

app = FastAPI(
    title="KKS Badge Management System / KKS工牌管理系统",
    description="""
## KKS P240 & P241 蓝牙工牌管理后台 API

### 功能模块 / Features:
- **设备管理 / Device Management** - 设备注册、状态监控
- **位置数据 / Location Data** - GPS/LBS/WIFI定位数据查询与轨迹回放
- **报警管理 / Alarm Management** - SOS、围栏、低电等报警处理
- **考勤管理 / Attendance** - 打卡记录查询与统计
- **指令管理 / Commands** - 远程指令下发与留言
- **蓝牙数据 / Bluetooth** - 蓝牙打卡与定位数据
- **信标管理 / Beacons** - 蓝牙信标注册与位置配置
- **心跳数据 / Heartbeats** - 设备心跳记录查询

### 认证 / Authentication:
设置 `API_KEY` 环境变量后，所有 `/api/` 请求需携带 `X-API-Key` 请求头。

### 通讯协议 / Protocol:
- TCP端口: {tcp_port} (设备连接)
- 支持协议: KKS P240/P241 通讯协议
    """.format(tcp_port=settings.TCP_PORT),
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan,
)

# Rate limiter — SlowAPIMiddleware applies default_limits to all routes
app.state.limiter = limiter
app.add_middleware(SlowAPIMiddleware)
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# CORS
_origins = [o.strip() for o in settings.CORS_ORIGINS.split(",") if o.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_origins,
    allow_credentials="*" not in _origins,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Global exception handler — prevent stack trace leaks
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    from fastapi.exceptions import RequestValidationError
    from starlette.exceptions import HTTPException as StarletteHTTPException
    # Let FastAPI handle its own exceptions
    if isinstance(exc, (StarletteHTTPException, RequestValidationError)):
        raise exc
    logger.exception("Unhandled exception on %s %s", request.method, request.url.path)
    return JSONResponse(
        status_code=500,
        content={"code": 500, "message": "Internal server error", "data": None},
    )


# Include API routers (all under /api/ prefix, protected by API key if configured)
_api_deps = [verify_api_key] if settings.API_KEY else []

app.include_router(devices.router, dependencies=[*_api_deps])
app.include_router(locations.router, dependencies=[*_api_deps])
app.include_router(alarms.router, dependencies=[*_api_deps])
app.include_router(attendance.router, dependencies=[*_api_deps])
app.include_router(commands.router, dependencies=[*_api_deps])
app.include_router(bluetooth.router, dependencies=[*_api_deps])
app.include_router(beacons.router, dependencies=[*_api_deps])
app.include_router(heartbeats.router, dependencies=[*_api_deps])

_STATIC_DIR = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(_STATIC_DIR)), name="static")


@app.get("/admin", response_class=HTMLResponse, tags=["Admin"])
async def admin_page():
    """管理后台页面 / Admin Dashboard"""
    html_path = _STATIC_DIR / "admin.html"
    return HTMLResponse(content=html_path.read_text(encoding="utf-8"))


@app.get("/", tags=["Root"])
async def root():
    return {
        "name": settings.APP_NAME,
        "version": "1.0.0",
        "docs": "/docs",
        "redoc": "/redoc",
        "admin": "/admin",
        "tcp_port": settings.TCP_PORT,
        "auth_enabled": settings.API_KEY is not None,
    }


@app.get("/health", tags=["Root"])
async def health():
    """Health check with database connectivity test."""
    db_ok = False
    try:
        from sqlalchemy import text
        async with engine.connect() as conn:
            await conn.execute(text("SELECT 1"))
        db_ok = True
    except Exception:
        logger.warning("Health check: database unreachable")

    status = "healthy" if db_ok else "degraded"
    return {
        "status": status,
        "database": "ok" if db_ok else "error",
        "connected_devices": len(tcp_manager.connections),
    }