///
/// Copyright (c) 2018 xdx. All rights reserved.
///
/// \file: sm4.go
///
/// \brief: SM4
///
/// \author: xdx
///

package sm4

// 这里都是SM4小包解密，大包的实现可以参照这里的多次调用CryptBlocks即可。
// 输入的dst可以是较大空间，返回dst的切片将其长度改变
// dstBuf = make([]byte, length, enoughCap)
// dst = ECBEncrypt(dstBuf[:0], key, src)
// now len(dst) = len(src)
// 输入dst是为了不在加密函数中分配空间。调用者可以预先分配好重复使用。
// SM4的相关函数如果确保输入数据长度都是正确的，那么可以不检查返回error

// TODO 大包加密 参照标准库实现。可以节省这里每次调用的轮密钥生成开销

// 如果不知道dst应该输入什么，那就输入nil
// 如果不知道dst应该输入什么，那就输入nil
// 如果不知道dst应该输入什么，那就输入nil

import (
	"crypto/cipher"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/internal/modes"
	"xdx.jelly/xgcl/internal/subtle"
	"xdx.jelly/xgcl/internal/tags"
	"xdx.jelly/xgcl/utils"
)

// 取错名字了，用下面的名字
var (
	EncryptECB = ECBEncrypt
	DecryptECB = ECBDecrypt

	EncryptCBC = CBCEncrypt
	DecryptCBC = CBCDecrypt

	EncryptOFB = OFBEncrypt
	DecryptOFB = OFBDecrypt

	EncryptCFB = CFBEncrypt
	DecryptCFB = CFBDecrypt
)

// export ECB BlockMode
var NewECBEncrypter = modes.NewECBEncrypter
var NewECBDecrypter = modes.NewECBDecrypter

// ECBEncrypt len(encData) must exactly equal to len(data)
// or return error
// we not realloc memory for encData!!!
// data和encData要么完全重叠，要么完全不重叠！！！
//
// @param dst[out]	密文预分配空间，可以为nil
// @param key[in]	16字节密钥
// @param src[in]	明文,长度必须是16的倍数，函数不负责padding
// @return []byte	输入密文。
// @note 如果输入dst的空间（cap）比src小，则可能引起内存地址得变化。
// @usage dst = ECBEncrypt(dst, key, src)
func ECBEncrypt(dst, key, src []byte) ([]byte, error) {

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input size must be a multiple of %d", BlockSize)
	}
	// dst = nil, cap(dst) = 0
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	if subtle.InexactOverlap(dst, src) {
		dst = make([]byte, len(src))
	}
	if len(src) == 0 {
		return dst[:0], nil
	}

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

	blockMode := modes.NewECBEncrypter(block)
	blockMode.CryptBlocks(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

// ECBDecrypt 解密
// params:
//
//	dst[in]		明文
//	src[in]		密文，长度必须是16的倍数，函数不负责解padding
//	key[in]		16字节密钥
func ECBDecrypt(dst, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}
	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := modes.NewECBDecrypter(block)
	blockMode.CryptBlocks(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

// CBCEncrypt cbc加密
// 函数返回时iv保存最后一个分组密文，因此对大包数据可以分组多次调用进行加密。
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节，返回时是最后一个分组密文
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func CBCEncrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCBCEncrypter(block, iv)
	blockMode.CryptBlocks(dst, src)

	// The last dst block is the return iv
	if len(dst) > 0 {
		copy(iv, dst[len(dst)-BlockSize:])
	}
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil

}

// CBCDecrypt cbc解密
// 函数返回时iv保存最后一个分组密文，因此对大包数据可以分组多次调用进行加密。
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节，返回时是最后一个分组
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func CBCDecrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCBCDecrypter(block, iv)
	blockMode.CryptBlocks(dst, src)

	// The last src block is the return iv
	if len(src) > 0 {
		copy(iv, src[len(src)-block.BlockSize():])
	}
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

/////////////////////////////////////////////////////
// CFB and OFB, not test carefully.

// CFBEncrypt cfb加密
// params:
//
//	dst[out]	密文，输入时分配好长度，必须和data的长度相等
//	iv[in/out]	初始向量, 16字节
//	key[in]		16字节密钥
//	src[in]		明文,长度必须是16的倍数，函数不负责padding
func CFBEncrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCFBEncrypter(block, iv)
	blockMode.XORKeyStream(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

// CFBDecrypt ...
func CFBDecrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewCFBDecrypter(block, iv)
	blockMode.XORKeyStream(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

// OFBEncrypt ofb加密
func OFBEncrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4EncFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewOFB(block, iv)
	blockMode.XORKeyStream(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}

// OFBDecrypt ofb解密
func OFBDecrypt(dst, iv, key, src []byte) ([]byte, error) {
	block, err := NewCipher(key)
	if err != nil {
		return dst, err
	}

	if len(iv) != BlockSize {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input IV size must be a multiple of %d", BlockSize)
	}

	if len(src)%BlockSize != 0 {
		return dst, gerrors.WithAnnotatingf(ErrSM4DecFailed, "input size must be a multiple of %d", BlockSize)
	}

	if dst == nil || cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[:len(src)]

	blockMode := cipher.NewOFB(block, iv)
	blockMode.XORKeyStream(dst, src)
	if tags.Gmt0028 {
		if c, ok := block.(utils.Clearable); ok {
			c.Clear()
		}
	}
	return dst, nil
}