Files
opencatd-open/internal/controller/team/team.go
2025-04-16 18:01:27 +08:00

564 lines
15 KiB
Go

package controller
import (
"errors"
"net/http"
"slices"
"strconv"
"strings"
"time"
"opencatd-open/internal/consts"
dto "opencatd-open/internal/dto/team"
"opencatd-open/internal/model"
service "opencatd-open/internal/service/team"
"opencatd-open/internal/utils"
"github.com/duke-git/lancet/v2/slice"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"gorm.io/gorm"
)
type Team struct {
db *gorm.DB
userService service.UserService
tokenService service.TokenService
keyService service.ApiKeyService
usageService service.UsageService
}
func NewTeam(userService service.UserService, tokenService service.TokenService, keyService service.ApiKeyService, usageService service.UsageService) *Team {
return &Team{
userService: userService,
tokenService: tokenService,
keyService: keyService,
usageService: usageService,
}
}
// initadmin
func (h *Team) InitAdmin(c *gin.Context) {
admin, err := h.userService.GetUser(c, 1)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
user := &model.User{
Name: "root",
Username: "root",
Password: "openteam",
Role: utils.ToPtr(consts.RoleRoot),
Tokens: []model.Token{
{
Name: "default",
Key: "sk-team-" + strings.ReplaceAll(uuid.New().String(), "-", ""),
UnlimitedQuota: utils.ToPtr(true),
},
},
}
if err := h.userService.CreateUser(c, user); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
var result = dto.UserInfo{
ID: user.ID,
Name: user.Username,
Token: user.Tokens[0].Key,
Status: utils.ToPtr(user.Status == consts.StatusEnabled),
}
c.JSON(http.StatusOK, result)
return
} else {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
if admin != nil {
c.JSON(http.StatusForbidden, gin.H{
"error": "super user already exists, use cli to reset password",
})
return
}
}
func (h *Team) Me(c *gin.Context) {
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "token not found"})
return
}
userToken := token.(*model.Token)
c.JSON(http.StatusOK, dto.UserInfo{
ID: userToken.UserID,
Name: userToken.User.Name,
Token: userToken.Key,
Status: utils.ToPtr(userToken.User.Status == consts.StatusEnabled),
})
}
// CreateUser 创建用户
func (h *Team) CreateUser(c *gin.Context) {
var userReq dto.UserInfo
if err := c.ShouldBindJSON(&userReq); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid input"})
return
}
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "Unauthorized"})
return
}
userToken := token.(*model.Token)
if *userToken.User.Role < consts.RoleAdmin { // 普通用户只能创建自己的token
create := &model.Token{
Name: userReq.Name,
Key: "sk-team-" + strings.ReplaceAll(uuid.New().String(), "-", ""),
}
if userReq.Token != "" {
_key := strings.ReplaceAll(userReq.Token, "-", "")
create.Key = "sk-team-" + strings.ReplaceAll(_key, " ", "")
}
if err := h.tokenService.Create(c.Request.Context(), create); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
} else {
user := &model.User{
Name: userReq.Name,
Username: userReq.Name,
Role: utils.ToPtr(consts.RoleUser),
Tokens: []model.Token{
{
Name: "default",
Key: "sk-team-" + strings.ReplaceAll(uuid.New().String(), "-", ""),
},
},
}
// 默认角色为普通用户
if err := h.userService.CreateUser(c.Request.Context(), user); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
c.JSON(http.StatusOK, gin.H{"message": "ok"})
}
// GetUser 获取用户信息
func (h *Team) GetUser(c *gin.Context) {
idStr := c.Param("id")
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
user, err := h.userService.GetUser(c.Request.Context(), id)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, user)
}
// UpdateUser 更新用户信息
func (h *Team) UpdateUser(c *gin.Context) {
var user model.User
if err := c.ShouldBindJSON(&user); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid input"})
return
}
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "Unauthorized"})
return
}
userToken := token.(*model.Token)
operatorID := userToken.UserID // 假设从上下文中获取操作者ID
if err := h.userService.UpdateUser(c.Request.Context(), &user, operatorID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "ok"})
}
// DeleteUser 删除用户
func (h *Team) DeleteUser(c *gin.Context) {
idStr := c.Param("id")
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "Unauthorized"})
return
}
userToken := token.(*model.Token)
if *userToken.User.Role < consts.RoleAdmin { // 用户只能删除自己的token
err := h.tokenService.Delete(c.Request.Context(), id)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
} else {
if err := h.userService.DeleteUser(c, id, userToken.UserID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
c.JSON(http.StatusOK, gin.H{"message": "ok"})
}
func (h *Team) ListUsages(c *gin.Context) {
fromStr := c.Query("from")
toStr := c.Query("to")
var from, to time.Time
loc, _ := time.LoadLocation("Local")
var listUsage []*dto.UsageInfo
var err error
if fromStr != "" && toStr != "" {
from, err = time.Parse("2006-01-02", fromStr)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid from date"})
return
}
to, err = time.Parse("2006-01-02", toStr)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid to date"})
return
}
} else {
year, month, _ := time.Now().In(loc).Date()
from = time.Date(year, month, 1, 0, 0, 0, 0, loc)
to = from.AddDate(0, 1, 0)
}
token, _ := c.Get("token")
userToken := token.(*model.Token)
if *userToken.User.Role < consts.RoleAdmin {
listUsage, err = h.usageService.ListByDateRange(c.Request.Context(), from, to, map[string]interface{}{"user_id": userToken.UserID})
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
} else {
listUsage, err = h.usageService.ListByDateRange(c.Request.Context(), from, to, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
c.JSON(http.StatusOK, listUsage)
}
// ListUsers 获取用户列表
func (h *Team) ListUsers(c *gin.Context) {
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100"))
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
active := c.DefaultQuery("active", "")
if !slices.Contains([]string{"true", "false", ""}, active) {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid active value"})
return
}
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "Unauthorized"})
return
}
userToken := token.(*model.Token)
if *userToken.User.Role < consts.RoleAdmin { // 用户只能获取自己的token
tokens, _, err := h.tokenService.Lists(c, limit, offset)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
var userDTOs []dto.UserInfo
for _, token := range tokens {
userDTOs = append(userDTOs, dto.UserInfo{
ID: token.User.ID,
Name: token.User.Name,
Token: token.Key,
Status: utils.ToPtr(token.User.Status == consts.StatusEnabled),
})
}
c.JSON(http.StatusOK, userDTOs)
return
}
users, err := h.userService.ListUsers(c, limit, offset, active)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
var userDTOs []dto.UserInfo
for _, user := range users {
useres := dto.UserInfo{
ID: user.ID,
Name: user.Name,
Status: utils.ToPtr(user.Status == consts.StatusEnabled),
}
if len(user.Tokens) > 0 {
useres.Token = user.Tokens[0].Key
}
userDTOs = append(userDTOs, useres)
}
c.JSON(http.StatusOK, userDTOs)
}
func (h *Team) ResetUserToken(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "Unauthorized"})
return
}
userToken := token.(*model.Token)
findtoken, err := h.tokenService.GetByUserID(c, id)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
findtoken.Key = "sk-team-" + strings.ReplaceAll(uuid.New().String(), "-", "")
if *userToken.User.Role < consts.RoleAdmin { // 非管理员只能修改自己的token
if *userToken.User.Role <= *findtoken.User.Role || userToken.UserID != findtoken.UserID {
c.JSON(http.StatusForbidden, gin.H{"error": "forbidden"})
return
}
err := h.tokenService.UpdateWithCondition(c, findtoken, map[string]interface{}{"user_id": userToken.UserID}, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
} else {
if err := h.tokenService.Update(c, findtoken); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
c.JSON(http.StatusOK, dto.UserInfo{
ID: findtoken.User.ID,
Name: findtoken.User.Name,
Token: findtoken.Key,
})
}
func (h *Team) CreateKey(c *gin.Context) {
token, exists := c.Get("token")
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "token not found"})
return
}
userToken := token.(*model.Token)
if *userToken.User.Role < consts.RoleAdmin {
c.JSON(http.StatusForbidden, gin.H{"error": "forbidden"})
return
}
var key dto.ApiKeyInfo
if err := c.ShouldBindJSON(&key); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
err := h.keyService.Create(&model.ApiKey{
Name: utils.ToPtr(key.Name),
ApiType: utils.ToPtr(key.ApiType),
ApiKey: utils.ToPtr(key.Key),
Endpoint: utils.ToPtr(key.Endpoint),
})
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, key)
}
func (h *Team) ListKeys(c *gin.Context) {
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
active := c.Query("active")
if !slice.Contain([]string{"true", "false", ""}, active) {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid active value"})
return
}
keys, err := h.keyService.List(limit, offset, active)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
var keysDTO []dto.ApiKeyInfo
for _, key := range keys {
keylength := len(*key.ApiKey) / 3
if keylength < 1 {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key length"})
return
}
keysDTO = append(keysDTO, dto.ApiKeyInfo{
ID: int(key.ID),
Name: *key.Name,
ApiType: *key.ApiType,
Endpoint: *key.Endpoint,
Key: *key.ApiKey,
})
}
c.JSON(http.StatusOK, keysDTO)
}
func (h *Team) UpdateKey(c *gin.Context) {
// 1. 获取并验证ID
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key id"})
return
}
// 2. 解析请求体
var updateKey dto.ApiKeyInfo // 更明确的命名
if err := c.ShouldBindJSON(&updateKey); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
// 3. 获取现有记录
existingKey, err := h.keyService.GetByID(id)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
// 4. 使用 UpdateFields 方法统一处理字段更新
updatedKey := updateKey.UpdateFields(existingKey)
// 5. 保存更新
if err := h.keyService.Update(updatedKey); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, updatedKey)
}
func (h *Team) DeleteKey(c *gin.Context) {
// 1. 获取并验证ID
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key id"})
return
}
// 2. 删除记录
if err := h.keyService.Delete(id); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "ok"})
}
// ChangePassword 修改密码
func (h *Team) ChangePassword(c *gin.Context) {
userID := c.GetInt64("userID") // 假设从上下文中获取用户ID
var req struct {
OldPassword string `json:"oldPassword"`
NewPassword string `json:"newPassword"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid input"})
return
}
if err := h.userService.ChangePassword(c.Request.Context(), userID, req.OldPassword, req.NewPassword); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "ok"})
}
// ResetPassword 重置密码
func (h *Team) ResetPassword(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
operatorID := int64(c.GetInt("userID")) // 假设从上下文中获取操作者ID
if err := h.userService.ResetPassword(c.Request.Context(), id, operatorID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "password reset successfully"})
}
// EnableUser 启用用户
func (h *Team) EnableUser(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
operatorID := int64(c.GetInt("userID")) // 假设从上下文中获取操作者ID
if err := h.userService.BatchEnableUsers(c, []int64{id}, operatorID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "user enabled successfully"})
}
// DisableUser 禁用用户
func (h *Team) DisableUser(c *gin.Context) {
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid user id"})
return
}
operatorID := int64(c.GetInt("userID")) // 假设从上下文中获取操作者ID
if err := h.userService.BatchDisableUsers(c.Request.Context(), []int64{id}, operatorID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"message": "user disabled successfully"})
}