path: root/store/store.go

// Store is not designed to be used as a database, or some high intensity
// key/value store, rather a low volume ad-hoc key value store for secrets
// inside of AWS.
//
// This should work out of the box with pretty much every AWS account.
// See also the bundled program.
package store

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ssm"
)

const SSM_MAX_SIZE = 4096 // This is dictated by AWS

const SSM_KEY_FORMAT = "%s-%04X" // Should be good up to 256MB ( 16^4 bytes... )

var (
	// TrimRegex is used to group the SSM keys inside of the Info struct under
	// ByKey. This is used so that we can have keys larger than 4KB.
	// See also SSM_KEY_FORMAT
	TrimRegex = regexp.MustCompile("-[0-9A-E][0-9A-E][0-9A-E][0-9A-E]$")

	// KMS_KEY_ID is optional, can be set set to utilize a specific KMS key if
	// desired.
	KMS_KEY_ID *string = nil

	// Tags are also optional, if used with the bundled program you can
	// simply set an environment variable. Otherwise, set them here
	// at the package level.
	Tags = []*ssm.Tag{}
)

// Info contains a few maps with pointers to all of the parameters, setup with
// different keys for easy lookup. `ByKey` is what you'd expect. `ByFullKey`
// has a dash and four hex digits appended to it for entries larger than 4K
// and actually reflects the keys you'll see in the parameter store console.
type Info struct {
	ByKey     map[string]*Entry
	ByFullKey map[string]*Entry
}

func (i *Info) init() {
	if i.ByKey == nil {
		i.ByKey = map[string]*Entry{}
	}
	if i.ByFullKey == nil {
		i.ByFullKey = map[string]*Entry{}
	}
}

func (i *Info) add(e *ssm.ParameterMetadata) {
	name := TrimRegex.ReplaceAllString(*e.Name, "")
	var entry *Entry

	// Doesn't exist, make a new entry
	if _, ok := i.ByKey[name]; !ok {
		entry = &Entry{
			Name: name,
			Keys: []*ssm.ParameterMetadata{e},
		}

		i.ByKey[name] = entry
		i.ByFullKey[*e.Name] = entry
		return
	}

	// Otherwise let's just update the one that's there
	entry = i.ByKey[name]
	entry.Keys = append(entry.Keys, e)
	i.ByFullKey[*e.Name] = entry
}

// Entry represents an entry in the store, and all of the actual parameters
// that it spans
type Entry struct {
	Name string
	Keys []*ssm.ParameterMetadata
}

// GetInfo returns a populated Info struct from the parameter store.
func GetInfo(svc *ssm.SSM) (*Info, error) {
	ret := &Info{
		ByKey:     map[string]*Entry{},
		ByFullKey: map[string]*Entry{},
	}

	var out = &ssm.DescribeParametersOutput{}
	var err error
	for {
		out, err = svc.DescribeParameters(&ssm.DescribeParametersInput{
			NextToken: out.NextToken,
		})
		if err != nil {
			return nil, err
		}
		for _, entry := range out.Parameters {
			ret.add(entry)
		}
		if out.NextToken == nil {
			break
		}
	}

	return ret, nil
}

// InsertParam will chuck data from the rdr into the parameter store under
// key, automatically chunking it into multiple parameters as needed.
func InsertParam(svc *ssm.SSM, rdr io.Reader, key string) error {
	buf := &bytes.Buffer{}
	enc := base64.NewEncoder(base64.StdEncoding, buf)

	_, err := io.Copy(enc, rdr)
	if err != nil {
		return fmt.Errorf("Error while reading stdin: %w", err)
	}

	err = enc.Close()
	if err != nil {
		return fmt.Errorf("While closing encoder: %w", err)
	}

	n := 1
	for {
		key := fmt.Sprintf(SSM_KEY_FORMAT, key, n)
		n++

		paramBuf := &bytes.Buffer{}

		n, err := io.CopyN(paramBuf, buf, SSM_MAX_SIZE)
		if err != io.EOF && err != nil {
			return fmt.Errorf("While writing output: %w", err)
		}
		// fmt.Fprintf(os.Stderr, "Wrote '%d' bytes to '%s'\n", n, key)

		_, err = svc.PutParameter(&ssm.PutParameterInput{
			KeyId: KMS_KEY_ID,
			Name:  aws.String(key),
			Type:  aws.String("SecureString"),
			Value: aws.String(paramBuf.String()),
			Tags:  Tags,
		})

		if err != nil {
			return fmt.Errorf("Failed putting prameter: %w", err)
		}

		if err == io.EOF || n < SSM_MAX_SIZE {
			break
		}
	}

	return nil
}

// GetParam will suck data out of parameter store for a key, automatically
// collecting all of the individual parameters needed to reconstruct the data
// and writes it out to the io.Writer
func GetParam(svc *ssm.SSM, wrtr io.Writer, key string) error {
	n := 1
	buf := &bytes.Buffer{}

	for {
		key := fmt.Sprintf(SSM_KEY_FORMAT, key, n)
		n++
		out, err := svc.GetParameter(&ssm.GetParameterInput{
			Name:           aws.String(key),
			WithDecryption: aws.Bool(true),
		})
		if err != nil {
			return fmt.Errorf("While fetching: '%s': %w", key, err)
		}

		w, err := buf.WriteString(*out.Parameter.Value)
		if err != nil {
			return fmt.Errorf("While writing to buffer: %w", err)
		}

		if w != SSM_MAX_SIZE {
			break
		}
	}

	dec := base64.NewDecoder(base64.StdEncoding, buf)
	_, err := io.Copy(wrtr, dec)
	if err != nil {
		return fmt.Errorf("Error writing output: %w", err)
	}

	return nil
}

// RemoveParam takes care of collecting all of the pieces for a given key,
// and removes all of them from the parameter store
func RemoveParam(svc *ssm.SSM, key string) error {
	info, err := GetInfo(svc)
	if err != nil {
		return fmt.Errorf("When fetching info: %w", err)
	}

	entry, ok := info.ByKey[key]
	if !ok {
		return fmt.Errorf("Entry '%s' not found in parameter store", key)
	}

	var e error
	for _, key := range entry.Keys {
		_, err := svc.DeleteParameter(&ssm.DeleteParameterInput{
			Name: key.Name,
		})
		if err != nil {
			if e == nil {
				e = err
			} else {
				e = fmt.Errorf("%s, %w", err, e)
			}
		}
	}

	return e
}