Add AWS KMS decryption support for environment variables

CHANGES

@@ -0,0 +1,80 @@
+Changes we want
+===============
+Curently ssm-env recognises variables beginning with ssm://, containing an SSM parameter
+store key after this prefix, and substitutes them with their stored value by making calls
+to the AWS SSM API. It then executes a child process with this new environment.
+
+We want it to also recognise KMS-encrypted parameters and have them substituted into
+their plaintext value in the child process environment. These variables are of the format
+
+    !kms '<encoded value>'
+
+where <encoded value> stands in for the base64-encoded KMS-encrypted secure test.
+
+Example
+=======
+Given the following example environment
+
+    VAR1=VALUE1
+    VAR2=ssm:///omnes/caeli
+    VAR3=!kms 'ABCDEFG=='
+
+currently SSM would detect that VAR2 has the appropriate prefix, query the SSM Parameter
+Database for the /omnes/caeli key and change VAR2 into the plain value it returns, while
+leaving VAR1 and VAR3 unchanged since they don't have the ssm:// prefix. Supposing the
+stored value of said key is 'plainvalue' (without quotes), the environment of the child
+process will be
+
+    VAR1=VALUE1
+    VAR2=plainvalue
+    VAR3=!kms 'ABCDEFG=='
+
+After the changes, for that sample enviroment we want ssm-env to detect that VAR3 has
+the appropriate format and query the AWS KMS service to decrypt the kms-encrypted value,
+getting the following modified environment for its child process (assuming that ABCDEFG==
+is the base64 kms-encrypted value, under certain key, of 'muchacha' (without quotes):
+
+    VAR1=VALUE
+    VAR2=plainvalue
+    VAR3=muchacha
+
+Note that the behaviour for neither VAR1 nor VAR2 is changed from the current
+implementation.
+
+Considerations
+==============
+* Express the new KMS-encrypted value format as a regular expression.
+* Attempt to mimic the current implementation as much as possible, substituting SSM API
+  calls for the necessary KMS API calls but keeping the larger structure and spirit
+  faithful to the project.
+* Add tests akin to the current ones, in which we mock the KMS service.
+
+Existing implementation
+=======================
+The wanted extra functionality is currently provided by a shell function:
+
+kms_env() {
+  env -0 | while IFS== read -r -d '' name value
+  do
+    if echo "$value" | grep -q '^!kms '
+    then
+      cipher="$(echo "$value" | sed -e 's/!kms //' | sed -e "s/'//g")"
+      # Don't quote cipher since it could already be quoted, and it's expected
+      # to be a base64-encrypted value
+      plain="$(aws kms decrypt --ciphertext-blob fileb://<(echo $cipher | base64 -d) --output text --query Plaintext | base64 -d)" \
+        || return 1
+      echo "export $name=$plain"
+    fi
+  done
+}
+
+$(kms_env) || error_exit 'Error decrypting environment with kms_env'
+
+where the last line represents the replacing of the environment for the next
+commands. Note that the behaviour is not exactly the same -- the idea of the
+kms_env() function is to be run in a shell script previous to the commands
+that will inherit the new environment, while ssm-env is expected to be run
+and passing the subsequent commands as parameters, which it will then exec.
+
+However, it's useful as a reference for the format and the means in which
+the KMS-encrypted values are received and expanded.

README.md

@@ -16,21 +16,28 @@ You can most likely find the downloaded binary in `~/go/bin/ssm-env`
 ssm-env [-template STRING] [-with-decryption] [-no-fail] COMMAND
 ```
 
+### Parameter Formats
+
+- **SSM Parameter Store**: `ssm:///parameter-path` (fetches the value from AWS SSM Parameter Store)
+- **KMS Encrypted**: `!kms 'base64EncodedEncryptedValue'` (decrypts base64-encoded value using AWS KMS)
+
 ## Details
 
 Given the following environment:
 
 ```
 RAILS_ENV=production
 COOKIE_SECRET=ssm://prod.app.cookie-secret
+API_KEY=!kms 'base64EncodedEncryptedValue=='
 ```
 
-You can run the application using `ssm-env` to automatically populate the `COOKIE_SECRET` env var from SSM:
+You can run the application using `ssm-env` to automatically populate the `COOKIE_SECRET` env var from SSM and decrypt the `API_KEY` using AWS KMS:
 
 ```console
 $ ssm-env env
 RAILS_ENV=production
 COOKIE_SECRET=super-secret
+API_KEY=decrypted-value
 ```
 
 You can also configure how the parameter name is determined for an environment variable, by using the `-template` flag:

main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/base64"
 	"flag"
 	"fmt"
 	"os"
@@ -13,6 +14,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/kms"
 	"github.com/aws/aws-sdk-go/service/ssm"
 )
 
