ymt_v3_generate-20260721-101148/ymt_v3_generate/crypto.go

package ymt_v3_generate

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"sort"
	"strings"
)

// PKCS7Padding 填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

// pkcs7UnPadding 去除填充
func pkcs7UnPadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("data is empty")
	}
	unpadding := int(data[length-1])
	if unpadding > length || unpadding == 0 {
		return nil, errors.New("invalid padding")
	}
	return data[:length-unpadding], nil
}

// aesECBEncrypt AES ECB 加密
func aesECBEncrypt(plaintext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	plaintext = pkcs7Padding(plaintext, blockSize)
	ciphertext := make([]byte, len(plaintext))
	// ECB 模式：逐块加密
	for start := 0; start < len(plaintext); start += blockSize {
		block.Encrypt(ciphertext[start:start+blockSize], plaintext[start:start+blockSize])
	}
	return ciphertext, nil
}

// aesECBDecrypt AES ECB 解密
func aesECBDecrypt(ciphertext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	if len(ciphertext)%blockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of block size")
	}
	plaintext := make([]byte, len(ciphertext))
	for start := 0; start < len(ciphertext); start += blockSize {
		block.Decrypt(plaintext[start:start+blockSize], ciphertext[start:start+blockSize])
	}
	return pkcs7UnPadding(plaintext)
}

// EncryptPlaintext 加密明文（AES ECB），返回 base64 字符串
func EncryptPlaintext(plaintext string, key []byte) (string, error) {
	ciphertext, err := aesECBEncrypt([]byte(plaintext), key)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// DecryptCiphertext 解密密文（base64 输入），返回明文字符串
func DecryptCiphertext(ciphertextBase64 string, key []byte) (string, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return "", err
	}
	plaintext, err := aesECBDecrypt(ciphertext, key)
	if err != nil {
		return "", err
	}
	return string(plaintext), nil
}

// Sign 生成签名：appid + timestamp + ciphertext，使用 RSA 私钥签名，返回 base64 签名
func Sign(appID, timestamp, ciphertext string, privateKey *rsa.PrivateKey) (string, error) {
	data := appID + timestamp + ciphertext
	hash := sha256.Sum256([]byte(data))
	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hash[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

// VerifySign 验证签名
func VerifySign(appID, timestamp, ciphertext, signBase64 string, publicKey *rsa.PublicKey) error {
	data := appID + timestamp + ciphertext
	hash := sha256.Sum256([]byte(data))
	signature, err := base64.StdEncoding.DecodeString(signBase64)
	if err != nil {
		return err
	}
	return rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], signature)
}

// ParsePrivateKey 解析 PEM 格式的 RSA 私钥
func ParsePrivateKey(pemStr string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, errors.New("failed to parse PEM block containing private key")
	}
	// 尝试 PKCS1
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err == nil {
		return key, nil
	}
	// 尝试 PKCS8
	key8, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err == nil {
		if rsaKey, ok := key8.(*rsa.PrivateKey); ok {
			return rsaKey, nil
		}
		return nil, errors.New("parsed key is not RSA")
	}
	return nil, fmt.Errorf("failed to parse private key: %v", err)
}

// ParsePublicKey 解析 PEM 格式的 RSA 公钥
func ParsePublicKey(pemStr string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, errors.New("failed to parse PEM block containing public key")
	}
	// 尝试 PKIX
	key, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err == nil {
		if rsaKey, ok := key.(*rsa.PublicKey); ok {
			return rsaKey, nil
		}
		return nil, errors.New("parsed key is not RSA")
	}
	// 尝试 PKCS1
	rsaKey, err := x509.ParsePKCS1PublicKey(block.Bytes)
	if err == nil {
		return rsaKey, nil
	}
	return nil, fmt.Errorf("failed to parse public key: %v", err)
}

// RemoveZeroValuesAndSort 将结构体转为 map，过滤零值，按键排序，返回 JSON 字符串
func RemoveZeroValuesAndSort(v interface{}) (string, error) {
	// 使用反射获取字段
	val := reflect.ValueOf(v)
	if val.Kind() == reflect.Ptr {
		val = val.Elem()
	}
	if val.Kind() != reflect.Struct {
		return "", errors.New("input must be a struct")
	}
	typ := val.Type()
	m := make(map[string]interface{})
	for i := 0; i < val.NumField(); i++ {
		field := typ.Field(i)
		jsonTag := field.Tag.Get("json")
		if jsonTag == "" || jsonTag == "-" {
			continue
		}
		// 处理 omitempty
		name := strings.Split(jsonTag, ",")[0]
		fieldVal := val.Field(i)
		// 检查零值
		if fieldVal.IsZero() {
			continue
		}
		m[name] = fieldVal.Interface()
	}
	// 按键排序
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	// 构建 JSON
	var buf bytes.Buffer
	buf.WriteByte('{')
	for i, k := range keys {
		if i > 0 {
			buf.WriteByte(',')
		}
		buf.WriteString(`"` + k + `":`)
		// 简单序列化（仅支持基本类型）
		v := m[k]
		switch val := v.(type) {
		case string:
			buf.WriteString(`"` + val + `"`)
		case int, int32, int64, uint, uint32, uint64:
			fmt.Fprintf(&buf, "%d", val)
		case float64:
			fmt.Fprintf(&buf, "%v", val)
		case bool:
			fmt.Fprintf(&buf, "%t", val)
		default:
			// 使用 json.Marshal 处理复杂类型
			b, err := json.Marshal(v)
			if err != nil {
				return "", err
			}
			buf.Write(b)
		}
	}
	buf.WriteByte('}')
	return buf.String(), nil
}