ymt_v3_generate-20260721-10.../ymt_v3_generate/res.md

7.1 KiB
Raw Permalink Blame History

package ymt_v3_generate

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

// EncryptType 加密类型
type EncryptType string

const (
	EncryptTypeAES EncryptType = "aes" // AES ECB模式
	EncryptTypeSM4 EncryptType = "sm4" // SM4 CBC模式暂未实现仅占位
)

// Crypto 加密工具
type Crypto struct {
	key         []byte
	encryptType EncryptType
}

// NewCrypto 创建加密工具
func NewCrypto(key string, encryptType EncryptType) *Crypto {
	return &Crypto{
		key:         []byte(key),
		encryptType: encryptType,
	}
}

// Encrypt 加密明文返回base64编码的密文
func (c *Crypto) Encrypt(plaintext []byte) (string, error) {
	switch c.encryptType {
	case EncryptTypeAES:
		return c.aesECBEncrypt(plaintext)
	case EncryptTypeSM4:
		return "", errors.New("SM4 encryption not implemented")
	default:
		return "", fmt.Errorf("unsupported encrypt type: %s", c.encryptType)
	}
}

// Decrypt 解密密文base64编码返回明文
func (c *Crypto) Decrypt(ciphertext string) ([]byte, error) {
	cipherData, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return nil, fmt.Errorf("base64 decode failed: %w", err)
	}
	switch c.encryptType {
	case EncryptTypeAES:
		return c.aesECBDecrypt(cipherData)
	case EncryptTypeSM4:
		return nil, errors.New("SM4 decryption not implemented")
	default:
		return nil, fmt.Errorf("unsupported encrypt type: %s", c.encryptType)
	}
}

// aesECBEncrypt AES ECB模式加密
func (c *Crypto) aesECBEncrypt(plaintext []byte) (string, error) {
	block, err := aes.NewCipher(c.key)
	if err != nil {
		return "", err
	}
	// PKCS7填充
	plaintext = pkcs7Pad(plaintext, block.BlockSize())
	ciphertext := make([]byte, len(plaintext))
	// ECB模式直接分组加密
	for start := 0; start < len(plaintext); start += block.BlockSize() {
		block.Encrypt(ciphertext[start:start+block.BlockSize()], plaintext[start:start+block.BlockSize()])
	}
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// aesECBDecrypt AES ECB模式解密
func (c *Crypto) aesECBDecrypt(ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(c.key)
	if err != nil {
		return nil, err
	}
	if len(ciphertext)%block.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 += block.BlockSize() {
		block.Decrypt(plaintext[start:start+block.BlockSize()], ciphertext[start:start+block.BlockSize()])
	}
	// 去除PKCS7填充
	plaintext, err = pkcs7Unpad(plaintext, block.BlockSize())
	if err != nil {
		return nil, err
	}
	return plaintext, nil
}

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

// pkcs7Unpad 去除PKCS7填充
func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	if len(data) == 0 {
		return nil, errors.New("data is empty")
	}
	padding := int(data[len(data)-1])
	if padding > blockSize || padding == 0 {
		return nil, errors.New("invalid padding")
	}
	for i := len(data) - padding; i < len(data); i++ {
		if int(data[i]) != padding {
			return nil, errors.New("invalid padding")
		}
	}
	return data[:len(data)-padding], nil
}

// Signer 签名工具
type Signer struct {
	privateKey *rsa.PrivateKey
}

// NewSigner 创建签名工具
func NewSigner(privateKeyPEM string) (*Signer, error) {
	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return nil, errors.New("failed to decode private key PEM")
	}
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		// 尝试PKCS1
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse private key: %w", err)
		}
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("private key is not RSA")
	}
	return &Signer{privateKey: rsaKey}, nil
}

// Sign 对数据进行签名返回base64编码的签名
func (s *Signer) Sign(data string) (string, error) {
	hash := sha256.Sum256([]byte(data))
	signature, err := rsa.SignPKCS1v15(rand.Reader, s.privateKey, crypto.SHA256, hash[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

// Verifier 验签工具
type Verifier struct {
	publicKey *rsa.PublicKey
}

// NewVerifier 创建验签工具
func NewVerifier(publicKeyPEM string) (*Verifier, error) {
	block, _ := pem.Decode([]byte(publicKeyPEM))
	if block == nil {
		return nil, errors.New("failed to decode public key PEM")
	}
	key, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %w", err)
	}
	rsaKey, ok := key.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("public key is not RSA")
	}
	return &Verifier{publicKey: rsaKey}, nil
}

// Verify 验证签名
func (v *Verifier) Verify(data, signatureBase64 string) error {
	sig, err := base64.StdEncoding.DecodeString(signatureBase64)
	if err != nil {
		return fmt.Errorf("base64 decode signature failed: %w", err)
	}
	hash := sha256.Sum256([]byte(data))
	return rsa.VerifyPKCS1v15(v.publicKey, crypto.SHA256, hash[:], sig)
}

// BuildPlaintext 构建待加密/签名的明文将结构体转为JSON去掉零值字段按key排序
func BuildPlaintext(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 or pointer to struct")
	}
	// 转换为map[string]interface{},过滤零值
	m := make(map[string]interface{})
	typ := val.Type()
	for i := 0; i < val.NumField(); i++ {
		field := typ.Field(i)
		jsonTag := field.Tag.Get("json")
		if jsonTag == "" || jsonTag == "-" {
			continue
		}
		// 处理json tag可能包含omitempty
		name := strings.Split(jsonTag, ",")[0]
		fieldVal := val.Field(i)
		if isZeroValue(fieldVal) {
			continue
		}
		m[name] = fieldVal.Interface()
	}
	// 按key排序
	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 idx, k := range keys {
		if idx > 0 {
			buf.WriteByte(',')
		}
		buf.WriteString(fmt.Sprintf(`"%s":`, k))
		// 将值转为JSON
		valBytes, err := json.Marshal(m[k])
		if err != nil {
			return "", err
		}
		buf.Write(valBytes)
	}
	buf.WriteByte('}')
	return buf.String(), nil
}

// isZeroValue 判断反射值是否为零值
func isZeroValue(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.String:
		return v.String() == ""
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v.Int() == 0
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v.Uint() == 0
	case reflect.Float32, reflect.Float64:
		return v.Float() == 0
	case reflect.Bool:
		return !v.Bool()
	case reflect.Slice, reflect.Map, reflect.Ptr, reflect.Interface:
		return v.IsNil()
	default:
		return false
	}
}