CostEvalPlatform/api/login_reg.py

# -*- coding: utf-8 -*-
# @Time    : 2024/11/19 下午7:27
# @Author  : 河瞬
# @FileName: login_reg.py
# @Software: PyCharm
from fastapi import HTTPException, Response, Depends, APIRouter
from typing import Optional, Annotated
from datetime import datetime, timedelta
from jose import JWTError, jwt

from sqlmodel import select

from models import Tenant, User, Project
from dependencies import *

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):
    # 查询用户
    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)

    # 关闭数据库会话
    session.close()

    return {"message": f"Login successful"}


@router.post("/api/s1/register")
async def register(data: dict, session: SessionDep):
    if session.exec(select(Tenant).where(Tenant.name == data['name'])):
        raise HTTPException(status_code=409, detail="租户名已存在")
    if session.exec(select(User).where(User.username == data['username'])):
        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.close()
    return {"detail": "注册成功"}