feat: 登录状态控制
This commit is contained in:
parent
bd4bcf0aaf
commit
46607a68df
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"is_init": true,
|
||||
"db_type": "sqlite",
|
||||
"db_type": "mysql",
|
||||
"mysql": {
|
||||
"host": "",
|
||||
"port": "",
|
||||
|
|
@ -8,9 +8,10 @@
|
|||
"password": "",
|
||||
"db_name": ""
|
||||
},
|
||||
"db_host": "string",
|
||||
"db_port": "string",
|
||||
"db_user": "string",
|
||||
"db_password": "string",
|
||||
"db_name": "your_workflow"
|
||||
"db_host": "localhost",
|
||||
"db_port": 3306,
|
||||
"db_user": "root",
|
||||
"db_password": "12345",
|
||||
"db_name": "your_workflow",
|
||||
"admin_password": "admin"
|
||||
}
|
||||
|
|
@ -1,13 +1,23 @@
|
|||
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":
|
||||
db_url = (
|
||||
f"mysql://{db_config['db_user']}:{db_config['db_password']}"
|
||||
f"@{db_config['db_host']}:{db_config['db_port']}/{db_config['db_name']}"
|
||||
)
|
||||
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:
|
||||
raise ValueError("Unsupported db_type")
|
||||
try:
|
||||
|
|
@ -16,9 +26,18 @@ async def init_tortoise(db_type: str, db_config: dict):
|
|||
modules={"models": ["models.db"]}, # 这里写你模型所在路径
|
||||
)
|
||||
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
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
raise e
|
||||
|
||||
|
||||
async def close_tortoise():
|
||||
await Tortoise.close_connections()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,33 @@
|
|||
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.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)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,43 @@
|
|||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from models.schemas.system import InitStatusResponse, InitConfigRequest
|
||||
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():
|
||||
|
|
@ -15,16 +48,23 @@ async def is_init():
|
|||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
|
||||
# 数据库
|
||||
if await is_init():
|
||||
logger.info("初始化数据库")
|
||||
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
|
||||
|
||||
logger.info("关闭数据库")
|
||||
scheduler.shutdown()
|
||||
logger.info("关闭定时器")
|
||||
|
||||
await close_tortoise()
|
||||
logger.info("关闭数据库")
|
||||
|
||||
|
||||
def create_app():
|
||||
|
|
@ -62,9 +102,10 @@ def create_app():
|
|||
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)
|
||||
init_success = await init_tortoise(request.db_type, init_config, request.admin_password)
|
||||
|
||||
if init_success:
|
||||
await save_init_config(init_config)
|
||||
|
|
@ -73,6 +114,26 @@ def create_app():
|
|||
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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@ from tortoise import fields
|
|||
from .system import BaseModel
|
||||
|
||||
class User(BaseModel):
|
||||
username = fields.CharField(max_length=50, unique=True)
|
||||
password = fields.CharField(max_length=128)
|
||||
username = fields.CharField(max_length=50, null=True)
|
||||
password = fields.CharField(max_length=128, null=True)
|
||||
level = fields.IntField(default=1, null=True, description="0:超级用户 1:普通用户")
|
||||
|
||||
class Meta:
|
||||
|
|
|
|||
|
|
@ -7,13 +7,13 @@ class BaseResponse(BaseModel):
|
|||
|
||||
|
||||
class InitConfigRequest(BaseModel):
|
||||
db_type: str
|
||||
admin_password: str
|
||||
db_host: str | None = None
|
||||
db_port: str | None = None
|
||||
db_user: str | None = None
|
||||
db_password: str | None = None
|
||||
db_name: str | None = None
|
||||
db_type: str = "sqlite"
|
||||
admin_password: str = "admin"
|
||||
db_host: str | None = "localhost"
|
||||
db_port: int | None = 3306
|
||||
db_user: str | None = "root"
|
||||
db_password: str | None = "12345"
|
||||
db_name: str | None = "your_workflow"
|
||||
|
||||
|
||||
class InitStatusResponse(BaseResponse):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -2,3 +2,4 @@ fastapi==0.116.1
|
|||
uvicorn==0.35.0
|
||||
tortoise-orm==0.25.1
|
||||
aiomysql==0.2.0
|
||||
apscheduler==3.11.0
|
||||
Binary file not shown.
Binary file not shown.
Loading…
Reference in New Issue