// 提供AES-CBC-256加密解密功能。
package aes

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"

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

type IVGenerator func([32]byte) [16]byte

// 截取给定密钥的前16字节作为IV。
func PrefixIVGenerator(key [32]byte) [16]byte {
	return ([16]byte)(key[:16])
}

// 对给定的密钥进行异或运算获取IV。
func XorIVGenerator(key [32]byte) [16]byte {
	var iv [16]byte
	for i := 0; i < 16; i++ {
		iv[i] = key[i] ^ key[i+16]
	}
	return iv
}

// 对给定的密钥进行Sha256哈希运算获取用于加密算法的32字节密钥。
func generateKey(key []byte) [32]byte {
	return sha256.Sum256(key)
}

// 对给定的数据进行AES加密。
// 参数padding指定了填充方式，可选值为NoPadding、ZeroPadding和PKCS7Padding，默认采用ZeroPadding。
func Encrypt(data []byte, key []byte, padding encryption.PaddingMode, ivGenerator ...IVGenerator) ([]byte, error) {
	keyBytes := generateKey(key)
	block, err := aes.NewCipher(keyBytes[:])
	if err != nil {
		return nil, err
	}
	iv := append(ivGenerator, XorIVGenerator)[0](keyBytes)
	plainText := encryption.Padding(data, block.BlockSize(), padding)

	cipherText := make([]byte, len(plainText))
	mode := cipher.NewCBCEncrypter(block, iv[:])
	mode.CryptBlocks(cipherText, plainText)

	return cipherText, nil
}

// 对给定的数据进行AES解密。
// 参数padding指定了填充方式，可选值为NoPadding、ZeroPadding和PKCS7Padding，默认采用ZeroPadding。
func Decrypt(data []byte, key []byte, padding encryption.PaddingMode, ivGenerator ...IVGenerator) ([]byte, error) {
	keyBytes := generateKey(key)
	block, err := aes.NewCipher(keyBytes[:])
	if err != nil {
		return nil, err
	}
	iv := append(ivGenerator, XorIVGenerator)[0](keyBytes)

	plainText := make([]byte, len(data))
	mode := cipher.NewCBCDecrypter(block, iv[:])
	mode.CryptBlocks(plainText, data)

	return encryption.Unpadding(plainText, padding), nil
}