Merge pull request #385 from segmentio/mckern/fix-env-vars

Fix `chamber env` and `chamber export -f dotenv`

cmd/env.go

@@ -3,13 +3,19 @@ package cmd
 import (
 	"fmt"
 	"regexp"
+	"sort"
 	"strings"
 
+	"github.com/alessio/shellescape"
 	analytics "github.com/segmentio/analytics-go/v3"
 	"github.com/segmentio/chamber/v2/utils"
+
 	"github.com/spf13/cobra"
 )
 
+// originally ported from github.com/joho/godotenv
+const doubleQuoteSpecialChars = "\\\n\r\"!$`"
+
 var (
 	// envCmd represents the env command
 	envCmd = &cobra.Command{
@@ -18,27 +24,54 @@ var (
 		Args:  cobra.ExactArgs(1),
 		RunE:  env,
 	}
-	pattern *regexp.Regexp
+	preserveCase   bool
+	escapeSpecials bool
 )
 
 func init() {
+	envCmd.Flags().SortFlags = false
+	envCmd.Flags().BoolVarP(&preserveCase, "preserve-case", "p", false, "preserve variable name case")
+	envCmd.Flags().BoolVarP(&escapeSpecials, "escape-strings", "e", false, "escape special characters in values")
 	RootCmd.AddCommand(envCmd)
-	pattern = regexp.MustCompile(`[^\w@%+=:,./-]`)
 }
 
+// Print all secrets to standard out as valid shell key-value
+// pairs or return an error if secrets cannot be safely
+// represented as shell words.
 func env(cmd *cobra.Command, args []string) error {
+	envVars, err := exportEnv(cmd, args)
+	if err != nil {
+		return err
+	}
+
+	for i := range envVars {
+		fmt.Println(envVars[i])
+	}
+
+	return nil
+}
+
+// Handle the actual work of retrieving and validating secrets.
+// Returns a []string, with each string being a `key=value` pair,
+// and returns any errors encountered along the way.
+// Keys will be converted into valid shell variable names,
+// and converted to uppercase unless --preserve is passed.
+// Key ordering is non-deterministic and unstable, as returned
+// value from a given secret store is non-deterministic and unstable.
+func exportEnv(cmd *cobra.Command, args []string) ([]string, error) {
 	service := utils.NormalizeService(args[0])
 	if err := validateService(service); err != nil {
-		return fmt.Errorf("Failed to validate service: %w", err)
+		return nil, fmt.Errorf("Failed to validate service: %w", err)
 	}
 
 	secretStore, err := getSecretStore()
 	if err != nil {
-		return fmt.Errorf("Failed to get secret store: %w", err)
+		return nil, fmt.Errorf("Failed to get secret store: %w", err)
 	}
+
 	rawSecrets, err := secretStore.ListRaw(service)
 	if err != nil {
-		return fmt.Errorf("Failed to list store contents: %w", err)
+		return nil, fmt.Errorf("Failed to list store contents: %w", err)
 	}
 
 	if analyticsEnabled && analyticsClient != nil {
@@ -53,24 +86,117 @@ func env(cmd *cobra.Command, args []string) error {
 		})
 	}
 
+	params := make(map[string]string)
 	for _, rawSecret := range rawSecrets {
-		fmt.Printf("export %s=%s\n",
-			strings.ToUpper(key(rawSecret.Key)),
-			shellescape(rawSecret.Value))
+		params[key(rawSecret.Key)] = rawSecret.Value
 	}
 
-	return nil
+	out, err := buildEnvOutput(params)
+	if err != nil {
+		return nil, err
+	}
+
+	// ensure output prints variable declarations as exported
+	for i := range out {
+		// Sprintf because each declaration already ends in a newline
+		out[i] = fmt.Sprintf("export %s", out[i])
+	}
+
+	return out, nil
+}
+
+// output will be returned lexically sorted by key name
+func buildEnvOutput(params map[string]string) ([]string, error) {
+	out := []string{}
+	for _, key := range sortedKeys(params) {
+		name := sanitizeKey(key)
+		if !preserveCase {
+			name = strings.ToUpper(name)
+		}
+
+		if err := validateShellName(name); err != nil {
+			return nil, err
+		}
+
+		// the default format prints all escape sequences as
+		// string literals, and wraps values in single quotes
+		// if they're unsafe or multi-line strings.
+		s := fmt.Sprintf(`%s=%s`, name, shellescape.Quote(params[key]))
+		if escapeSpecials {
+			// this format collapses special characters like newlines
+			// or carriage returns. requires escape sequences to be interpolated
+			// by whatever parses our key="value" pairs.
+			s = fmt.Sprintf(`%s="%s"`, name, doubleQuoteEscape(params[key]))
+		}
+
+		// don't rely on printf to handle properly quoting or
+		// escaping shell output -- just white-knuckle it ourselves.
+
+		out = append(out, s)
+	}
+
+	return out, nil
 }
 
