desungongpai/app/main.py

from pathlib import Path

from fastapi import Depends, 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, api_keys, ws, geocoding
from app.dependencies import verify_api_key, require_write, require_admin

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


async def run_data_cleanup():
    """Delete records older than DATA_RETENTION_DAYS."""
    from datetime import datetime, timedelta
    from app.config import BEIJING_TZ
    from sqlalchemy import delete
    from app.models import LocationRecord, HeartbeatRecord, AlarmRecord, AttendanceRecord, BluetoothRecord

    cutoff = datetime.now(BEIJING_TZ) - timedelta(days=settings.DATA_RETENTION_DAYS)
    total_deleted = 0
    async with async_session() as session:
        async with session.begin():
            for model, time_col in [
                (LocationRecord, LocationRecord.created_at),
                (HeartbeatRecord, HeartbeatRecord.created_at),
                (AlarmRecord, AlarmRecord.created_at),
                (AttendanceRecord, AttendanceRecord.created_at),
                (BluetoothRecord, BluetoothRecord.created_at),
            ]:
                result = await session.execute(
                    delete(model).where(time_col < cutoff)
                )
                if result.rowcount:
                    total_deleted += result.rowcount
                    logger.info("Cleanup: deleted %d old %s records", result.rowcount, model.__tablename__)
    return total_deleted


async def _data_cleanup_loop():
    """Background task that runs cleanup periodically."""
    while True:
        try:
            await asyncio.sleep(settings.DATA_CLEANUP_INTERVAL_HOURS * 3600)
            deleted = await run_data_cleanup()
            if deleted:
                logger.info("Data cleanup completed: %d records removed", deleted)
        except asyncio.CancelledError:
            break
        except Exception:
            logger.exception("Data cleanup error")


@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")
    # Create missing indexes (safe for existing databases)
    try:
        from sqlalchemy import text as sa_text
        async with engine.begin() as conn:
            for stmt in [
                "CREATE INDEX IF NOT EXISTS ix_alarm_type ON alarm_records(alarm_type)",
                "CREATE INDEX IF NOT EXISTS ix_alarm_ack ON alarm_records(acknowledged)",
                "CREATE INDEX IF NOT EXISTS ix_bt_beacon_mac ON bluetooth_records(beacon_mac)",
                "CREATE INDEX IF NOT EXISTS ix_loc_type ON location_records(location_type)",
                "CREATE INDEX IF NOT EXISTS ix_att_type ON attendance_records(attendance_type)",
            ]:
                await conn.execute(sa_text(stmt))
        logger.info("Database indexes verified/created")
    except Exception:
        logger.exception("Failed to create indexes")

    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))
    cleanup_task = asyncio.create_task(_data_cleanup_loop())
    yield
    # Shutdown
    cleanup_task.cancel()
    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])
app.include_router(api_keys.router, dependencies=[*_api_deps])
app.include_router(ws.router)  # WebSocket handles auth internally
app.include_router(geocoding.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")

    from app.websocket_manager import ws_manager
    status = "healthy" if db_ok else "degraded"
    return {
        "status": status,
        "database": "ok" if db_ok else "error",
        "connected_devices": len(tcp_manager.connections),
        "websocket_connections": ws_manager.connection_count,
    }


@app.post("/api/system/cleanup", tags=["System / 系统管理"], dependencies=[Depends(require_admin)] if settings.API_KEY else [])
async def manual_cleanup():
    """手动触发数据清理 / Manually trigger data cleanup (admin only)."""
    try:
        deleted = await run_data_cleanup()
        return {"code": 0, "message": f"Cleanup completed: {deleted} records removed", "data": {"deleted": deleted}}
    except Exception as e:
        logger.exception("Manual cleanup failed")
        return {"code": 500, "message": f"Cleanup failed: {str(e)}", "data": None}