feat: 登录状态控制

This commit is contained in:
李如威 2025-08-22 18:44:12 +08:00
parent bd4bcf0aaf
commit 46607a68df
11 changed files with 148 additions and 27 deletions

View File

@ -1,6 +1,6 @@
{ {
"is_init": true, "is_init": true,
"db_type": "sqlite", "db_type": "mysql",
"mysql": { "mysql": {
"host": "", "host": "",
"port": "", "port": "",
@ -8,9 +8,10 @@
"password": "", "password": "",
"db_name": "" "db_name": ""
}, },
"db_host": "string", "db_host": "localhost",
"db_port": "string", "db_port": 3306,
"db_user": "string", "db_user": "root",
"db_password": "string", "db_password": "12345",
"db_name": "your_workflow" "db_name": "your_workflow",
"admin_password": "admin"
} }

View File

@ -1,13 +1,23 @@
from tortoise.contrib.fastapi import Tortoise from tortoise.contrib.fastapi import Tortoise
from .logger import logger
from pathlib import Path
from models.db.user import User
import hashlib
async def init_tortoise(db_type: str, db_config: dict):
async def init_tortoise(
db_type: str, db_config: dict, admin_password: str | None = None
):
if db_type == "mysql": if db_type == "mysql":
db_url = ( db_url = (
f"mysql://{db_config['db_user']}:{db_config['db_password']}" f"mysql://{db_config['db_user']}:{db_config['db_password']}"
f"@{db_config['db_host']}:{db_config['db_port']}/{db_config['db_name']}" f"@{db_config['db_host']}:{db_config['db_port']}/{db_config['db_name']}"
) )
elif db_type == "sqlite": elif db_type == "sqlite":
db_url = f"sqlite://{db_config['db_name']}" BASE_DIR = Path(__file__).parent.parent.resolve() / "sqlite"
BASE_DIR.mkdir(parents=True, exist_ok=True)
db_file = BASE_DIR / db_config["db_name"]
db_url = f"sqlite://{db_file}"
else: else:
raise ValueError("Unsupported db_type") raise ValueError("Unsupported db_type")
try: try:
@ -16,9 +26,18 @@ async def init_tortoise(db_type: str, db_config: dict):
modules={"models": ["models.db"]}, # 这里写你模型所在路径 modules={"models": ["models.db"]}, # 这里写你模型所在路径
) )
await Tortoise.generate_schemas() await Tortoise.generate_schemas()
if admin_password:
admin_user = await User.filter(username="admin").first()
if not admin_user:
await User(
password=hashlib.md5(admin_password.encode()).hexdigest(), username="admin", level=0
).save()
return True return True
except Exception as e: except Exception as e:
logger.error(e)
raise e raise e
async def close_tortoise(): async def close_tortoise():
await Tortoise.close_connections() await Tortoise.close_connections()

View File

@ -1,7 +1,33 @@
import logging import logging
logging.basicConfig( # 定义颜色
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s:\n%(message)s" RESET = "\033[0m"
) COLORS = {
"DEBUG": "\033[36m", # 青色
"INFO": "\033[32m", # 绿色
"WARNING": "\033[33m", # 黄色
"ERROR": "\033[31m", # 红色
"CRITICAL": "\033[41m", # 红底
}
class ColorFormatter(logging.Formatter):
def format(self, record):
color = COLORS.get(record.levelname, RESET)
record.levelname = f"{color}{record.levelname}{RESET}"
return super().format(record)
# 创建 logger
logger = logging.getLogger("workflow") logger = logging.getLogger("workflow")
logger.setLevel(logging.DEBUG)
# 创建控制台 handler
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
# 设置带颜色的 formatter
formatter = ColorFormatter("%(levelname)s: \t %(name)s - %(asctime)s - %(message)s")
ch.setFormatter(formatter)
logger.addHandler(ch)

View File

@ -1,10 +1,43 @@
import uvicorn import uvicorn
from fastapi import FastAPI import models.db as DB
from models.schemas.system import InitStatusResponse, InitConfigRequest 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.config import load_init_config, save_init_config
from core.db import init_tortoise, close_tortoise from core.db import init_tortoise, close_tortoise
from core.logger import logger from core.logger import logger
from contextlib import asynccontextmanager 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(): async def is_init():
@ -15,16 +48,23 @@ async def is_init():
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
# 数据库
if await is_init(): if await is_init():
logger.info("初始化数据库")
init_config = await load_init_config() init_config = await load_init_config()
await init_tortoise(init_config["db_type"], 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 yield
logger.info("关闭数据库") scheduler.shutdown()
logger.info("关闭定时器")
await close_tortoise() await close_tortoise()
logger.info("关闭数据库")
def create_app(): def create_app():
@ -62,9 +102,10 @@ def create_app():
init_config["db_user"] = request.db_user init_config["db_user"] = request.db_user
init_config["db_password"] = request.db_password init_config["db_password"] = request.db_password
init_config["db_name"] = request.db_name init_config["db_name"] = request.db_name
init_config["admin_password"] = request.admin_password
init_config["is_init"] = True init_config["is_init"] = True
init_success = await init_tortoise(request.db_type, init_config) init_success = await init_tortoise(request.db_type, init_config, request.admin_password)
if init_success: if init_success:
await save_init_config(init_config) await save_init_config(init_config)
@ -73,6 +114,26 @@ def create_app():
return InitStatusResponse(success=False, message="初始化数据库失败") return InitStatusResponse(success=False, message="初始化数据库失败")
except Exception as e: except Exception as e:
return InitStatusResponse(success=False, message=str(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 return app

View File

@ -2,8 +2,8 @@ from tortoise import fields
from .system import BaseModel from .system import BaseModel
class User(BaseModel): class User(BaseModel):
username = fields.CharField(max_length=50, unique=True) username = fields.CharField(max_length=50, null=True)
password = fields.CharField(max_length=128) password = fields.CharField(max_length=128, null=True)
level = fields.IntField(default=1, null=True, description="0:超级用户 1:普通用户") level = fields.IntField(default=1, null=True, description="0:超级用户 1:普通用户")
class Meta: class Meta:

View File

@ -7,13 +7,13 @@ class BaseResponse(BaseModel):
class InitConfigRequest(BaseModel): class InitConfigRequest(BaseModel):
db_type: str db_type: str = "sqlite"
admin_password: str admin_password: str = "admin"
db_host: str | None = None db_host: str | None = "localhost"
db_port: str | None = None db_port: int | None = 3306
db_user: str | None = None db_user: str | None = "root"
db_password: str | None = None db_password: str | None = "12345"
db_name: str | None = None db_name: str | None = "your_workflow"
class InitStatusResponse(BaseResponse): class InitStatusResponse(BaseResponse):

View File

@ -0,0 +1,13 @@
from pydantic import BaseModel
from .system import BaseResponse
import hashlib
class LoginRequest(BaseModel):
username: str = "admin"
password: str = "admin"
def get_md5_password(self):
return hashlib.md5(self.password.encode()).hexdigest()
class LoginResponse(BaseResponse):
data: dict | None

View File

@ -2,3 +2,4 @@ fastapi==0.116.1
uvicorn==0.35.0 uvicorn==0.35.0
tortoise-orm==0.25.1 tortoise-orm==0.25.1
aiomysql==0.2.0 aiomysql==0.2.0
apscheduler==3.11.0

Binary file not shown.

Binary file not shown.