// 提供3DES-EDE加解密算法，本解密算法采用`(K1, K2, K3)`的安全长密钥。
package tdes

import (
	"crypto/cipher"
	"crypto/des"
	"crypto/sha256"

	"archgrid.xyz/ag/toolsbox/encryption"
)

type KeyGenerator func([]byte) [3][8]byte

// 直接采用分组的方式，将密钥分为4个部分，取前三个部分作为3DES的密钥。
func PartitionKeyGenerator(key []byte) [3][8]byte {
	hashedKey := sha256.Sum256(key)
	return [3][8]byte{
		([8]byte)(hashedKey[0:8]),
		([8]byte)(hashedKey[8:16]),
		([8]byte)(hashedKey[16:24]),
	}
}

// 对给定密钥进行Sha256散列以后，将散列结果分为三部分，分别作为3DES的密钥，前三组分别与第四组异或，得到3组新的密钥。
func XorKeyGenerator(key []byte) [3][8]byte {
	hashedKey := sha256.Sum256(key)
	var result [3][8]byte
	for i := 0; i < 3; i++ {
		for j := 0; j < 8; j++ {
			result[i][j] = hashedKey[i*8+j] ^ hashedKey[3*8+j]
		}
	}
	return result
}

// 对给定的数据使用给定的密钥进行加密。IV与密钥相同。
func rawDesEncrypt(data []byte, key [8]byte) ([]byte, error) {
	block, err := des.NewCipher(key[:])
	if err != nil {
		return nil, err
	}
	iv := key[:]
	cipherText := make([]byte, len(data))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(cipherText, data)
	return cipherText, nil
}

// 对给定的数据使用给定的密钥进行解密。IV与密钥相同。
func rawDesDecrypt(data []byte, key [8]byte) ([]byte, error) {
	block, err := des.NewCipher(key[:])
	if err != nil {
		return nil, err
	}
	iv := key[:]
	plainText := make([]byte, len(data))
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(plainText, data)
	return plainText, nil
}

// 对给定的数据进行加密。
// 3DES-EDE加密IV为8字节，默认与运算后的密钥相同。
// 参数padding指定了填充方式，可选值为NoPadding、ZeroPadding和PKCS7Padding，默认采用ZeroPadding。
func Encrypt(data []byte, key []byte, padding encryption.PaddingMode, keyGenerator ...KeyGenerator) ([]byte, error) {
	desKeys := append(keyGenerator, XorKeyGenerator)[0](key)
	plainText := encryption.Padding(data, des.BlockSize, padding)
	cipher1Text, err := rawDesEncrypt(plainText, desKeys[0])
	if err != nil {
		return nil, err
	}
	plain2Text, err := rawDesDecrypt(cipher1Text, desKeys[1])
	if err != nil {
		return nil, err
	}
	cipher3Text, err := rawDesEncrypt(plain2Text, desKeys[2])
	return cipher3Text, err
}

// 对给定的数据进行解密。
// 3DES-EDE解密IV为8字节，默认与运算后的密钥相同。
// 参数padding指定了填充方式，可选值为NoPadding、ZeroPadding和PKCS7Padding，默认采用ZeroPadding。
func Decrypt(data []byte, key []byte, padding encryption.PaddingMode, keyGenerator ...KeyGenerator) ([]byte, error) {
	desKeys := append(keyGenerator, XorKeyGenerator)[0](key)
	plainText, err := rawDesDecrypt(data, desKeys[2])
	if err != nil {
		return nil, err
	}
	cipher2Text, err := rawDesEncrypt(plainText, desKeys[1])
	if err != nil {
		return nil, err
	}
	plain3Text, err := rawDesDecrypt(cipher2Text, desKeys[0])
	return encryption.Unpadding(plain3Text, des.BlockSize, padding), err
}