package algorithm

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
)

// PKCS7 for aes
type PKCS7 struct {
}

func NewPKCS7() *PKCS7 {
	return &PKCS7{}
}

func (pk *PKCS7) padding(ciphertext []byte, blockSize int) []byte {
	return pkcs7Padding(ciphertext, blockSize)
}

func (pk *PKCS7) unPadding(data []byte) []byte {
	return pkcs7UnPadding(data)
}

// aes加密，填充秘钥key的16位，24,32分别对应AES-128, AES-192, or AES-256.
func (pk *PKCS7) aesCbcEncrypt(raw, key []byte, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	//填充原文
	blockSize := block.BlockSize()
	raw = pk.padding(raw, blockSize)

	//初始向量IV必须是唯一，但不需要保密
	var cipherText []byte
	if len(iv) < blockSize {
		//block大小 16
		cipherText = make([]byte, blockSize+len(raw))
		iv = cipherText[:blockSize]
		if _, er := io.ReadFull(rand.Reader, iv); er != nil {
			return nil, er
		}
	} else {
		cipherText = make([]byte, len(raw))
		blockSize = 0
	}
	//block大小和初始向量大小一定要一致
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(cipherText[blockSize:], raw)

	return cipherText, nil
}

func (pk *PKCS7) aesCbcDecrypt(encryptData, aesKey []byte, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}

	blockSize := block.BlockSize()

	if len(encryptData) < blockSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	if len(iv) == 0 {
		iv = encryptData[:blockSize]
		encryptData = encryptData[blockSize:]
	}
	// CBC mode always works in whole blocks.
	if len(encryptData)%blockSize != 0 {
		return nil, fmt.Errorf("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	// CryptBlocks can work in-place if the two arguments are the same.
	mode.CryptBlocks(encryptData, encryptData)

	//解填充
	return pk.unPadding(encryptData), nil
}

// Encrypt PKCS interface
func (pk *PKCS7) Encrypt(raw, key []byte) (string, error) {
	return pk.EncryptIv(raw, key, nil)
}

// EncryptIv
// 带IV
func (pk *PKCS7) EncryptIv(raw, aesKey []byte, iv []byte) (string, error) {
	data, err := pk.aesCbcEncrypt(raw, aesKey, iv)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

// Decrypt PKCS interface
func (pk *PKCS7) Decrypt(raw string, key []byte) ([]byte, error) {
	return pk.DecryptIv(raw, key, nil)
}

// DecryptIv
// 带IV
func (pk *PKCS7) DecryptIv(raw string, aesKey []byte, iv []byte) ([]byte, error) {
	data, err := base64.StdEncoding.DecodeString(raw)
	if err != nil {
		return nil, err
	}
	return pk.aesCbcDecrypt(data, aesKey, iv)
}