crypto-sm/sm4.go

package sm

import (
	"crypto/cipher"
	"errors"

	"apigo.cc/go/crypto"
	"apigo.cc/go/safe"
	"github.com/emmansun/gmsm/sm4"
)

type SM4Cipher struct {
	useGCM bool
}

var SM4CBC = &SM4Cipher{useGCM: false}
var SM4GCM = &SM4Cipher{useGCM: true}

func NewSM4CBC(safeKeyBuf, safeIvBuf *safe.SafeBuf) (*crypto.Symmetric, error) {
	return crypto.NewSymmetric(SM4CBC, safeKeyBuf, safeIvBuf)
}
func NewSM4CBCAndEraseKey(key, iv []byte) (*crypto.Symmetric, error) {
	return crypto.NewSymmetricAndEraseKey(SM4CBC, key, iv)
}
func NewSM4CBCWithoutEraseKey(key, iv []byte) (*crypto.Symmetric, error) {
	return crypto.NewSymmetricWithoutEraseKey(SM4CBC, key, iv)
}

func NewSM4GCM(safeKeyBuf, safeIvBuf *safe.SafeBuf) (*crypto.Symmetric, error) {
	return crypto.NewSymmetric(SM4GCM, safeKeyBuf, safeIvBuf)
}
func NewSM4GCMAndEraseKey(key, iv []byte) (*crypto.Symmetric, error) {
	return crypto.NewSymmetricAndEraseKey(SM4GCM, key, iv)
}
func NewSM4GCMWithoutEraseKey(key, iv []byte) (*crypto.Symmetric, error) {
	return crypto.NewSymmetricWithoutEraseKey(SM4GCM, key, iv)
}

func NewSM4CBCByPassword(password, salt []byte) (*crypto.Symmetric, error) {
	derived := crypto.DeriveKey(password, salt, 16+16)
	defer safe.ZeroMemory(derived)
	return NewSM4CBCAndEraseKey(derived[:16], derived[16:])
}

func NewSM4GCMByPassword(password, salt []byte) (*crypto.Symmetric, error) {
	derived := crypto.DeriveKey(password, salt, 16+12)
	defer safe.ZeroMemory(derived)
	return NewSM4GCMAndEraseKey(derived[:16], derived[16:])
}

func (s *SM4Cipher) Encrypt(data []byte, key []byte, iv []byte) ([]byte, error) {

	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if s.useGCM {
		sm4gcm, err := cipher.NewGCM(block)
		if err != nil {
			return nil, err
		}
		return sm4gcm.Seal(nil, iv[:sm4gcm.NonceSize()], data, nil), nil
	} else {
		blockSize := block.BlockSize()
		paddedData := crypto.Pkcs5Padding(data, blockSize)
		blockMode := cipher.NewCBCEncrypter(block, iv[:blockSize])
		crypted := make([]byte, len(paddedData))
		blockMode.CryptBlocks(crypted, paddedData)
		return crypted, nil
	}
}

func (s *SM4Cipher) Decrypt(data []byte, key []byte, iv []byte) ([]byte, error) {
	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if s.useGCM {
		sm4gcm, err := cipher.NewGCM(block)
		if err != nil {
			return nil, err
		}
		return sm4gcm.Open(nil, iv[:sm4gcm.NonceSize()], data, nil)
	} else {
		blockSize := block.BlockSize()
		if len(data)%blockSize != 0 {
			return nil, errors.New("ciphertext is not a multiple of block size")
		}
		blockMode := cipher.NewCBCDecrypter(block, iv[:blockSize])
		plainText := make([]byte, len(data))
		blockMode.CryptBlocks(plainText, data)
		unpadded := crypto.Pkcs5UnPadding(plainText)
		if unpadded == nil {
			return nil, errors.New("padding error")
		}
		return unpadded, nil
	}
}