CostEvalPlatform/main.py

56 lines
1.5 KiB
Python

from fastapi import FastAPI, HTTPException, Response, Depends
from typing import Optional
from datetime import datetime, timedelta
from jose import JWTError, jwt
from models import *
app = FastAPI()
# 创建数据库引擎
engine = create_engine('sqlite:///test.db')
# 创建所有表
Base.metadata.create_all(engine)
# 创建会话
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
session = SessionLocal()
# 用于生成和验证JWT的密钥
SECRET_KEY = "your_secret_key"
ALGORITHM = "HS256"
# 生成JWT token
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
to_encode = data.copy()
if expires_delta:
expire = datetime.utcnow() + expires_delta
else:
expire = datetime.utcnow() + timedelta(minutes=15)
to_encode.update({"exp": expire})
encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
return encoded_jwt
# 登录路由
@app.post("/api/s1/login")
async def login(response: Response, user_data: dict):
# 查询用户
user = session.query(User).filter(User.name == user_data['username']).first()
# 验证用户名和密码
if not user or user.password != user_data['password']:
raise HTTPException(status_code=401, detail="Login failed")
# 生成JWT token
token = create_access_token(data={"sub": user.name})
# 设置cookie
response.set_cookie(key="session_token", value=token, httponly=True)
# 关闭数据库会话
session.close()
return {"message": "Login successful"}