@@ -24,6 +26,9 @@ const (
 	// defaultBatchSize is the default number of parameters to fetch at once.
 	// The SSM API limits this to a maximum of 10 at the time of writing.
 	defaultBatchSize = 10
+
+	// KMS prefix for variables that contain KMS-encrypted values
+	KMSPrefix = "!kms "
 )
 
 // TemplateFuncs are helper functions provided to the template.
@@ -77,6 +82,7 @@ func main() {
 		batchSize: defaultBatchSize,
 		t:         t,
 		ssm:       &lazySSMClient{},
+		kms:       &lazyKMSClient{},
 		os:        os,
 	}
 	must(e.expandEnviron(*decrypt, *nofail))
@@ -93,7 +99,7 @@ type lazySSMClient struct {
 func (c *lazySSMClient) GetParameters(input *ssm.GetParametersInput) (*ssm.GetParametersOutput, error) {
 	// Initialize the SSM client (and AWS session) if it hasn't been already.
 	if c.ssm == nil {
-		sess, err := c.awsSession()
+		sess, err := awsSession()
 		if err != nil {
 			return nil, err
 		}
@@ -102,7 +108,27 @@ func (c *lazySSMClient) GetParameters(input *ssm.GetParametersInput) (*ssm.GetPa
 	return c.ssm.GetParameters(input)
 }
 
-func (c *lazySSMClient) awsSession() (*session.Session, error) {
+// lazyKMSClient wraps the AWS SDK KMS client such that the AWS session and
+// KMS client are not actually initialized until Decrypt is called for
+// the first time.
+type lazyKMSClient struct {
+	kms kmsClient
+}
+
+func (c *lazyKMSClient) Decrypt(input *kms.DecryptInput) (*kms.DecryptOutput, error) {
+	// Initialize the KMS client (and AWS session) if it hasn't been already.
+	if c.kms == nil {
+		sess, err := awsSession()
+		if err != nil {
+			return nil, err
+		}
+		c.kms = kms.New(sess)
+	}
+	return c.kms.Decrypt(input)
+}
+
+// awsSession creates and configures an AWS session with region detection
+func awsSession() (*session.Session, error) {
 	sess, err := session.NewSession(&aws.Config{
 		CredentialsChainVerboseErrors: aws.Bool(true),
 	})
@@ -132,6 +158,10 @@ type ssmClient interface {
 	GetParameters(*ssm.GetParametersInput) (*ssm.GetParametersOutput, error)
 }
 
+type kmsClient interface {
+	Decrypt(*kms.DecryptInput) (*kms.DecryptOutput, error)
+}
+
 type environ interface {
 	Environ() []string
 	Setenv(key, vale string)
@@ -155,6 +185,7 @@ type ssmVar struct {
 type expander struct {
 	t         *template.Template
 	ssm       ssmClient
+	kms       kmsClient
 	os        environ
 	batchSize int
 }
@@ -172,14 +203,33 @@ func (e *expander) parameter(k, v string) (*string, error) {
 	return nil, nil
 }
 
+type kmsVar struct {
+	envvar  string
+	encoded string
+}
+
 func (e *expander) expandEnviron(decrypt bool, nofail bool) error {
 	// Environment variables that point to some SSM parameters.
 	var ssmVars []ssmVar
-
+
+	// Environment variables that are KMS encrypted.
+	var kmsVars []kmsVar
+
 	uniqNames := make(map[string]bool)
 	for _, envvar := range e.os.Environ() {
 		k, v := splitVar(envvar)
 
+		// Check if this is a KMS encrypted value
+		if strings.HasPrefix(v, KMSPrefix) {
+			// Extract the base64 value by removing the prefix and any quotes
+			encodedPart := strings.TrimPrefix(v, KMSPrefix)
+			// Remove leading and trailing quotes if present
+			encodedPart = strings.Trim(encodedPart, "'\" ")
+
+			kmsVars = append(kmsVars, kmsVar{k, encodedPart})
+			continue
+		}
+
 		parameter, err := e.parameter(k, v)
 		if err != nil {
 			// TODO: Should this _also_ not error if nofail is passed?
@@ -197,34 +247,47 @@ func (e *expander) expandEnviron(decrypt bool, nofail bool) error {
 		}
 	}
 
-	if len(uniqNames) == 0 {
-		// Nothing to do, no SSM parameters.
-		return nil
-	}
+	// Process SSM parameters
+	if len(uniqNames) > 0 {
+		names := make([]string, len(uniqNames))
+		i := 0
+		for k := range uniqNames {
+			names[i] = k
+			i++
+		}
 
-	names := make([]string, len(uniqNames))
-	i := 0
-	for k := range uniqNames {
-		names[i] = k
-		i++
-	}
+		for i := 0; i < len(names); i += e.batchSize {
+			j := i + e.batchSize
+			if j > len(names) {
+				j = len(names)
+			}
 
-	for i := 0; i < len(names); i += e.batchSize {
-		j := i + e.batchSize
-		if j > len(names) {
-			j = len(names)
-		}
+			values, err := e.getParameters(names[i:j], decrypt, nofail)
+			if err != nil {
+				return err
+			}
 
-		values, err := e.getParameters(names[i:j], decrypt, nofail)
-		if err != nil {
-			return err
+			for _, v := range ssmVars {
+				val, ok := values[v.parameter]
+				if ok {
+					e.os.Setenv(v.envvar, val)
+				}
+			}
 		}
+	}
 
-		for _, v := range ssmVars {
-			val, ok := values[v.parameter]
-			if ok {
-				e.os.Setenv(v.envvar, val)
+	// Process KMS encrypted values
+	if len(kmsVars) > 0 {
+		for _, kv := range kmsVars {
+			decryptedValue, err := e.decryptKmsValue(kv.encoded, nofail)
+			if err != nil {
+				if nofail {
+					fmt.Fprintf(os.Stderr, "ssm-env: failed to decrypt KMS value: %v\n", err)
+					continue
+				}
+				return fmt.Errorf("failed to decrypt KMS value: %v", err)
 			}
+			e.os.Setenv(kv.envvar, decryptedValue)
 		}
 	}
 
@@ -287,6 +350,32 @@ func (e *invalidParametersError) Error() string {
 	return fmt.Sprintf("invalid parameters: %v", e.InvalidParameters)
 }
 
+// decryptKmsValue decrypts a base64-encoded KMS-encrypted value.
+func (e *expander) decryptKmsValue(encodedValue string, nofail bool) (string, error) {
+	// Add padding to base64 if needed
+	// Base64 encoding requires the string length to be a multiple of 4
+	padding := len(encodedValue) % 4
+	if padding != 0 {
+		encodedValue = encodedValue + strings.Repeat("=", 4-padding)
+	}
+
+	decodedBytes, err := base64.StdEncoding.DecodeString(encodedValue)
+	if err != nil {
+		return "", fmt.Errorf("failed to decode base64 value: %v", err)
+	}
+
+	input := &kms.DecryptInput{
+		CiphertextBlob: decodedBytes,
+	}
+
+	result, err := e.kms.Decrypt(input)
+	if err != nil {
+		return "", err
+	}
+
+	return string(result.Plaintext), nil
+}
+
 func splitVar(v string) (key, val string) {
 	parts := strings.Split(v, "=")
 	return parts[0], parts[1]