aws-env.go

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"

	jsoniter "github.com/json-iterator/go" //native json lib doesn't support marshalling map[interface{}]interface{}

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

type dict map[interface{}]interface{}

var version string

func main() {

	if len(os.Args) > 1 {
		if os.Args[1] == "-v" {
			fmt.Printf("awsenv version %s\n", version)
			os.Exit(0)
		}
	}

	// Set log level if defind, default level info
	if os.Getenv("LOG_LEVEL") != "" {
		golog.SetLevel(os.Getenv("LOG_LEVEL"))
	}

	keys := strings.Split(os.Getenv("SSM_PATH"), "/")
	params := make(map[string]string)

	// Remove the empty string created by the split
	if keys[0] == "" {
		keys = keys[1:]
	}

	path := ""
	// Loop through the sub paths and retrieve parameters
	for i := range keys {
		path = path + "/" + keys[i]
		golog.Infof("Retrieving parameters in path %s", path)
		ExportVariables(path, "", params)
	}

	toStdout := strings.ToLower(os.Getenv("TO_STDOUT"))

	if toStdout == "true" {
		ParametersToStdout(params)
	} else {
		ParametersToFile(params)
	}
}

func CreateClient() *ssm.SSM {
	maxRetries := GetEnv("AWSENV_MAX_RETRIES", "3")
	maxRetriesInt, _ := strconv.Atoi(maxRetries)

	config := &aws.Config{MaxRetries: aws.Int(maxRetriesInt)}
	session := session.Must(session.NewSession(config))
	return ssm.New(session)
}

func ExportVariables(path string, nextToken string, params map[string]string) {
	client := CreateClient()

	input := &ssm.GetParametersByPathInput{
		Path:           &path,
		WithDecryption: aws.Bool(true),
	}

	if nextToken != "" {
		input.SetNextToken(nextToken)
	}

	output, err := client.GetParametersByPath(input)

	if err != nil {
		golog.Fatal(err)
	}

	for _, element := range output.Parameters {
		env, value := TrimParameter(path, element)
		golog.Debugf("Found parameter %s in path %s", env, path)
		params[env] = value
	}

	if output.NextToken != nil {
		ExportVariables(path, *output.NextToken, params)
	}
}

func TrimParameter(path string, parameter *ssm.Parameter) (string, string) {
	name := *parameter.Name
	value := *parameter.Value

	env := strings.Trim(name[len(path):], "/")
	value = strings.Replace(value, "\n", "\\n", -1)

	return env, value
}

func ParametersToStdout(params map[string]string) {
	golog.Debug("Printing params to stdout")
	for key, value := range params {
		key = strings.ToUpper(key)
		fmt.Printf("export %s=$'%s'\n", key, value)
	}
}

func ParametersToFile(params map[string]string) {
	var fileContent []byte

	format := GetEnv("FORMAT", "export")
	if format == "json" {
		fileContent = ParametersToJson(params)
	} else if format == "json-nested" {
		fileContent = ParametersToNestedJson(params)
	} else {
		fileContent = ParametersToFormattedString(format, params)
	}

	dir := GetEnv("DIRECTORY", "/ssm")
	// Create directory if it doesn't exist
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		golog.Debugf("Creating directory %s", dir)
		os.MkdirAll(dir, 0755)
	}

	filename := GetEnv("AWSENV_FILENAME", ".env")
	golog.Debugf("Writing params to %s/%s", dir, filename)

	err := ioutil.WriteFile(dir+"/"+filename, fileContent, 0744)
	if err != nil {
		golog.Fatal(err)
	}
}

func ParametersToJson(params map[string]string) []byte {
	json, err := jsoniter.MarshalIndent(params, "", " ")
	if err != nil {
		golog.Fatal(err)
	}

	return json
}

func ParametersToNestedJson(params map[string]string) []byte {
	nestedParams := make(dict)
	for key, value := range params {
		keys := strings.Split(key, "_")
		SetNestedParamKeys(nestedParams, keys, value)
	}

	json, err := jsoniter.MarshalIndent(nestedParams, "", " ")
	if err != nil {
		golog.Fatal(err)
	}

	return json
}

func SetNestedParamKeys(d dict, keys []string, value string) {
	if len(keys) == 1 {
		d[keys[0]] = value
		return
	}
	v, ok := d[keys[0]]
	if !ok {
		v = dict{}
		d[keys[0]] = v
	}
	SetNestedParamKeys(v.(dict), keys[1:], value)
}

func ParametersToFormattedString(format string, params map[string]string) []byte {
	var buffer bytes.Buffer

	for key, value := range params {
		buffer.WriteString(FormatParameter(key, value, format))
	}

	return buffer.Bytes()
}

func FormatParameter(key string, value string, format string) string {
	switch format {
	case "shell":
		return fmt.Sprintf("%s='%s'\n", key, value)
	case "unquoted-shell":
		return fmt.Sprintf("%s=%s\n", key, value)
	default:
		return fmt.Sprintf("export %s=$'%s'\n", key, value)
	}
}

func GetEnv(envvar string, defaultValue string) string {
	value, exists := os.LookupEnv(envvar)
	if !exists {
		return defaultValue
	}
	return value
}