cmautobuy/admin/service/ai_secret.go

package service

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"unicode/utf8"

	"github.com/goccy/go-yaml"
)

// AISecretStore 把 AI API Key 隔离在数据库、data 和发布目录之外。
type AISecretStore interface {
	Get(providerID string) (string, error)
	Set(providerID, apiKey string) error
	Clear(providerID string) error
	Status(providerID string) (configured bool, suffix string, err error)
}

type aiSecretFile struct {
	Version   int               `yaml:"version"`
	Providers map[string]string `yaml:"providers"`
}

// FileAISecretStore 使用同目录临时文件 + fsync + 原子替换保存密钥。
type FileAISecretStore struct {
	path string
	mu   sync.Mutex
}

func NewFileAISecretStore(path string) *FileAISecretStore {
	return &FileAISecretStore{path: strings.TrimSpace(path)}
}

func (s *FileAISecretStore) Get(providerID string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	data, err := s.readLocked()
	if err != nil {
		return "", err
	}
	return data.Providers[strings.TrimSpace(providerID)], nil
}

func (s *FileAISecretStore) Status(providerID string) (bool, string, error) {
	secret, err := s.Get(providerID)
	if err != nil {
		return false, "", err
	}
	if secret == "" {
		return false, "", nil
	}
	runes := []rune(secret)
	start := len(runes) - 4
	if start < 0 {
		start = 0
	}
	return true, "••••" + string(runes[start:]), nil
}

func (s *FileAISecretStore) Set(providerID, apiKey string) error {
	providerID = strings.TrimSpace(providerID)
	apiKey = strings.TrimSpace(apiKey)
	if providerID == "" {
		return fmt.Errorf("服务商编号不能为空")
	}
	if utf8.RuneCountInString(apiKey) < 8 {
		return fmt.Errorf("API Key 至少需要 8 个字符")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	data, err := s.readLocked()
	if err != nil {
		return err
	}
	data.Providers[providerID] = apiKey
	return s.writeLocked(data)
}

func (s *FileAISecretStore) Clear(providerID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	data, err := s.readLocked()
	if err != nil {
		return err
	}
	delete(data.Providers, strings.TrimSpace(providerID))
	return s.writeLocked(data)
}

func (s *FileAISecretStore) readLocked() (aiSecretFile, error) {
	result := aiSecretFile{Version: 1, Providers: map[string]string{}}
	path, err := s.safePath()
	if err != nil {
		return result, err
	}
	raw, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return result, nil
	}
	if err != nil {
		return result, fmt.Errorf("读取 AI 密钥文件失败")
	}
	if err := yaml.Unmarshal(raw, &result); err != nil {
		return aiSecretFile{}, fmt.Errorf("解析 AI 密钥文件失败")
	}
	if result.Providers == nil {
		result.Providers = map[string]string{}
	}
	return result, nil
}

func (s *FileAISecretStore) writeLocked(data aiSecretFile) error {
	path, err := s.safePath()
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return fmt.Errorf("准备 AI 密钥目录失败")
	}
	raw, err := yaml.Marshal(data)
	if err != nil {
		return fmt.Errorf("编码 AI 密钥文件失败")
	}
	tmp, err := os.CreateTemp(filepath.Dir(path), ".ai-secrets-*.tmp")
	if err != nil {
		return fmt.Errorf("创建 AI 密钥临时文件失败")
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)
	if err := tmp.Chmod(0o600); err != nil {
		tmp.Close()
		return fmt.Errorf("设置 AI 密钥临时文件权限失败")
	}
	if _, err := tmp.Write(raw); err != nil {
		tmp.Close()
		return fmt.Errorf("写入 AI 密钥临时文件失败")
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return fmt.Errorf("同步 AI 密钥临时文件失败")
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("关闭 AI 密钥临时文件失败")
	}
	if err := os.Rename(tmpName, path); err != nil {
		return fmt.Errorf("替换 AI 密钥文件失败")
	}
	if err := os.Chmod(path, 0o600); err != nil {
		return fmt.Errorf("设置 AI 密钥文件权限失败")
	}
	if dir, err := os.Open(filepath.Dir(path)); err == nil {
		_ = dir.Sync()
		_ = dir.Close()
	}
	return nil
}

func (s *FileAISecretStore) safePath() (string, error) {
	if s == nil || s.path == "" {
		return "", fmt.Errorf("未配置 CMAUTOBUY_AI_SECRETS_PATH，不能保存或读取 AI 密钥")
	}
	if !filepath.IsAbs(s.path) {
		return "", fmt.Errorf("AI 密钥文件必须使用绝对路径")
	}
	abs, err := filepath.Abs(s.path)
	if err != nil {
		return "", fmt.Errorf("解析 AI 密钥文件路径失败")
	}
	for _, root := range protectedAISecretRoots() {
		if pathWithin(abs, root) {
			return "", fmt.Errorf("AI 密钥文件不能位于仓库、data 或发布目录内")
		}
	}
	return abs, nil
}

func protectedAISecretRoots() []string {
	var roots []string
	if cwd, err := os.Getwd(); err == nil {
		roots = append(roots, cwd)
	}
	if exe, err := os.Executable(); err == nil {
		roots = append(roots, filepath.Dir(exe))
	}
	return roots
}

func pathWithin(path, root string) bool {
	path, root = filepath.Clean(path), filepath.Clean(root)
	rel, err := filepath.Rel(root, path)
	return err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
}