From e33c1a3690d84b4ac47b9783c54a9241f9b42ed8 Mon Sep 17 00:00:00 2001
From: Miles Maddox <j.milesmaddox@gmail.com>
Date: Mon, 27 Aug 2018 22:23:43 -0500
Subject: [PATCH] support encrypted strings

---
 lib/Diff.go            | 14 ++++++++++----
 lib/ParameterStates.go | 12 ++++++++----
 lib/commands_test.go   |  6 +++---
 lib/util.go            |  9 +++++++++
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/lib/Diff.go b/lib/Diff.go
index 174ceca..6608fcc 100644
--- a/lib/Diff.go
+++ b/lib/Diff.go
@@ -59,20 +59,26 @@ func (diff *Diff) AppendDeleteChange(path string) error {
 }
 
 // AppendAddChange appends a line to delete
-func (diff *Diff) AppendAddChange(path, desiredValue, currentValue string) error {
+func (diff *Diff) AppendAddChange(path, desiredValue, currentValue string, encryptionKey *string) error {
 	if currentValue != "" {
 		diff.asVisual = append(diff.asVisual, color.YellowString(fmt.Sprintf("~\t%s\t%s  -->  %s", path, currentValue, desiredValue)))
 	} else {
 		diff.asVisual = append(diff.asVisual, color.GreenString(fmt.Sprintf("+\t%s\t%s", path, desiredValue)))
 	}
 
-	diff.additions = append(diff.additions, &ssm.PutParameterInput{
-		// KeyId: ""
+	add := &ssm.PutParameterInput{
 		Name:      aws.String(path),
 		Overwrite: aws.Bool(true),
 		Type:      aws.String("String"),
 		Value:     aws.String(desiredValue),
-	})
+	}
+
+	if encryptionKey != nil {
+		add.KeyId = encryptionKey
+		add.Type = aws.String("SecureString")
+	}
+
+	diff.additions = append(diff.additions, add)
 
 	return nil
 }
diff --git a/lib/ParameterStates.go b/lib/ParameterStates.go
index ae38326..0252603 100644
--- a/lib/ParameterStates.go
+++ b/lib/ParameterStates.go
@@ -52,18 +52,22 @@ func (p *ParameterStates) diff(current ParameterStates) (diffs Diff, err error)
 
 	for path, ps := range *p {
 		for key, value := range ps.Parameters {
-
+			var encryptionKey *string
+			pathAndKey := fmt.Sprintf("%s/%s", path, key)
+			if stringInSlice(key, ps.EncryptedKeys) {
+				encryptionKey = ps.EncryptionKey
+			}
 			// Add because the path does not exist in current
 			if current[path] == nil {
-				diffs.AppendAddChange(fmt.Sprintf("%s/%s", path, key), value, "")
+				diffs.AppendAddChange(pathAndKey, value, "", encryptionKey)
 
 				// Add because the key does not exist in current
 			} else if current[path].Parameters[key] == "" {
-				diffs.AppendAddChange(fmt.Sprintf("%s/%s", path, key), value, "")
+				diffs.AppendAddChange(pathAndKey, value, "", encryptionKey)
 
 				// Add because the key is not up to date in current
 			} else if value != current[path].Parameters[key] {
-				diffs.AppendAddChange(fmt.Sprintf("%s/%s", path, key), value, current[path].Parameters[key])
+				diffs.AppendAddChange(pathAndKey, value, current[path].Parameters[key], encryptionKey)
 
 			}
 		}
diff --git a/lib/commands_test.go b/lib/commands_test.go
index e075d4a..4067829 100644
--- a/lib/commands_test.go
+++ b/lib/commands_test.go
@@ -5,7 +5,7 @@ import (
 )
 
 func TestRun(t *testing.T) {
-	CMDPull([]string{"/dev", "/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
-	// CMDPush([]string{"/dev", "/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
-	// CMDDiff([]string{"/dev", "/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
+	CMDPull([]string{"/ops", "/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
+	CMDPush([]string{"/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
+	CMDDiff([]string{"/ops"}, "yaml", "/home/justmiles/go/src/github.com/justmiles/ssm-parameter-store/scratch")
 }
diff --git a/lib/util.go b/lib/util.go
index 1d425eb..e0e03e1 100644
--- a/lib/util.go
+++ b/lib/util.go
@@ -112,3 +112,12 @@ func Check(err error) {
 		os.Exit(1)
 	}
 }
+
+func stringInSlice(a string, list []string) bool {
+	for _, b := range list {
+		if b == a {
+			return true
+		}
+	}
+	return false
+}