opencatd-open/internal/controller/apikey.go

package controller

import (
	"bytes"
	"encoding/json"
	"net/http"
	"opencatd-open/internal/consts"
	"opencatd-open/internal/dto"
	"opencatd-open/internal/model"
	"opencatd-open/internal/utils"
	"strconv"
	"strings"

	"github.com/duke-git/lancet/v2/slice"
	"github.com/gin-gonic/gin"
)

func (a Api) CreateApiKey(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	newkey := new(model.ApiKey)
	err := c.ShouldBind(newkey)
	if err != nil {
		dto.Fail(c, 400, err.Error())
	}
	if slice.Contain([]string{"openai", "azure", "claude"}, *newkey.ApiType) {
		sma, err := utils.FetchKeyModel(a.db, newkey)
		if err == nil {
			newkey.SupportModelsArray = sma
			var buf = new(bytes.Buffer)
			json.NewEncoder(buf).Encode(sma) //nolint:errcheck
			newkey.SupportModels = utils.ToPtr(buf.String())
		}
	}

	err = a.keyService.CreateApiKey(c, newkey)
	if err != nil {
		dto.Fail(c, 400, err.Error())
	} else {
		dto.Success(c, nil)
	}

}

func (a Api) GetApiKey(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	id, err := strconv.ParseInt(c.Param("id"), 10, 64)
	if err != nil {
		dto.Fail(c, 400, err.Error())
		return
	}
	key, err := a.keyService.GetApiKey(c, id)
	if err != nil {
		dto.Fail(c, 400, err.Error())
	} else {
		dto.Success(c, key)
	}
}

func (a Api) ListApiKey(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	limit, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20"))
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	offset := (page - 1) * limit
	active := c.QueryArray("active[]")
	if !slice.ContainSubSlice([]string{"true", "false"}, active) {
		dto.Fail(c, http.StatusBadRequest, "active must be true or false")
		return
	}

	keys, total, err := a.keyService.ListApiKey(c, limit, offset, active)
	if err != nil {
		dto.Fail(c, 500, err.Error())
	} else {
		for _, key := range keys {
			str := *key.ApiKey
			slen := len(str)
			if slen > 20 {
				slen = 20
			}
			str = str[:slen]
			key.ApiKey = &str

			var sma []string
			json.NewDecoder(strings.NewReader(*key.SupportModels)).Decode(&sma) //nolint:errcheck
			key.SupportModelsArray = sma
		}
		dto.Success(c, gin.H{
			"total": total,
			"keys":  keys,
		})
	}
}

func (a Api) DeleteApiKey(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	var batchid dto.BatchIDRequest
	err := c.ShouldBind(&batchid)
	if err != nil {
		dto.Fail(c, 400, err.Error())
		return
	}

	err = a.keyService.DeleteApiKey(c, batchid.IDs)
	if err != nil {
		dto.Fail(c, 500, err.Error())
	} else {
		dto.Success(c, nil)
	}
}

func (a Api) UpdateApiKey(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	var req model.ApiKey
	err := c.ShouldBind(&req)
	if err != nil {
		dto.Fail(c, 400, err.Error())
		return
	}

	err = a.keyService.UpdateApiKey(c, &req)
	if err != nil {
		dto.Fail(c, 500, err.Error())
	} else {
		dto.Success(c, nil)
	}
}

func (a Api) ApiKeyOption(c *gin.Context) {
	role := c.MustGet("user_role").(*consts.UserRole)
	if *role < consts.RoleAdmin {
		dto.Fail(c, 403, "Permission denied")
		return
	}
	option := strings.ToLower(c.Param("option"))
	var batchid dto.BatchIDRequest
	err := c.ShouldBind(&batchid)
	if err != nil {
		dto.Fail(c, 400, err.Error())
		return
	}
	switch option {
	case "enable":
		err = a.keyService.EnableApiKey(c, batchid.IDs)
	case "disable":
		err = a.keyService.DisableApiKey(c, batchid.IDs)
	case "delete":
		err = a.keyService.DeleteApiKey(c, batchid.IDs)
	default:
		dto.Fail(c, 400, "invalid option, only support enable, disable, delete")
		return
	}
	if err != nil {
		dto.Fail(c, 400, err.Error())
		return
	}
	dto.Success(c, nil)
}