Merge pull request #2 from kevinglasson/add-env-command

Add ability to dump parameters to environment for running a command

README.md

@@ -4,6 +4,7 @@
 
 ## Contents
 - [Contents](#contents)
+- [Demo](#demo)
 - [Installation](#installation)
   - [Using go get](#using-go-get)
   - [Pre-built binaries](#pre-built-binaries)
@@ -21,6 +22,9 @@
     - [File format support](#file-format-support)
 - [Why?](#why)
 - [Acknowledgements](#acknowledgements)
+
+## Demo
+[![asciicast](https://asciinema.org/a/GTP4YjvB3TWdcOSPD9swooZGa.svg)](https://asciinema.org/a/GTP4YjvB3TWdcOSPD9swooZGa)
 ## Installation
 ### Using go get
 To install use `go get` with or without -u to have goss installed in your `$GOBIN`.

cmd/completion.go

@@ -0,0 +1,88 @@
+/*
+Copyright © 2020 NAME HERE <EMAIL ADDRESS>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package cmd
+
+import (
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+// completionCmd represents the completion command
+var completionCmd = &cobra.Command{
+	Use:   "completion [bash|zsh|fish|powershell]",
+	Short: "Generate completion script",
+	Long: `To load completions:
+
+Bash:
+
+$ source <(goss completion bash)
+
+# To load completions for each session, execute once:
+Linux:
+  $ goss completion bash > /etc/bash_completion.d/goss
+MacOS:
+  $ goss completion bash > /usr/local/etc/bash_completion.d/goss
+
+Zsh:
+
+# If shell completion is not already enabled in your environment you will need
+# to enable it.  You can execute the following once:
+
+$ echo "autoload -U compinit; compinit" >> ~/.zshrc
+
+# To load completions for each session, execute once:
+$ goss completion zsh > "${fpath[1]}/_goss"
+
+# You will need to start a new shell for this setup to take effect.
+
+Fish:
+
+$ goss completion fish | source
+
+# To load completions for each session, execute once:
+$ goss completion fish > ~/.config/fish/completions/goss.fish
+`,
+	DisableFlagsInUseLine: true,
+	ValidArgs:             []string{"bash", "zsh", "fish", "powershell"},
+	Args:                  cobra.ExactValidArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		switch args[0] {
+		case "bash":
+			cmd.Root().GenBashCompletion(os.Stdout)
+		case "zsh":
+			cmd.Root().GenZshCompletion(os.Stdout)
+		case "fish":
+			cmd.Root().GenFishCompletion(os.Stdout, true)
+		case "powershell":
+			cmd.Root().GenPowerShellCompletion(os.Stdout)
+		}
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(completionCmd)
+
+	// Here you will define your flags and configuration settings.
+
+	// Cobra supports Persistent Flags which will work for this command
+	// and all subcommands, e.g.:
+	// completionCmd.PersistentFlags().String("foo", "", "A help for foo")
+
+	// Cobra supports local flags which will only run when this command
+	// is called directly, e.g.:
+	// completionCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
+}

cmd/env.go

@@ -26,109 +26,175 @@ package cmd
 
 import (
 	"fmt"
-
+	"log"
+	"os"
+	"os/exec"
+	"os/signal"
+	"runtime"
+	"strings"
+	"syscall"
+
+	"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/spf13/cobra"
 )
 
-// envCmd represents the env command
-var envCmd = &cobra.Command{
-	Use:   "env",
-	Short: "A brief description of your command",
-	Run: func(cmd *cobra.Command, args []string) {
-		fmt.Println("env called")
-	},
-}
+var (
+	// For flags.
+	envPath string
+
+	// envCmd represents the env command.
+	envCmd = &cobra.Command{
+		Use:   "env",
+		Short: "Load parameters into the environment and run a command",
+		Args:  cobra.MinimumNArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Println("env called")
+			envRun(args[0], args[1:], envPath)
+		},
+	}
+)
 
 func init() {
-	// rootCmd.AddCommand(envCmd)
+	rootCmd.AddCommand(envCmd)
+
+	envCmd.Flags().StringVarP(
+		&envPath, "path", "p", "", "parameter path",
+	)
+	envCmd.MarkFlagRequired("path")
+
+}
 
-	// Here you will define your flags and configuration settings.
+func envRun(command string, args []string, path string) error {
+	fmt.Printf("Command: %s\nArgs: %s\nPath: %s\n", command, args, path)
 
-	// Cobra supports Persistent Flags which will work for this command
-	// and all subcommands, e.g.:
-	// envCmd.PersistentFlags().String("foo", "", "A help for foo")
+	// Fetch the parameters.
+	env := environ(os.Environ())
+	mp, err := getParameters(path)
+	if err != nil {
+		return fmt.Errorf("Failed to get parameters: %w", err)
+	}
 
-	// Cobra supports local flags which will only run when this command
-	// is called directly, e.g.:
-	// envCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
+	// Set the parameters in the environment.
+	for k, v := range mp {
+		env.Set(k, v)
+	}
+
+	if !supportsExecSyscall() {
+		return execCmd(command, args, env)
+	}
+
+	return execSyscall(command, args, env)
 }
 
-// func execEnvironment(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
-// 	val, err := creds.Get()
-// 	if err != nil {
-// 		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
-// 	}
-
-// 	env := environ(os.Environ())
-// 	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)
-
-// 	log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
-// 	env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
-// 	env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)
-
-// 	if val.SessionToken != "" {
-// 		log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
-// 		env.Set("AWS_SESSION_TOKEN", val.SessionToken)
-// 		env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
-// 	}
-// 	if expiration, err := creds.ExpiresAt(); err == nil {
-// 		log.Println("Setting subprocess env: AWS_SESSION_EXPIRATION")
-// 		env.Set("AWS_SESSION_EXPIRATION", iso8601.Format(expiration))
-// 	}
-
-// 	if !supportsExecSyscall() {
-// 		return execCmd(input.Command, input.Args, env)
-// 	}
-
-// 	return execSyscall(input.Command, input.Args, env)
-// }
-
-// func execCmd(command string, args []string, env []string) error {
-// 	log.Printf("Starting child process: %s %s", command, strings.Join(args, " "))
-
-// 	cmd := exec.Command(command, args...)
-// 	cmd.Stdin = os.Stdin
-// 	cmd.Stdout = os.Stdout
-// 	cmd.Stderr = os.Stderr
-// 	cmd.Env = env
-
-// 	sigChan := make(chan os.Signal, 1)
-// 	signal.Notify(sigChan)
-
-// 	if err := cmd.Start(); err != nil {
-// 		return err
-// 	}
-
-// 	go func() {
-// 		for {
-// 			sig := <-sigChan
-// 			cmd.Process.Signal(sig)
-// 		}
-// 	}()
-
-// 	if err := cmd.Wait(); err != nil {
-// 		cmd.Process.Signal(os.Kill)
-// 		return fmt.Errorf("Failed to wait for command termination: %v", err)
-// 	}
-
-// 	waitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus)
-// 	os.Exit(waitStatus.ExitStatus())
-// 	return nil
-// }
-
-// func execSyscall(command string, args []string, env []string) error {
-// 	log.Printf("Exec command %s %s", command, strings.Join(args, " "))
-
-// 	argv0, err := exec.LookPath(command)
-// 	if err != nil {
-// 		return fmt.Errorf("Couldn't find the executable '%s': %w", command, err)
-// 	}
-
-// 	log.Printf("Found executable %s", argv0)
-
-// 	argv := make([]string, 0, 1+len(args))
-// 	argv = append(argv, command)
-// 	argv = append(argv, args...)
-
-// 	return syscall.Exec(argv0, argv, env)
-// }
+// getParameters is a virtual copy of listParameters - it needs to be refactored
+func getParameters(
+	path string,
+) (map[string]string, error) {
+	// Create Session.
+	sess, err := session.NewSession()
+	if err != nil {
+		return nil, fmt.Errorf("Session error: %w", err)
+	}
+
+	// Create SSM service.
+	svc := ssm.New(sess)
+
+	// Retrieve parameters.
+	res, err := svc.GetParametersByPath(
+		&ssm.GetParametersByPathInput{
+			Path:           aws.String(path),
+			Recursive:      aws.Bool(true),
+			WithDecryption: aws.Bool(true),
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("SSM request error: %w", err)
+	}
+
+	// Put the variables into a k-v map.
+	mp := make(map[string]string)
+	for _, v := range res.Parameters {
+		ss := strings.Split(*v.Name, "/")
+		key := ss[len(ss)-1]
+		if key != "" {
+			mp[key] = *v.Value
+		}
+	}
+
+	return mp, nil
+}
+
+func supportsExecSyscall() bool {
+	return runtime.GOOS == "linux" || runtime.GOOS == "darwin" || runtime.GOOS == "freebsd"
+}
+
+func execCmd(command string, args []string, env []string) error {
+	log.Printf("Starting child process: %s %s", command, strings.Join(args, " "))
+
+	cmd := exec.Command(command, args...)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Env = env
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan)
+
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+
+	go func() {
+		for {
+			sig := <-sigChan
+			cmd.Process.Signal(sig)
+		}
+	}()
+
+	if err := cmd.Wait(); err != nil {
+		cmd.Process.Signal(os.Kill)
+		return fmt.Errorf("Failed to wait for command termination: %v", err)
+	}
+
+	waitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus)
+	os.Exit(waitStatus.ExitStatus())
+	return nil
+}
+
+func execSyscall(command string, args []string, env []string) error {
+	argv0, err := exec.LookPath(command)
+	if err != nil {
+		return fmt.Errorf("Couldn't find the executable '%s': %w", command, err)
+	}
+
+	argv := make([]string, 0, 1+len(args))
+	argv = append(argv, command)
+	argv = append(argv, args...)
+
+	return syscall.Exec(argv0, argv, env)
+}
+
+// environ is a slice of strings representing the environment, in the form "key=value".
+type environ []string
+
+// Unset an environment variable by key
+func (e *environ) Unset(key string) {
+	for i := range *e {
+		// If we found the key
+		if strings.HasPrefix((*e)[i], key+"=") {
+			// Move the last value to replace the key
+			(*e)[i] = (*e)[len(*e)-1]
+			// Slice of the last value as we moved it to 'i'
+			*e = (*e)[:len(*e)-1]
+			break
+		}
+	}
+}
+
+// Set adds an environment variable, replacing any existing ones of the same key
+func (e *environ) Set(key, val string) {
+	e.Unset(key)
+	*e = append(*e, key+"="+val)
+}