-// shellescape returns a shell-escaped version of the string s. The returned value
-// is a string that can safely be used as one token in a shell command line.
-func shellescape(s string) string {
-	if len(s) == 0 {
-		return "''"
+// The name of a variable can contain only letters (a-z, case insensitive),
+// numbers (0-9) or the underscore character (_). It may only begin with
+// a letter or an underscore.
+func validateShellName(s string) error {
+	shellChars := regexp.MustCompile(`^[A-Za-z0-9_]+$`).MatchString
+	validShellName := regexp.MustCompile(`^[A-Za-z_]{1}`).MatchString
+
+	if !shellChars(s) {
+		return fmt.Errorf("cmd: %q contains invalid characters for a shell variable name", s)
 	}
-	if pattern.MatchString(s) {
-		return "'" + strings.Replace(s, "'", "'\"'\"'", -1) + "'"
+
+	if !validShellName(s) {
+		return fmt.Errorf("cmd: shell variable name %q must start with a letter or underscore", s)
 	}
 
+	return nil
+}
+
+// note that all character width will be preserved; a single space
+// (or period, tab, or newline) will be replaced with a single underscore.
+// no squeezing/collapsing of replaced characters is performed at all.
+func sanitizeKey(s string) string {
+	// I promise, we don't actually care about allocations here.
+	// allocate *away*.
+	s = strings.TrimSpace(s)
+	s = strings.ReplaceAll(s, "-", "_")
+	s = strings.ReplaceAll(s, ".", "_")
+	// whitespace gets a visit from The Big Hammer that is regex.
+	s = regexp.MustCompile(`[[:space:]]`).ReplaceAllString(s, "_")
+
 	return s
 }
+
+// originally ported from github.com/joho/godotenv
+func doubleQuoteEscape(line string) string {
+	for _, c := range doubleQuoteSpecialChars {
+		toReplace := "\\" + string(c)
+		if c == '\n' {
+			toReplace = `\n`
+		}
+		if c == '\r' {
+			toReplace = `\r`
+		}
+		line = strings.Replace(line, string(c), toReplace, -1)
+	}
+	return line
+}
+
+// return the keys from params, sorted by keyname.
+// note that sort.Strings() is not case insensitive.
+// e.g. []string{"A", "b", "cat", "Dog", "dog"} will sort as:
+// []string{"A", "Dog", "b", "cat", "dog"}. That doesn't
+// really matter here but it may lead to surprises.
+func sortedKeys(params map[string]string) []string {
+	keys := []string{}
+
+	for key := range params {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	return keys
+}

cmd/env_test.go

@@ -0,0 +1,54 @@
+package cmd
+
+import (
+	"fmt"
+	"testing"
+)
+
+func Test_validateShellName(t *testing.T) {
+	tests := []struct {
+		name       string
+		str        string
+		shouldFail bool
+	}{
+		{name: "strings with spaces should fail", str: "invalid strings", shouldFail: true},
+		{name: "strings with only underscores should pass", str: "valid_string", shouldFail: false},
+		{name: "strings with dashes should fail", str: "validish-string", shouldFail: true},
+		{name: "strings that start with numbers should fail", str: "1invalidstring", shouldFail: true},
+		{name: "strings that start with underscores should pass", str: "_1validstring", shouldFail: false},
+		{name: "strings that contain periods should fail", str: "invalid.string", shouldFail: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := validateShellName(tt.str); (err != nil) != tt.shouldFail {
+				t.Errorf("validateShellName error: %v, expect wantErr %v", err, tt.shouldFail)
+			}
+		})
+	}
+}
+
+func Test_sanitizeKey(t *testing.T) {
+	tests := []struct {
+		given    string
+		expected string
+	}{
+		{given: "invalid strings", expected: "invalid_strings"},
+		{given: "extremely  invalid  strings", expected: "extremely__invalid__strings"},
+		{given: fmt.Sprintf("\nunbelievably\tinvalid\tstrings\n"), expected: "unbelievably_invalid_strings"},
+		{given: "valid_string", expected: "valid_string"},
+		{given: "validish-string", expected: "validish_string"},
+		{given: "valid.string", expected: "valid_string"},
+		// the following two strings should not be corrected, simply returned as-is.
+		{given: "1invalidstring", expected: "1invalidstring"},
+		{given: "_1validstring", expected: "_1validstring"},
+	}
+
+	for _, tt := range tests {
+		t.Run("test sanitizing key names", func(t *testing.T) {
+			if got := sanitizeKey(tt.given); got != tt.expected {
+				t.Errorf("shellName error: want %q, got %q", tt.expected, got)
+			}
+		})
+	}
+}

cmd/export.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"sort"
 	"strings"
 
 	"github.com/magiconair/properties"
@@ -18,8 +17,6 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
-const doubleQuoteSpecialChars = "\\\n\r\"!$`"
-
 // exportCmd represents the export command
 var (
 	exportFormat string
@@ -34,8 +31,10 @@ var (
 )
 
 func init() {
+	exportCmd.Flags().SortFlags = false
 	exportCmd.Flags().StringVarP(&exportFormat, "format", "f", "json", "Output format (json, yaml, java-properties, csv, tsv, dotenv, tfvars)")
 	exportCmd.Flags().StringVarP(&exportOutput, "output-file", "o", "", "Output file (default is standard output)")
+
 	RootCmd.AddCommand(exportCmd)
 }
 
@@ -103,7 +102,7 @@ func runExport(cmd *cobra.Command, args []string) error {
 	case "dotenv":
 		err = exportAsEnvFile(params, w)
 	case "tfvars":
-		err = exportAsTfvars(params, w)
+		err = exportAsTFvars(params, w)
 	default:
 		err = fmt.Errorf("Unsupported export format: %s", exportFormat)
 	}
@@ -115,22 +114,36 @@ func runExport(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+// this is fundamentally broken, in that there is no actual .env file
+// spec. some parsers support values spanned over multiple lines
+// as long as they're quoted, others only support character literals
+// inside of quotes. we should probably offer the option to control
+// which spec we adhere to, or use a marshaler that provides a
+// spec instead of hoping for the best.
 func exportAsEnvFile(params map[string]string, w io.Writer) error {
-	// Env File like:
-	// KEY=VAL
-	// OTHER=OTHERVAL
-	for _, k := range sortedKeys(params) {
-		key := strings.ToUpper(k)
-		key = strings.Replace(key, "-", "_", -1)
-		w.Write([]byte(fmt.Sprintf(`%s="%s"`+"\n", key, doubleQuoteEscape(params[k]))))
+	// use top-level escapeSpecials variable to ensure that
+	// the dotenv format prints escaped values every time
+	escapeSpecials = true
+	out, err := buildEnvOutput(params)
+	if err != nil {
+		return err
 	}
+
+	for i := range out {
+		_, err := w.Write([]byte(fmt.Sprintln(out[i])))
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
-func exportAsTfvars(params map[string]string, w io.Writer) error {
+func exportAsTFvars(params map[string]string, w io.Writer) error {
 	// Terraform Variables is like dotenv, but removes the TF_VAR and keeps lowercase
 	for _, k := range sortedKeys(params) {
-		key := strings.TrimPrefix(k, "tf_var_")
+		key := sanitizeKey(strings.TrimPrefix(k, "tf_var_"))
+
 		w.Write([]byte(fmt.Sprintf(`%s = "%s"`+"\n", key, doubleQuoteEscape(params[k]))))
 	}
 	return nil
@@ -173,43 +186,21 @@ func exportAsCsv(params map[string]string, w io.Writer) error {
 	defer csvWriter.Flush()
 	for _, k := range sortedKeys(params) {
 		if err := csvWriter.Write([]string{k, params[k]}); err != nil {
-			return fmt.Errorf("Failed to write param %s to CSV file: %w", k, err)
+			return fmt.Errorf("Failed to write param %q to CSV file: %w", k, err)
 		}
 	}
 	return nil
 }
 
 func exportAsTsv(params map[string]string, w io.Writer) error {
 	// TSV (Tab Separated Values) like:
+	tsvWriter := csv.NewWriter(w)
+	tsvWriter.Comma = '\t'
+	defer tsvWriter.Flush()
 	for _, k := range sortedKeys(params) {
-		if _, err := fmt.Fprintf(w, "%s\t%s\n", k, params[k]); err != nil {
-			return fmt.Errorf("Failed to write param %s to TSV file: %w", k, err)
+		if err := tsvWriter.Write([]string{k, params[k]}); err != nil {
+			return fmt.Errorf("Failed to write param %q to TSV file: %w", k, err)
 		}
 	}
 	return nil
 }
-
-func sortedKeys(params map[string]string) []string {
-	keys := make([]string, len(params))
-	i := 0
-	for k := range params {
-		keys[i] = k
-		i++
-	}
-	sort.Strings(keys)
-	return keys
-}
-
-func doubleQuoteEscape(line string) string {
-	for _, c := range doubleQuoteSpecialChars {
-		toReplace := "\\" + string(c)
-		if c == '\n' {
-			toReplace = `\n`
-		}
-		if c == '\r' {
-			toReplace = `\r`
-		}
-		line = strings.Replace(line, string(c), toReplace, -1)
-	}
-	return line
-}