package vault

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/vault/physical"
)

const (
	// initialKeyTerm is the hard coded initial key term. This is
	// used only for values that are not encrypted with the keyring.
	initialKeyTerm = 1

	// termSize the number of bytes used for the key term.
	termSize = 4

	// aesgcmVersionByte is prefixed to a message to allow for
	// future versioning of barrier implementations.
	aesgcmVersionByte = 0x1
)

// barrierInit is the JSON encoded value stored
type barrierInit struct {
	Version int    // Version is the current format version
	Key     []byte // Key is the primary encryption key
}

// AESGCMBarrier is a SecurityBarrier implementation that uses the AES
// cipher core and the Galois Counter Mode block mode. It defaults to
// the golang NONCE default value of 12 and a key size of 256
// bit. AES-GCM is high performance, and provides both confidentiality
// and integrity.
type AESGCMBarrier struct {
	backend physical.Backend

	l      sync.RWMutex
	sealed bool

	// keyring is used to maintain all of the encryption keys, including
	// the active key used for encryption, but also prior keys to allow
	// decryption of keys encrypted under previous terms.
	keyring *Keyring

	// cache is used to reduce the number of AEAD constructions we do
	cache     map[uint32]cipher.AEAD
	cacheLock sync.RWMutex
}

// NewAESGCMBarrier is used to construct a new barrier that uses
// the provided physical backend for storage.
func NewAESGCMBarrier(physical physical.Backend) (*AESGCMBarrier, error) {
	b := &AESGCMBarrier{
		backend: physical,
		sealed:  true,
		cache:   make(map[uint32]cipher.AEAD),
	}
	return b, nil
}

// Initialized checks if the barrier has been initialized
// and has a master key set.
func (b *AESGCMBarrier) Initialized() (bool, error) {
	// Read the keyring file
	out, err := b.backend.Get(keyringPath)
	if err != nil {
		return false, fmt.Errorf("failed to check for initialization: %v", err)
	}
	if out != nil {
		return true, nil
	}

	// Fallback, check for the old sentinel file
	out, err = b.backend.Get(barrierInitPath)
	if err != nil {
		return false, fmt.Errorf("failed to check for initialization: %v", err)
	}
	return out != nil, nil
}

// Initialize works only if the barrier has not been initialized
// and makes use of the given master key.
func (b *AESGCMBarrier) Initialize(key []byte) error {
	// Verify the key size
	min, max := b.KeyLength()
	if len(key) < min || len(key) > max {
		return fmt.Errorf("Key size must be %d or %d", min, max)
	}

	// Check if already initialized
	if alreadyInit, err := b.Initialized(); err != nil {
		return err
	} else if alreadyInit {
		return ErrBarrierAlreadyInit
	}

	// Generate encryption key
	encrypt, err := b.GenerateKey()
	if err != nil {
		return fmt.Errorf("failed to generate encryption key: %v", err)
	}

	// Create a new keyring, install the keys
	keyring := NewKeyring()
	keyring = keyring.SetMasterKey(key)
	keyring, err = keyring.AddKey(&Key{
		Term:    1,
		Version: 1,
		Value:   encrypt,
	})
	if err != nil {
		return fmt.Errorf("failed to create keyring: %v", err)
	}
	return b.persistKeyring(keyring)
}

// persistKeyring is used to write out the keyring using the
// master key to encrypt it.
func (b *AESGCMBarrier) persistKeyring(keyring *Keyring) error {
	// Create the keyring entry
	buf, err := keyring.Serialize()
	if err != nil {
		return fmt.Errorf("failed to serialize keyring: %v", err)
	}
	defer memzero(buf)

	// Create the AES-GCM
	gcm, err := b.aeadFromKey(keyring.MasterKey())
	if err != nil {
		return err
	}

	// Encrypt the barrier init value
	value := b.encrypt(initialKeyTerm, gcm, buf)

	// Create the keyring physical entry
	pe := &physical.Entry{
		Key:   keyringPath,
		Value: value,
	}
	if err := b.backend.Put(pe); err != nil {
		return fmt.Errorf("failed to persist keyring: %v", err)
	}
	return nil
}

// GenerateKey is used to generate a new key
func (b *AESGCMBarrier) GenerateKey() ([]byte, error) {
	// Generate a 256bit key
	buf := make([]byte, 2*aes.BlockSize)
	_, err := rand.Read(buf)
	return buf, err
}

// KeyLength is used to sanity check a key
func (b *AESGCMBarrier) KeyLength() (int, int) {
	return aes.BlockSize, 2 * aes.BlockSize
}

// Sealed checks if the barrier has been unlocked yet. The Barrier
// is not expected to be able to perform any CRUD until it is unsealed.
func (b *AESGCMBarrier) Sealed() (bool, error) {
	b.l.RLock()
	defer b.l.RUnlock()
	return b.sealed, nil
}

