go-base/util/crypto/crypto_service.go

package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strings"
)

// CryptoService 加密服务结构体
type CryptoService struct {
	key    []byte
	prefix string // 加密文本前缀，用于识别是否已加密
}

// NewCryptoService 创建加密服务
func NewCryptoService(key string) (*CryptoService, error) {
	keyBytes := []byte(key)

	// 如果密钥长度不足，使用PBKDF2派生密钥
	if len(keyBytes) < 32 {
		derivedKey, err := deriveKey(keyBytes, 32)
		if err != nil {
			return nil, err
		}
		keyBytes = derivedKey
	} else if len(keyBytes) > 32 {
		keyBytes = keyBytes[:32] // 截取前32字节
	}

	return &CryptoService{
		key:    keyBytes,
		prefix: "ENC::", // 加密文本前缀
	}, nil
}

// deriveKey 使用PBKDF2派生密钥
func deriveKey(password []byte, keyLen int) ([]byte, error) {
	// 这里简化实现，实际应该使用crypto/sha256和更复杂的盐
	// 为了简单示例，使用固定盐（生产环境应该使用随机盐）
	salt := []byte("fixed-salt-for-demo")

	hash := make([]byte, keyLen)
	for i := 0; i < keyLen; i++ {
		if i < len(password) {
			hash[i] = password[i] ^ salt[i%len(salt)]
		} else {
			hash[i] = salt[i%len(salt)]
		}
	}
	return hash, nil
}

// Encrypt 加密文本
func (cs *CryptoService) Encrypt(plaintext string) (string, error) {
	// 如果已经是加密文本，直接返回
	if cs.IsEncrypted(plaintext) {
		return plaintext, nil
	}

	block, err := aes.NewCipher(cs.key)
	if err != nil {
		return "", fmt.Errorf("创建密码块失败: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("创建GCM模式失败: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", fmt.Errorf("生成随机数失败: %w", err)
	}

	ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil)
	encoded := base64.StdEncoding.EncodeToString(ciphertext)

	// 添加前缀标识
	return cs.prefix + encoded, nil
}

// Decrypt 解密文本
func (cs *CryptoService) Decrypt(input string) (string, error) {
	// 如果不是加密文本，直接返回原文本
	if !cs.IsEncrypted(input) {
		return input, nil
	}

	// 移除前缀
	encoded := strings.TrimPrefix(input, cs.prefix)

	ciphertext, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		// 如果不是有效的Base64，返回原文本
		return input, nil
	}

	block, err := aes.NewCipher(cs.key)
	if err != nil {
		return "", fmt.Errorf("创建密码块失败: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("创建GCM模式失败: %w", err)
	}

	if len(ciphertext) < gcm.NonceSize() {
		return "", errors.New("密文太短")
	}

	nonce, ciphertext := ciphertext[:gcm.NonceSize()], ciphertext[gcm.NonceSize():]

	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		// 如果解密失败，返回原文本（可能是格式相似但不是真正的加密文本）
		return input, nil
	}

	return string(plaintext), nil
}

// IsEncrypted 检查文本是否已加密
func (cs *CryptoService) IsEncrypted(text string) bool {
	return strings.HasPrefix(text, cs.prefix)
}

// EncryptOrDecrypt 智能加密或解密
// 如果输入是明文则加密，如果已是密文则解密
func (cs *CryptoService) EncryptOrDecrypt(input string) (string, error) {
	if cs.IsEncrypted(input) {
		return cs.Decrypt(input)
	}
	return cs.Encrypt(input)
}

// EncryptIfNeeded 如果需要则加密
func (cs *CryptoService) EncryptIfNeeded(input string) (string, error) {
	if cs.IsEncrypted(input) {
		return input, nil
	}
	return cs.Encrypt(input)
}

// DecryptIfNeeded 如果需要则解密
func (cs *CryptoService) DecryptIfNeeded(input string) (string, error) {
	if cs.IsEncrypted(input) {
		return cs.Decrypt(input)
	}
	return input, nil
}