CostEvalPlatform/api/login_reg.py

# -*- 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(days=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": "注册成功"}