// Unseal is used to provide the master key which permits the barrier
// to be unsealed. If the key is not correct, the barrier remains sealed.
func (b *AESGCMBarrier) Unseal(key []byte) error {
	b.l.Lock()
	defer b.l.Unlock()

	// Do nothing if already unsealed
	if !b.sealed {
		return nil
	}

	// Create the AES-GCM
	gcm, err := b.aeadFromKey(key)
	if err != nil {
		return err
	}

	// Read in the keyring
	out, err := b.backend.Get(keyringPath)
	if err != nil {
		return fmt.Errorf("failed to check for keyring: %v", err)
	}
	if out != nil {
		// Decrypt the barrier init key
		plain, err := b.decrypt(gcm, out.Value)
		if err != nil {
			if strings.Contains(err.Error(), "message authentication failed") {
				return ErrBarrierInvalidKey
			}
			return err
		}
		defer memzero(plain)

		// Recover the keyring
		keyring, err := DeserializeKeyring(plain)
		if err != nil {
			return fmt.Errorf("keyring deserialization failed: %v", err)
		}

		// Setup the keyring and finish
		b.keyring = keyring
		b.sealed = false
		return nil
	}

	// Read the barrier initialization key
	out, err = b.backend.Get(barrierInitPath)
	if err != nil {
		return fmt.Errorf("failed to check for initialization: %v", err)
	}
	if out == nil {
		return ErrBarrierNotInit
	}

	// Decrypt the barrier init key
	plain, err := b.decrypt(gcm, out.Value)
	if err != nil {
		if strings.Contains(err.Error(), "message authentication failed") {
			return ErrBarrierInvalidKey
		}
		return err
	}
	defer memzero(plain)

	// Unmarshal the barrier init
	var init barrierInit
	if err := json.Unmarshal(plain, &init); err != nil {
		return fmt.Errorf("failed to unmarshal barrier init file")
	}
	defer memzero(init.Key)

	// Setup a new keyring, this is for backwards compatability
	keyring := NewKeyring()
	keyring = keyring.SetMasterKey(key)
	keyring, err = keyring.AddKey(&Key{
		Term:    1,
		Version: 1,
		Value:   init.Key,
	})
	if err != nil {
		return fmt.Errorf("failed to create keyring: %v", err)
	}
	if err := b.persistKeyring(keyring); err != nil {
		return err
	}

	// Delete the old barrier entry
	if err := b.backend.Delete(barrierInitPath); err != nil {
		return fmt.Errorf("failed to delete barrier init file: %v", err)
	}

	// Set the vault as unsealed
	b.keyring = keyring
	b.sealed = false
	return nil
}

// Seal is used to re-seal the barrier. This requires the barrier to
// be unsealed again to perform any further operations.
func (b *AESGCMBarrier) Seal() error {
	b.l.Lock()
	defer b.l.Unlock()

	// Remove the primary key, and seal the vault
	b.cache = make(map[uint32]cipher.AEAD)
	b.keyring = nil
	b.sealed = true
	return nil
}

// Rotate is used to create a new encryption key. All future writes
// should use the new key, while old values should still be decryptable.
func (b *AESGCMBarrier) Rotate() error {
	b.l.Lock()
	defer b.l.Unlock()
	if b.sealed {
		return ErrBarrierSealed
	}

	// Generate a new key
	encrypt, err := b.GenerateKey()
	if err != nil {
		return fmt.Errorf("failed to generate encryption key: %v", err)
	}

	// Get the next term
	term := b.keyring.ActiveTerm()

	// Add a new encryption key
	newKeyring, err := b.keyring.AddKey(&Key{
		Term:    term + 1,
		Version: 1,
		Value:   encrypt,
	})
	if err != nil {
		return fmt.Errorf("failed to add new encryption key: %v", err)
	}

	// Persist the new keyring
	if err := b.persistKeyring(newKeyring); err != nil {
		return err
	}

	// Swap the keyrings
	b.keyring = newKeyring
	return nil
}

// Put is used to insert or update an entry
func (b *AESGCMBarrier) Put(entry *Entry) error {
	defer metrics.MeasureSince([]string{"barrier", "put"}, time.Now())
	b.l.RLock()
	defer b.l.RUnlock()
	if b.sealed {
		return ErrBarrierSealed
	}

	term := b.keyring.ActiveTerm()
	primary, err := b.aeadForTerm(term)
	if err != nil {
		return err
	}

	pe := &physical.Entry{
		Key:   entry.Key,
		Value: b.encrypt(term, primary, entry.Value),
	}
	return b.backend.Put(pe)
}

