package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"fmt"
)

func PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	m := blockSize - len(ciphertext)%blockSize
	n := bytes.Repeat([]byte{byte(m)}, m)
	return append(ciphertext, n...)
}

func PKCS7UnPadding(origData []byte) []byte {
	m := len(origData)
	n := int(origData[m-1])
	if m > n {
		return origData[:(m - n)]
	} else {
		return origData
	}
}

func AesEncrypt(buf, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	defer func() {
		if v := recover(); v != nil {
			err = fmt.Errorf("decrypt error %v", v)
		}
	}()
	blockSize := block.BlockSize()
	buf = PKCS7Padding(buf, blockSize)
	blockMode := cipher.NewCBCEncrypter(block, key[:blockSize])
	tmp := make([]byte, len(buf))
	blockMode.CryptBlocks(tmp, buf)
	return tmp, nil
}

func AesDecrypt(buf, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	defer func() {
		if v := recover(); v != nil {
			err = fmt.Errorf("decrypt error %v", v)
		}
	}()
	blockSize := block.BlockSize()
	blockMode := cipher.NewCBCDecrypter(block, key[:blockSize])
	origData := make([]byte, len(buf))
	blockMode.CryptBlocks(origData, buf)
	origData = PKCS7UnPadding(origData)
	return origData, nil
}