feat: 登录状态控制
This commit is contained in:
parent
bd4bcf0aaf
commit
46607a68df
|
|
@ -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"
|
||||||
}
|
}
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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):
|
||||||
|
|
|
||||||
|
|
@ -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
|
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.
Loading…
Reference in New Issue