1 files changed, 207 insertions, 0 deletions
diff --git a/store/store.go b/store/store.go
new file mode 100644
index 0000000..fc42193
--- /dev/null
+++ b/store/store.go
@@ -0,0 +1,207 @@
+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
+
+// ((16^4)*4096)/1024/1024
+// If we ever need more than 256 MB in parameter store, we've done something
+// very wrong.
+const SSM_KEY_FORMAT = "%s-%04X" //
+
+var (
+	// TrimRegex is used to group the SSM keys inside of the Info struct under
+	// ByKey. This will only be used for params that exceed 4KB.
+	TrimRegex = regexp.MustCompile("-[0-9A-E][0-9A-E][0-9A-E][0-9A-E]$")
+
+	// Optional, can be set set to utilize a specific KMS key if desired.
+	KMS_KEY_ID *string = nil
+
+	Tags = []*ssm.Tag{}
+)
+
+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
+}
+
+type Entry struct {
+	Name string
+	Keys []*ssm.ParameterMetadata
+}
+
+// GetInfo returns a populated Info struct from the SSM
+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
+}
+
+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
+}
+
+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
+}
+
+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
+}