package util

import (
    "bytes"
    "crypto/aes"
    "crypto/cipher"
    "encoding/base64"
    "errors"
)

func aesDecrypt(cipherText []byte, key, iv []byte) ([]byte, error) {
    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, errors.New("illegal secret")
    }
    // 检验密文长度是否正确
    blockSize := block.BlockSize()

    if len(cipherText) < blockSize {
        return nil, errors.New("illegal cipherText: too small")
    }

    if len(cipherText)%blockSize != 0 {
        return nil, errors.New("illegal cipherText")
    }
    // 创建CBC解密模式
    blockModel := cipher.NewCBCDecrypter(block, iv)

    // 解密生成平文
    plainText := make([]byte, len(cipherText))
    blockModel.CryptBlocks(plainText, cipherText)

    // 去掉最后一个平文块的填充内容
    plainText = bytes.TrimSpace(plainText)
    return plainText, nil
}

func AesDecrypt(cipherTextWithBase64 []byte, key, iv []byte) ([]byte, error) {
    cipherText := make([]byte, base64.StdEncoding.DecodedLen(len(cipherTextWithBase64)))
    n, err := base64.StdEncoding.Decode(cipherText, cipherTextWithBase64)
    if err != nil {
        return nil, err
    }

    return aesDecrypt(cipherText[:n], key, iv)
}

// plainText: 明文
// key: 密钥
// iv： 初始化向量
func aesEncrypt(plainText []byte, key []byte, iv []byte) ([]byte, error) {

    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, errors.New("illegal secret")
    }

    plainText = spacePadding(plainText, block.BlockSize())
    // 创建CBC加密模式
    blockMode := cipher.NewCBCEncrypter(block, iv)
    cipherText := make([]byte, len(plainText))
    blockMode.CryptBlocks(cipherText, plainText)
    return cipherText, nil
}

func spacePadding(src []byte, blockSize int) []byte {
    padding := blockSize - len(src)%blockSize
    padText := bytes.Repeat([]byte(" "), padding)
    return append(src, padText...)
}

func AesEncrypt(plainText, key, iv []byte) ([]byte, error) {
    cipherText, err := aesEncrypt(plainText, key, iv)

    if err != nil {
        return nil, err
    }
    cipherTextBase64 := make([]byte, base64.StdEncoding.EncodedLen(len(cipherText)))
    base64.StdEncoding.Encode(cipherTextBase64, cipherText)
    return cipherTextBase64, nil
}