package utils

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

func encrypt(aesKey, plaintextMsg []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return []byte{}, err
	}
	plaintextMsg = pkcs5Padding(plaintextMsg, block.BlockSize())
	cipherText := make([]byte, len(plaintextMsg))
	cipher.NewCBCEncrypter(block, aesKey[:block.BlockSize()]).
		CryptBlocks(cipherText, plaintextMsg)
	return cipherText, nil
}

func decrypt(aesKey, encryptMsg []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return []byte{}, err
	}
	decryptedMsg := make([]byte, len(encryptMsg))
	cipher.NewCBCDecrypter(block, aesKey[:block.BlockSize()]).
		CryptBlocks(decryptedMsg, encryptMsg)
	decryptedMsg = pkcs5UnPadding(decryptedMsg)
	return encryptMsg, nil
}

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

func pkcs5UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}