// Get is used to fetch an entry
func (b *AESGCMBarrier) Get(key string) (*Entry, error) {
	defer metrics.MeasureSince([]string{"barrier", "get"}, time.Now())
	b.l.RLock()
	defer b.l.RUnlock()
	if b.sealed {
		return nil, ErrBarrierSealed
	}

	// Read the key from the backend
	pe, err := b.backend.Get(key)
	if err != nil {
		return nil, err
	} else if pe == nil {
		return nil, nil
	}

	// Decrypt the ciphertext
	plain, err := b.decryptKeyring(pe.Value)
	if err != nil {
		return nil, fmt.Errorf("decryption failed: %v", err)
	}

	// Wrap in a logical entry
	entry := &Entry{
		Key:   key,
		Value: plain,
	}
	return entry, nil
}

// Delete is used to permanently delete an entry
func (b *AESGCMBarrier) Delete(key string) error {
	defer metrics.MeasureSince([]string{"barrier", "delete"}, time.Now())
	b.l.RLock()
	defer b.l.RUnlock()
	if b.sealed {
		return ErrBarrierSealed
	}

	return b.backend.Delete(key)
}

// List is used ot list all the keys under a given
// prefix, up to the next prefix.
func (b *AESGCMBarrier) List(prefix string) ([]string, error) {
	defer metrics.MeasureSince([]string{"barrier", "list"}, time.Now())
	b.l.RLock()
	defer b.l.RUnlock()
	if b.sealed {
		return nil, ErrBarrierSealed
	}

	return b.backend.List(prefix)
}

// aeadForTerm returns the AES-GCM AEAD for the given term
func (b *AESGCMBarrier) aeadForTerm(term uint32) (cipher.AEAD, error) {
	// Check for the keyring
	keyring := b.keyring
	if keyring == nil {
		return nil, nil
	}

	// Check the cache for the aead
	b.cacheLock.RLock()
	aead, ok := b.cache[term]
	b.cacheLock.RUnlock()
	if ok {
		return aead, nil
	}

	// Read the underlying key
	key := keyring.TermKey(term)
	if key == nil {
		return nil, nil
	}

	// Create a new aead
	aead, err := b.aeadFromKey(key.Value)
	if err != nil {
		return nil, err
	}

	// Update the cache
	b.cacheLock.Lock()
	b.cache[term] = aead
	b.cacheLock.Unlock()
	return aead, nil
}

// aeadFromKey returns an AES-GCM AEAD using the given key.
func (b *AESGCMBarrier) aeadFromKey(key []byte) (cipher.AEAD, error) {
	// Create the AES cipher
	aesCipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %v", err)
	}

	// Create the GCM mode AEAD
	gcm, err := cipher.NewGCM(aesCipher)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize GCM mode")
	}
	return gcm, nil
}

// encrypt is used to encrypt a value
func (b *AESGCMBarrier) encrypt(term uint32, gcm cipher.AEAD, plain []byte) []byte {
	// Allocate the output buffer with room for tern, version byte,
	// nonce, GCM tag and the plaintext
	capacity := termSize + 1 + gcm.NonceSize() + gcm.Overhead() + len(plain)
	size := termSize + 1 + gcm.NonceSize()
	out := make([]byte, size, capacity)

	// Set the key term
	binary.BigEndian.PutUint32(out[:4], term)

	// Set the version byte
	out[4] = aesgcmVersionByte

	// Generate a random nonce
	nonce := out[5 : 5+gcm.NonceSize()]
	rand.Read(nonce)

	// Seal the output
	out = gcm.Seal(out, nonce, plain, nil)
	return out
}

// decrypt is used to decrypt a value
func (b *AESGCMBarrier) decrypt(gcm cipher.AEAD, cipher []byte) ([]byte, error) {
	// Verify the term is always just one
	term := binary.BigEndian.Uint32(cipher[:4])
	if term != initialKeyTerm {
		return nil, fmt.Errorf("term mis-match")
	}

	// Verify the version byte
	if cipher[4] != aesgcmVersionByte {
		return nil, fmt.Errorf("version bytes mis-match")
	}

	// Capture the parts
	nonce := cipher[5 : 5+gcm.NonceSize()]
	raw := cipher[5+gcm.NonceSize():]
	out := make([]byte, 0, len(raw)-gcm.NonceSize())

	// Attempt to open
	return gcm.Open(out, nonce, raw, nil)
}

// decryptKeyring is used to decrypt a value using the keyring
func (b *AESGCMBarrier) decryptKeyring(cipher []byte) ([]byte, error) {
	// Verify the term
	term := binary.BigEndian.Uint32(cipher[:4])

	// Verify the version byte
	if cipher[4] != aesgcmVersionByte {
		return nil, fmt.Errorf("version bytes mis-match")
	}

	// Get the GCM by term
	gcm, err := b.aeadForTerm(term)
	if err != nil {
		return nil, err
	}
	if gcm == nil {
		return nil, fmt.Errorf("no decryption key available for term %d", term)
	}

	// Capture the parts
	nonce := cipher[5 : 5+gcm.NonceSize()]
	raw := cipher[5+gcm.NonceSize():]
	out := make([]byte, 0, len(raw)-gcm.NonceSize())

	// Attempt to open
	return gcm.Open(out, nonce, raw, nil)
}