cpa-plugin/internal/plugin/key_management.go

package plugin

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"
	"unicode"

	managedaccess "cpa-ext/internal/access"
	"cpa-ext/internal/repository"
)

type createManagedKeyRequest struct {
	Name              string    `json:"name"`
	Secret            string    `json:"secret,omitempty"`
	RouteMode         string    `json:"route_mode,omitempty"`
	UpstreamAccountID string    `json:"upstream_account_id,omitempty"`
	AllModels         *bool     `json:"all_models,omitempty"`
	Models            *[]string `json:"models,omitempty"`
}

type updateManagedKeyRequest struct {
	ID                string    `json:"id"`
	Name              *string   `json:"name,omitempty"`
	Status            *string   `json:"status,omitempty"`
	RouteMode         *string   `json:"route_mode,omitempty"`
	UpstreamAccountID *string   `json:"upstream_account_id,omitempty"`
	AllModels         *bool     `json:"all_models,omitempty"`
	Models            *[]string `json:"models,omitempty"`
}

type archiveManagedKeyRequest struct {
	ID string `json:"id"`
}

type managedKeyStatsResponse struct {
	Stats  managedaccess.KeyStats `json:"stats"`
	Recent []usageListItem        `json:"recent"`
}

func (a *App) listManagedKeys(includeArchived bool) ManagementResponse {
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
	}
	keys, err := store.ListManagedKeys(context.Background(), includeArchived)
	if err != nil {
		return managementError(http.StatusInternalServerError, "database_error", err.Error())
	}
	return jsonManagementResponse(http.StatusOK, map[string]any{"keys": keys})
}

func (a *App) createManagedKey(body []byte) ManagementResponse {
	var req createManagedKeyRequest
	if err := json.Unmarshal(body, &req); err != nil {
		return managementError(http.StatusBadRequest, "invalid_request", "无法解析 Key 请求")
	}
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
	}
	name := strings.TrimSpace(req.Name)
	if name == "" || len(name) > 64 {
		return managementError(http.StatusBadRequest, "invalid_name", "Key 名称必须为 1-64 个字符")
	}
	secret := strings.TrimSpace(req.Secret)
	if secret == "" {
		var generateErr error
		secret, generateErr = generatedToken("cpa_")
		if generateErr != nil {
			return managementError(http.StatusInternalServerError, "random_unavailable", "无法生成 Key")
		}
	}
	if !validManagedSecret(secret) {
		return managementError(http.StatusBadRequest, "invalid_key", "Key 必须为 1-256 个不含空白或控制字符的字符")
	}
	keyID, err := generatedToken("key_")
	if err != nil {
		return managementError(http.StatusInternalServerError, "random_unavailable", "无法生成 Key ID")
	}
	key := managedaccess.ManagedKey{
		ID: keyID, Name: name, Secret: secret,
		Status: managedaccess.StatusActive, RouteMode: managedaccess.RouteAuto, AllModels: true,
	}
	a.mu.RLock()
	defaultName := a.config.BootstrapName
	a.mu.RUnlock()
	if template, err := store.ManagedKeyByName(context.Background(), defaultName); err == nil {
		key.RouteMode = template.RouteMode
		key.UpstreamAccountID = template.UpstreamAccountID
		key.AllModels = template.AllModels
		key.Models = append([]string(nil), template.Models...)
	}
	if req.RouteMode != "" {
		key.RouteMode = req.RouteMode
	}
	if req.UpstreamAccountID != "" {
		key.UpstreamAccountID = strings.TrimSpace(req.UpstreamAccountID)
	}
	if req.AllModels != nil {
		key.AllModels = *req.AllModels
	}
	if req.Models != nil {
		key.Models = normalizeModels(*req.Models)
	}
	if key.RouteMode == managedaccess.RouteAuto {
		key.UpstreamAccountID = ""
	}
	if key.AllModels {
		key.Models = nil
	}
	if err := validateManagedKeyRule(store, key); err != nil {
		return managementError(http.StatusBadRequest, "invalid_rule", err.Error())
	}
	if err := store.CreateManagedKey(context.Background(), key); err != nil {
		return managementError(http.StatusConflict, "key_conflict", err.Error())
	}
	created, _ := store.ManagedKeyByID(context.Background(), key.ID)
	return jsonManagementResponse(http.StatusCreated, created)
}

func (a *App) updateManagedKey(body []byte) ManagementResponse {
	var req updateManagedKeyRequest
	if err := json.Unmarshal(body, &req); err != nil || strings.TrimSpace(req.ID) == "" {
		return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
	}
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
	}
	key, err := store.ManagedKeyByID(context.Background(), req.ID)
	if err != nil {
		return managementError(http.StatusNotFound, "key_not_found", "Key 不存在")
	}
	if req.Name != nil {
		key.Name = strings.TrimSpace(*req.Name)
	}
	if req.Status != nil {
		key.Status = strings.TrimSpace(*req.Status)
	}
	if req.RouteMode != nil {
		key.RouteMode = strings.TrimSpace(*req.RouteMode)
	}
	if req.UpstreamAccountID != nil {
		key.UpstreamAccountID = strings.TrimSpace(*req.UpstreamAccountID)
	}
	if req.AllModels != nil {
		key.AllModels = *req.AllModels
	}
	if req.Models != nil {
		key.Models = normalizeModels(*req.Models)
	}
	if key.RouteMode == managedaccess.RouteAuto {
		key.UpstreamAccountID = ""
	}
	if key.AllModels {
		key.Models = nil
	}
	if key.Name == "" || len(key.Name) > 64 {
		return managementError(http.StatusBadRequest, "invalid_name", "Key 名称必须为 1-64 个字符")
	}
	if key.Status != managedaccess.StatusActive && key.Status != managedaccess.StatusDisabled {
		return managementError(http.StatusBadRequest, "invalid_status", "状态只能是 active 或 disabled")
	}
	if err := validateManagedKeyRule(store, key); err != nil {
		return managementError(http.StatusBadRequest, "invalid_rule", err.Error())
	}
	if err := store.UpdateManagedKey(context.Background(), key); err != nil {
		status := http.StatusConflict
		if errors.Is(err, repository.ErrManagedKeyNotFound) {
			status = http.StatusNotFound
		}
		return managementError(status, "update_failed", err.Error())
	}
	updated, _ := store.ManagedKeyByID(context.Background(), key.ID)
	return jsonManagementResponse(http.StatusOK, updated)
}

