package manager

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"

	"code.justin.tv/systems/sandstorm/logging"
	"code.justin.tv/systems/sandstorm/pkg/envelope"
	"code.justin.tv/systems/sandstorm/pkg/stat/statiface"
	"code.justin.tv/systems/sandstorm/resource"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/kms"
)

// An Envelope encrypts and decrypts secrets with single-use KMS data keys using
// AES-256-GCM.
type Envelope struct {
	clientLoader *clientLoader
	kmsKey       resource.KMSKey
	Logger       logging.Logger
	statter      statiface.API
}

// Seal generates a 256-bit data key using KMS and encrypts the given plaintext
// with AES-256-GCM using a random nonce. The ciphertext is appended to the
// nonce, which is in turn appended to the KMS data key ciphertext and returned.
func (e *Envelope) Seal(ctxt map[string]string, plaintext []byte) (encryptedDataKey envelope.EncryptedDataKey, ciphertext []byte, err error) {

	var dataKeyOutput *kms.GenerateDataKeyOutput
	if err = e.statter.WithMeasuredResult("GenerateDataKey", []*cloudwatch.Dimension{{
		Name:  aws.String("KMSKeyARN"),
		Value: aws.String(e.kmsKey.KeyARN),
	}}, func() error {
		dataKeyOutput, err = e.clientLoader.KMS(e.kmsKey.Region).GenerateDataKey(&kms.GenerateDataKeyInput{
			EncryptionContext: e.context(ctxt),
			KeySpec:           aws.String("AES_256"),
			KeyId:             aws.String(e.kmsKey.KeyARN),
		})
		return err
	}); err != nil {
		return
	}

	ciphertext, err = encrypt(dataKeyOutput.Plaintext, plaintext, []byte(aws.StringValue(dataKeyOutput.KeyId)))
	if err != nil {
		return
	}

	encryptedDataKey = envelope.EncryptedDataKey{
		Value:  dataKeyOutput.CiphertextBlob,
		Region: e.kmsKey.Region,
		KeyARN: aws.StringValue(dataKeyOutput.KeyId),
	}

	return
}

// Open takes the output of Seal and decrypts it. If any part of the ciphertext
// or context is modified, Seal will return an error instead of the decrypted
// data.
func (e *Envelope) Open(dataKey envelope.EncryptedDataKey, ciphertext []byte, ctxt map[string]string) (plaintext []byte, err error) {

	defer Zero(plaintext)

	var decryptOutput *kms.DecryptOutput

	if err = e.statter.WithMeasuredResult("Decrypt", []*cloudwatch.Dimension{{
		Name:  aws.String("KMSKeyARN"),
		Value: aws.String(dataKey.KeyARN),
	}}, func() (err error) {
		decryptOutput, err = e.clientLoader.KMS(dataKey.Region).Decrypt(&kms.DecryptInput{
			CiphertextBlob:    dataKey.Value,
			EncryptionContext: e.context(ctxt),
		})
		return
	}); err != nil {
		return
	}

	plaintext, err = decrypt(decryptOutput.Plaintext, ciphertext, []byte(dataKey.KeyARN))
	if err != nil {
		return
	}

	if apiErr, ok := err.(awserr.Error); ok {
		if apiErr.Code() == "InvalidCiphertextException" {
			err = fmt.Errorf("unable to decrypt data key")
		}
	}
	return
}

func (e *Envelope) context(c map[string]string) map[string]*string {
	ctxt := make(map[string]*string)
	for k, v := range c {
		ctxt[k] = aws.String(v)
	}
	return ctxt
}

func encrypt(key, plaintext, data []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

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

	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}

	return gcm.Seal(nonce, nonce, plaintext, data), nil
}

func decrypt(key, ciphertext, data []byte) ([]byte, error) {

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce, ciphertext := ciphertext[:gcm.NonceSize()], ciphertext[gcm.NonceSize():]

	return gcm.Open(nil, nonce, ciphertext, data)
}

// Zero replaces every byte in b with 0
func Zero(b []byte) {
	for i := range b {
		b[i] = 0
	}
}
