import uvicorn import models.db as DB import uuid import time from fastapi import FastAPI, Depends, HTTPException from models.schemas.system import InitStatusResponse, InitConfigRequest, BaseResponse from models.schemas.user import LoginRequest, LoginResponse from core.config import load_init_config, save_init_config from core.db import init_tortoise, close_tortoise from core.logger import logger from contextlib import asynccontextmanager from fastapi.security import OAuth2PasswordBearer from apscheduler.schedulers.asyncio import AsyncIOScheduler # 简单的登陆状态存储 sessions: dict[str, tuple[float, DB.User]] = {} oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") scheduler = AsyncIOScheduler() # 定时清理过期session def clear_session(): TTL = 60 * 60 * 24 now = time.time() expired_keys = [token for token, (expires_at, _) in sessions.items() if expires_at + TTL < now] logger.info(f"检测 session, {now} {sessions}") for token in expired_keys: sessions.pop(token, None) if expired_keys: logger.info(f"清理过期 now: {now} session: {expired_keys}") def get_current_user(token:str = Depends(oauth2_scheme)) -> DB.User | None: if not token or token not in sessions: raise HTTPException(status_code=401, detail="用户未登陆") _, user = sessions[token] sessions[token] = time.time(), user return user async def is_init(): init_config = await load_init_config() init_status = init_config.get("is_init") return init_status @asynccontextmanager async def lifespan(app: FastAPI): # 数据库 if await is_init(): init_config = await load_init_config() await init_tortoise(init_config["db_type"], init_config, init_config.get("admin_password")) logger.info("初始化数据库") # 定时器 scheduler.add_job(clear_session, "interval", seconds=5) scheduler.start() logger.info("初始化定时器") yield scheduler.shutdown() logger.info("关闭定时器") await close_tortoise() logger.info("关闭数据库") def create_app(): app = FastAPI(title="Workflow API", version="1.0.0", lifespan=lifespan) @app.get("/") async def root(): return {"message": "hello"} @app.get( "/check_init", tags=["System"], description="系统是否初始化", response_model=InitStatusResponse, ) async def check_init(): return InitStatusResponse(data=await is_init()) @app.post( "/init_system", tags=["System"], description="系统初始化", response_model=InitStatusResponse, ) async def init_system(request: InitConfigRequest): logger.info("init_system") try: init_config = await load_init_config() if init_config.get('is_init'): return InitStatusResponse(data=True) # 保存配置 init_config["db_type"] = request.db_type init_config["db_host"] = request.db_host init_config["db_port"] = request.db_port init_config["db_user"] = request.db_user init_config["db_password"] = request.db_password init_config["db_name"] = request.db_name init_config["admin_password"] = request.admin_password init_config["is_init"] = True init_success = await init_tortoise(request.db_type, init_config, request.admin_password) if init_success: await save_init_config(init_config) return InitStatusResponse(data=True) return InitStatusResponse(success=False, message="初始化数据库失败") except Exception as e: return InitStatusResponse(success=False, message=str(e)) @app.post("/login", tags=["User"], description="用户登录", response_model=LoginResponse) async def login(login_info: LoginRequest): user = await DB.User.filter(username=login_info.username, password=login_info.get_md5_password()).first() if user: token = uuid.uuid4().hex sessions[token] = time.time(), user logger.info(sessions) return LoginResponse(data={ "token": token, "user": user.json(del_columns=["password"]) }) return LoginResponse(success=False, message="帐号或密码错误") @app.post("/logout", tags=["User"], description="退出登录", response_model=BaseResponse) async def logout(token:str): sessions.pop(token, None) logger.info(sessions) return BaseResponse() return app app = create_app() if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True)