openteam/server/internal/api/keys.go

package api

import (
	"encoding/json"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/pkg/apikey"
	"github.com/openteam/server/internal/pkg/resp"
	"github.com/openteam/server/internal/store"
)

type createKeyReq struct {
	Name                *string  `json:"name" binding:"required,min=1,max=64"`
	QuotaTokensPerDay   *int64   `json:"quota_tokens_per_day"`
	QuotaRequestsPerDay *int     `json:"quota_requests_per_day"`
	AllowedModels       []string `json:"allowed_models"`
	ExpiresAt           *string  `json:"expires_at"` // RFC3339
}

// CreateKey POST /api/v1/keys — 创建密钥，明文仅此一次返回。
func (h *Handler) CreateKey(c *gin.Context) {
	u, ok := userFromContext(c)
	if !ok {
		resp.Fail(c, http.StatusUnauthorized, "authentication required")
		return
	}
	var req createKeyReq
	if err := c.ShouldBindJSON(&req); err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid input: "+err.Error())
		return
	}
	plain, hash, prefix, err := apikey.Generate()
	if err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to generate key")
		return
	}
	k := store.APIKey{
		UserID:              u.ID,
		Name:                *req.Name,
		KeyHash:             hash,
		KeyPrefix:           prefix,
		QuotaTokensPerDay:   req.QuotaTokensPerDay,
		QuotaRequestsPerDay: req.QuotaRequestsPerDay,
		AllowedModels:       req.AllowedModels,
		Status:              store.KeyStatusActive,
	}
	// 限额填 0 视为不限制（同 PatchKey），存成 0 会让密钥一发都发不出去
	if k.QuotaTokensPerDay != nil && *k.QuotaTokensPerDay <= 0 {
		k.QuotaTokensPerDay = nil
	}
	if k.QuotaRequestsPerDay != nil && *k.QuotaRequestsPerDay <= 0 {
		k.QuotaRequestsPerDay = nil
	}
	if req.ExpiresAt != nil {
		t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
		if err != nil {
			resp.Fail(c, http.StatusBadRequest, "expires_at must be RFC3339")
			return
		}
		k.ExpiresAt = &t
	}
	if err := h.a.DB.Create(&k).Error; err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to create key")
		return
	}
	resp.Created(c, gin.H{
		"id":         k.ID,
		"name":       k.Name,
		"key":        plain, // 仅此一次
		"key_prefix": k.KeyPrefix,
		"created_at": k.CreatedAt,
	})
}

// keyJSON 密钥的对外表示（列表与 PATCH 回显共用）。
func keyJSON(k store.APIKey) gin.H {
	return gin.H{
		"id":                     k.ID,
		"name":                   k.Name,
		"key_prefix":             k.KeyPrefix,
		"quota_tokens_per_day":   k.QuotaTokensPerDay,
		"quota_requests_per_day": k.QuotaRequestsPerDay,
		"allowed_models":         k.AllowedModels,
		"expires_at":             k.ExpiresAt,
		"status":                 k.Status,
		"last_used_at":           k.LastUsedAt,
		"created_at":             k.CreatedAt,
	}
}

// ListKeys GET /api/v1/keys
func (h *Handler) ListKeys(c *gin.Context) {
	u, ok := userFromContext(c)
	if !ok {
		resp.Fail(c, http.StatusUnauthorized, "authentication required")
		return
	}
	var keys []store.APIKey
	if err := h.a.DB.Where("user_id = ?", u.ID).Order("id DESC").Find(&keys).Error; err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to load keys")
		return
	}
	out := make([]gin.H, 0, len(keys))
	for _, k := range keys {
		out = append(out, keyJSON(k))
	}
	resp.OK(c, gin.H{"items": out})
}

// PatchKey PATCH /api/v1/keys/:id — 改名、限额、白名单、启停。
func (h *Handler) PatchKey(c *gin.Context) {
	u, ok := userFromContext(c)
	if !ok {
		resp.Fail(c, http.StatusUnauthorized, "authentication required")
		return
	}
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid key id")
		return
	}
	var req struct {
		Name                *string   `json:"name"`
		QuotaTokensPerDay   *int64    `json:"quota_tokens_per_day"`
		QuotaRequestsPerDay *int      `json:"quota_requests_per_day"`
		AllowedModels       *[]string `json:"allowed_models"`
		Status              *string   `json:"status"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid input")
		return
	}
	var k store.APIKey
	if err := h.a.DB.Where("id = ? AND user_id = ?", id, u.ID).First(&k).Error; err != nil {
		resp.Fail(c, http.StatusNotFound, "key not found")
		return
	}
	updates := map[string]any{}
	if req.Name != nil {
		updates["name"] = *req.Name
	}
	// 限额：>0 生效，0 或负数视为"不限制"（写 NULL）。
	// 不能存 0：token 配额判定是 `已用 >= 上限`，0 会让密钥第一次请求就被 429 打死。
	if req.QuotaTokensPerDay != nil {
		if *req.QuotaTokensPerDay > 0 {
			updates["quota_tokens_per_day"] = *req.QuotaTokensPerDay
		} else {
			updates["quota_tokens_per_day"] = nil
		}
	}
	if req.QuotaRequestsPerDay != nil {
		if *req.QuotaRequestsPerDay > 0 {
			updates["quota_requests_per_day"] = *req.QuotaRequestsPerDay
		} else {
			updates["quota_requests_per_day"] = nil
		}
	}
	if req.Status != nil {
		if *req.Status != store.KeyStatusActive && *req.Status != store.KeyStatusRevoked {
			resp.Fail(c, http.StatusBadRequest, "status must be active or revoked")
			return
		}
		updates["status"] = *req.Status
	}
	if len(updates) > 0 {
		if err := h.a.DB.Model(&k).Updates(updates).Error; err != nil {
			resp.Fail(c, http.StatusInternalServerError, "failed to update key")
			return
		}
	}
	// allowed_models 是 jsonb：手动序列化为 JSON 字符串写入（跨 SQLite/Postgres）
	// 显式传 [] 表示"清空白名单"（前端编辑框预填现值，字段内容即最终意图）
	if req.AllowedModels != nil {
		raw, _ := json.Marshal(*req.AllowedModels)
		if err := h.a.DB.Model(&k).Update("allowed_models", string(raw)).Error; err != nil {
			resp.Fail(c, http.StatusInternalServerError, "failed to update key")
			return
		}
	}
	// 重新读出，回显更新后的完整密钥（含被清零的限额）
	var updated store.APIKey
	if err := h.a.DB.First(&updated, k.ID).Error; err != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to reload key")
		return
	}
	resp.OK(c, keyJSON(updated))
}

// DeleteKey DELETE /api/v1/keys/:id — 删除（硬删，立即失效不可恢复）。
func (h *Handler) DeleteKey(c *gin.Context) {
	u, ok := userFromContext(c)
	if !ok {
		resp.Fail(c, http.StatusUnauthorized, "authentication required")
		return
	}
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		resp.Fail(c, http.StatusBadRequest, "invalid key id")
		return
	}
	res := h.a.DB.Where("id = ? AND user_id = ?", id, u.ID).Delete(&store.APIKey{})
	if res.Error != nil {
		resp.Fail(c, http.StatusInternalServerError, "failed to delete key")
		return
	}
	if res.RowsAffected == 0 {
		resp.Fail(c, http.StatusNotFound, "key not found")
		return
	}
	resp.OK(c, gin.H{"ok": true})
}