package main

import (
	"bytes"
	"crypto/rand"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/gen/kms"
	"golang.org/x/crypto/nacl/secretbox"
)

const (
	region = "eu-west-1"          // <-- change if needed
	keyId  = "INSERT-KEY-ID-HERE" // <-- insert key id here
)

func main() {
	plaintext, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		panic(err)
	}
	fmt.Println("Plaintext: ", plaintext)

	// Create client
	creds, err := aws.ProfileCreds("", "", 10*time.Minute)
	if err != nil {
		panic(err)
	}
	kmsClient := kms.New(creds, region, nil)

	// Encrypt plaintext
	encrypted, err := Encrypt(kmsClient, plaintext) // <-- to be implemented
	if err != nil {
		panic(err)
	}
	fmt.Println("Encrypted: ", encrypted)

	// Decrypt ciphertext
	decrypted, err := Decrypt(kmsClient, encrypted) // <-- to be implemented
	if err != nil {
		panic(err)
	}
	fmt.Println("Decrypted: ", decrypted)
}

const (
	keyLength   = 32
	nonceLength = 24
)

type payload struct {
	Key     []byte
	Nonce   *[nonceLength]byte
	Message []byte
}

func Encrypt(kmsClient *kms.KMS, plaintext []byte) ([]byte, error) {
	// Generate data key
	rsp, err := kmsClient.GenerateDataKey(&kms.GenerateDataKeyRequest{
		KeyID:         aws.String(keyId),
		NumberOfBytes: aws.Integer(keyLength),
	})
	if err != nil {
		return nil, err
	}

	// Initialize payload
	p := &payload{
		Key:   rsp.CiphertextBlob,
		Nonce: &[nonceLength]byte{},
	}

	// Set nonce
	if _, err = rand.Read(p.Nonce[:]); err != nil {
		return nil, err
	}

	// Create key
	key := &[keyLength]byte{}
	copy(key[:], rsp.Plaintext)

	// Encrypt message
	p.Message = secretbox.Seal(p.Message, plaintext, p.Nonce, key)

	buf := &bytes.Buffer{}
	if err := gob.NewEncoder(buf).Encode(p); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func Decrypt(kmsClient *kms.KMS, ciphertext []byte) ([]byte, error) {
	// Decode ciphertext with gob
	var p payload
	gob.NewDecoder(bytes.NewReader(ciphertext)).Decode(&p)

	// Decrypt key
	decryptRsp, err := kmsClient.Decrypt(&kms.DecryptRequest{
		CiphertextBlob: p.Key,
	})
	if err != nil {
		return nil, err
	}
	key := &[keyLength]byte{}
	copy(key[:], decryptRsp.Plaintext)

	// Decrypt message
	var plaintext []byte
	plaintext, ok := secretbox.Open(plaintext, p.Message, p.Nonce, key)
	if !ok {
		return nil, fmt.Errorf("Failed to open secretbox")
	}
	return plaintext, nil
}