package KFEncryption

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"io"
)

// Encrypt takes a plaintext byte array and encrypts it using the provided key.
// it returns the ciphertext and an error if one occurs.
// it generates the cipher text using AES-CGM and a nonce.
// a nonce is a number that can only be used once and will be unique for each encryption, but it is not a secret..
// it is safe to store the nonce along the cipher text.
// the nonce is prepended to the cipher text.
func Encrypt(plaintext []byte, key string) ([]byte, error) {
	block, err := aes.NewCipher(validateKey([]byte(key), 32))
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, 12)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	newCGM, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	cipherText := newCGM.Seal(nonce, nonce, plaintext, nil)
	return cipherText, nil
}

// Decrypt takes a cipher text byte array and decrypts it using the provided key
func Decrypt(cipherText []byte, key string) ([]byte, error) {
	block, err := aes.NewCipher(validateKey([]byte(key), 32))
	if err != nil {
		return nil, err
	}

	newCGM, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := newCGM.NonceSize()
	if len(cipherText) < nonceSize {
		return nil, err
	}

	nonce, cipherText := cipherText[:nonceSize], cipherText[nonceSize:]
	plainText, err := newCGM.Open(nil, nonce, cipherText, nil)
	if err != nil {
		return nil, err
	}

	return plainText, nil
}

// validateKey takes a key and a length and returns a key of the correct length by truncating or padding it
func validateKey(key []byte, length int) []byte {
	if len(key) == length {
		return key
	}

	newKey := make([]byte, length)
	copy(newKey, key)
	return newKey
}