package go_encryption

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
)

func NewAES128EncryptionProvider() IEncryptionProvider {
	key, _ := hex.DecodeString(defaultEncryptionKey)
	return &aes128encryption{
		encryptionKey: key,
		_context:      defaultContext,
	}
}

type aes128encryption struct {
	encryptionKey []byte
	_context      []byte
}

// SetEncryptionKey allows to modify the encryption key.
func (e *aes128encryption) SetEncryptionKey(key []byte) {
	// Make sure we have a valid key length
	if keyLen := len(key); keyLen != 16 &&
		keyLen != 24 &&
		keyLen != 32 {
		if keyLen == 0 {
			key, _ = hex.DecodeString(defaultEncryptionKey)
		} else if keyLen < 32 {
			keyHash := sha256.Sum256(key)
			key = keyHash[:]
		} else {
			keyHash := sha512.Sum512_256(key)
			key = keyHash[:]
		}
	}
	e.encryptionKey = key
}

// SetContext allows to modify the encryption context.
// N.B.: The context must be between 12 and 16 bytes long.
func (e *aes128encryption) SetContext(context []byte) (err error) {
	length := len(context)
	if 12 > length || length > 16 {
		err = fmt.Errorf("encryption: unsupported context length")
	} else {
		e._context = e._context[:0]
		e._context = append(e._context, context...)
	}
	return
}

// EncryptData returns the given data as JSON in encrypted form.
func (e *aes128encryption) EncryptData(data interface{}) (encrypted []byte, err error) {
	var binary []byte
	if binary, err = json.Marshal(data); err == nil {
		encrypted, err = e.Encrypt(binary)
	}
	return
}

// EncryptToFile saves the given data to the given filepath as JSON in encrypted form.
func (e *aes128encryption) EncryptToFile(data interface{}, filepath string) (err error) {
	var encrypted []byte
	if encrypted, err = e.EncryptData(data); err == nil {
		err = os.WriteFile(filepath, encrypted, os.ModePerm)
	}
	return
}

// DecryptData tries to decrypt the given binary and to unmarshal the resulting JSON to the target interface.
func (e *aes128encryption) DecryptData(binary []byte, target interface{}) (err error) {
	var decrypted []byte
	if decrypted, err = e.Decrypt(binary); err == nil {
		err = json.Unmarshal(decrypted, target)
		for i := 0; i < len(decrypted); i++ {
			decrypted[i] = 0
		}
	}
	return
}

// DecryptData tries to decrypt the content of the file specified and to unmarshal the resulting JSON to the target interface.
func (e *aes128encryption) DecryptFile(path string, target interface{}) (err error) {
	var binary []byte
	if binary, err = os.ReadFile(path); err == nil {
		err = e.DecryptData(binary, target)
	}
	return
}

// getCrypter returns an AES cipher interface based on the current 16-, 24- or 32-byte key.
func (e *aes128encryption) getCrypter() (crypter cipher.AEAD, err error) {
	var key []byte
	key = append(key, e.encryptionKey...)
	var block cipher.Block
	if block, err = aes.NewCipher(key); err == nil {
		crypter, err = cipher.NewGCMWithNonceSize(block, len(e._context))
	}
	for i := 0; i < len(key); i++ {
		key[i] = 0
	}
	return
}

// getNonce returns a copy of the conext data
func (e *aes128encryption) getNonce() (nonce []byte) {
	nonce = append(nonce, e._context...)
	return
}

// Encrypt tries to encrypt the given binary with the appropriate AES cipher.
// N.B.: The data in the given binary will be nulled upon successful encryption.
func (e *aes128encryption) Encrypt(binary []byte) (encrypted []byte, err error) {
	var crypter cipher.AEAD
	if crypter, err = e.getCrypter(); err == nil {
		nonce := e.getNonce()
		encrypted = crypter.Seal(nonce, nonce, binary, nil)
		for i := 0; i < len(binary); i++ {
			binary[i] = 0
		}
	}
	return
}

// Decrypt tries to decrypt the given binary with the appropriate AES cipher.
func (e *aes128encryption) Decrypt(binary []byte) (decrypted []byte, err error) {
	var crypter cipher.AEAD
	if crypter, err = e.getCrypter(); err == nil {
		nonceSize := crypter.NonceSize()
		if len(binary) > nonceSize {
			context, encrypted := binary[:nonceSize], binary[nonceSize:]
			nonce := e.getNonce()
			if bytes.Equal(nonce, context) {
				decrypted, err = crypter.Open(nil, nonce, encrypted, nil)
			} else {
				err = fmt.Errorf("encryption: bad context")
			}
		} else {
			err = fmt.Errorf("encryption: not enough data")
		}
	}
	return
}