package crypto

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

// AES aes
type AES struct {
	key []byte
	iv  []byte
}

var (
	defaultAes AES
)

// NewAES new
func NewAES(key, iv string) *AES {
	return &AES{
		key: []byte(key),
		iv:  []byte(iv),
	}
}

// SetKeyIV set key & iv
func (c *AES) SetKeyIV(key, iv string) {
	c.key = []byte(key)
	c.iv = []byte(iv)
}

// Encrypt encrypt data
func (c *AES) Encrypt(bs []byte) (crypted []byte, err error) {
	block, err := aes.NewCipher(c.key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	bs = pkcs7Padding(bs, blockSize)

	blockMode := cipher.NewCBCEncrypter(block, c.iv[:blockSize])
	crypted = make([]byte, len(bs))

	blockMode.CryptBlocks(crypted, bs)
	return
}

// Decrypt decrypt data
func (c *AES) Decrypt(bs []byte) (ds []byte, err error) {
	block, err := aes.NewCipher(c.key)
	if err != nil {
		return nil, err
	}
	n := len(bs)
	blockSize := block.BlockSize()
	if n%blockSize != 0 {
		return nil, errors.New(`crypto/cipher: input not full blocks`)
	}

	blockMode := cipher.NewCBCDecrypter(block, c.iv[:blockSize])
	ds = make([]byte, n)
	blockMode.CryptBlocks(ds, bs)
	ds = pkcs7UnPadding(ds)
	return
}

func pkcs7Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padtext...)
}

func pkcs7UnPadding(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	if unpadding >= length {
		return []byte(``)
	}
	return data[:(length - unpadding)]
}

// SetAESDefaultKeyIV set default key & iv
func SetAESDefaultKeyIV(key, iv string) {
	defaultAes = AES{
		key: []byte(key),
		iv:  []byte(iv),
	}
}

// AesEncrypt default key & iv encrypt
func AesEncrypt(bs []byte) ([]byte, error) {
	return defaultAes.Encrypt(bs)
}

// AesDecrypt default key & iv decrypt
func AesDecrypt(bs []byte) ([]byte, error) {
	return defaultAes.Decrypt(bs)
}