func (a *App) archiveManagedKey(body []byte) ManagementResponse {
	var req archiveManagedKeyRequest
	if err := json.Unmarshal(body, &req); err != nil || strings.TrimSpace(req.ID) == "" {
		return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
	}
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
	}
	if err := store.ArchiveManagedKey(context.Background(), req.ID); err != nil {
		return managementError(http.StatusNotFound, "key_not_found", "Key 不存在或已经归档")
	}
	return jsonManagementResponse(http.StatusOK, map[string]bool{"archived": true})
}

func (a *App) managedKeyStats(id string) ManagementResponse {
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
	}
	if _, err := store.ManagedKeyByID(context.Background(), id); err != nil {
		return managementError(http.StatusNotFound, "key_not_found", "Key 不存在")
	}
	location, _ := time.LoadLocation("Asia/Shanghai")
	now := time.Now().In(location)
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, location)
	stats, err := store.KeyStats(context.Background(), id, today)
	if err != nil {
		return managementError(http.StatusInternalServerError, "database_error", err.Error())
	}
	records, err := store.ListRecent(context.Background(), 1000)
	if err != nil {
		return managementError(http.StatusInternalServerError, "database_error", err.Error())
	}
	recent := make([]usageListItem, 0, 50)
	for _, record := range records {
		if record.ManagedKeyID == id {
			recent = append(recent, usageItem(record))
			if len(recent) == 50 {
				break
			}
		}
	}
	return jsonManagementResponse(http.StatusOK, managedKeyStatsResponse{Stats: stats, Recent: recent})
}

func (a *App) upstreamAccounts() ManagementResponse {
	syncErr := a.syncUpstreamAccounts(context.Background())
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "上游账号数据库尚未初始化")
	}
	if observedErr := a.flushObservedUpstreams(context.Background(), store); observedErr != nil && syncErr == nil {
		syncErr = observedErr
	}
	accounts, err := store.ListUpstreamAccounts(context.Background())
	if err != nil {
		return managementError(http.StatusInternalServerError, "database_error", err.Error())
	}
	response := map[string]any{"accounts": accounts}
	if syncErr != nil {
		response["warning"] = syncErr.Error()
	}
	return jsonManagementResponse(http.StatusOK, response)
}

func (a *App) modelSuggestions() ManagementResponse {
	store, ok := a.currentStore()
	if !ok {
		return managementError(http.StatusServiceUnavailable, "database_unavailable", "数据库尚未初始化")
	}
	models, err := store.ModelSuggestions(context.Background())
	if err != nil {
		return managementError(http.StatusInternalServerError, "database_error", err.Error())
	}
	return jsonManagementResponse(http.StatusOK, map[string]any{"models": models})
}

func (a *App) currentStore() (*repository.SQLiteUsageRepository, bool) {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.store, !a.closed && a.store != nil
}

func validateManagedKeyRule(store *repository.SQLiteUsageRepository, key managedaccess.ManagedKey) error {
	if key.RouteMode != managedaccess.RouteAuto && key.RouteMode != managedaccess.RouteStrict {
		return errors.New("路由模式只能是 auto 或 strict")
	}
	if key.RouteMode == managedaccess.RouteAuto {
		key.UpstreamAccountID = ""
	} else {
		if strings.TrimSpace(key.UpstreamAccountID) == "" {
			return errors.New("严格路由必须选择上游账号")
		}
		if _, err := store.UpstreamAccountByID(context.Background(), key.UpstreamAccountID); err != nil {
			return errors.New("选择的上游账号不存在")
		}
	}
	if !key.AllModels && len(normalizeModels(key.Models)) == 0 {
		return errors.New("未允许全部模型时至少配置一个模型")
	}
	return nil
}

func normalizeModels(models []string) []string {
	seen := make(map[string]struct{})
	result := make([]string, 0, len(models))
	for _, raw := range models {
		model := strings.ToLower(strings.TrimSpace(raw))
		if model == "" {
			continue
		}
		if _, exists := seen[model]; exists {
			continue
		}
		seen[model] = struct{}{}
		result = append(result, model)
	}
	return result
}

func generatedToken(prefix string) (string, error) {
	raw := make([]byte, 24)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	return prefix + base64.RawURLEncoding.EncodeToString(raw), nil
}

func validManagedSecret(secret string) bool {
	if secret == "" || len(secret) > 256 {
		return false
	}
	for _, value := range secret {
		if unicode.IsSpace(value) || unicode.IsControl(value) {
			return false
		}
	}
	return true
}

func managementError(status int, code, message string) ManagementResponse {
	return jsonManagementResponse(status, map[string]any{"error": map[string]string{"code": code, "message": message}})
}