67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
# -*- 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": "注册成功"}
|