CostEvalPlatform/api/login_reg.py
2024-11-20 22:08:21 +08:00

75 lines
2.5 KiB
Python

# -*- coding: utf-8 -*-
# @Time : 2024/11/19 下午7:27
# @Author : 河瞬
# @FileName: login_reg.py
# @Software: PyCharm
from datetime import datetime, timedelta
from typing import Optional
from fastapi import APIRouter
from dependencies import *
from models import Tenant
router = APIRouter()
# 生成JWT token
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None, settings: SettingsDep = SettingsDep):
to_encode = data.copy()
if expires_delta:
expire = datetime.utcnow() + expires_delta
else:
expire = datetime.utcnow() + timedelta(minutes=15)
to_encode.update({"exp": expire})
print(settings, type(settings))
encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
return encoded_jwt
# 登录路由
@router.post("/api/s1/login")
async def login(response: Response, user_data: dict, session: SessionDep):
if user_data.get('username') is None or user_data.get('password') is None:
raise HTTPException(status_code=401, detail="用户名或密码不能为空")
# 查询用户
user = session.exec(select(User).where(User.username == user_data['username'])).first()
# 验证用户名和密码
if not user or user.password != user_data['password']:
raise HTTPException(status_code=401, detail="登录失败,用户名或密码错误")
# 生成JWT token
token = create_access_token(data={"id": user.id, "role": user.role, "tanant_id": user.tenant.id})
# 设置cookie
response.set_cookie(
key="session_token",
value=token,
httponly=True,
# domain=".ecs2.heshunme.xyz", # 确保域正确
samesite="none", # 根据需要设置samesite属性
secure=True,
)
# 关闭数据库会话
session.close()
return {"message": f"Login successful", "role": user.role}
@router.post("/api/s1/register")
async def register(data: dict, session: SessionDep):
if session.exec(select(Tenant).where(Tenant.name == data['name'])).first():
raise HTTPException(status_code=409, detail="租户名已存在")
if session.exec(select(User).where(User.username == data['username'])).first():
raise HTTPException(status_code=409, detail="用户名已存在")
tenant = Tenant(name=data['name'])
user = User(username=data['username'], password=data['password'], role=1, tenant=tenant)
session.add(tenant)
session.add(user)
session.commit()
session.close()
return {"detail": "注册成功"}