Feature/xargs (#23)

* initial attempt at xargs support

* simplify account lookup logic

* simplify xargs using new speculate

* remove unused exports

* fix cred caching

* add multi implementation

* finish multi initial design

* fix linting

* add changelog

Author: akerl-unpriv

CHANGELOG.md

@@ -1,3 +1,8 @@
+# 2.1.0 / 2019-08-19
+
+* [FEATURE] Add xargs subcommand
+* [FEATURE] Mutex for thread-safe credential lookups
+
 # 2.0.9 / 2019-08-15
 
 * [FEATURE] Update speculate to handle platform-specific env vars

cmd/travel.go

@@ -84,15 +84,15 @@ func travelRunner(cmd *cobra.Command, args []string) error {
 		}}
 	}
 
-	creds, err := path.TraverseWithOptions(opts)
+	c, err := path.TraverseWithOptions(opts)
 	if err != nil {
 		return err
 	}
 
-	for _, line := range creds.ToEnvVars() {
+	for _, line := range c.ToEnvVars() {
 		fmt.Println(line)
 	}
-	url, err := creds.ToCustomConsoleURL(servicePath)
+	url, err := c.ToCustomConsoleURL(servicePath)
 	if err != nil {
 		return err
 	}

cmd/xargs.go

@@ -0,0 +1,107 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/akerl/voyager/v2/cartogram"
+	"github.com/akerl/voyager/v2/multi"
+	"github.com/akerl/voyager/v2/travel"
+	"github.com/akerl/voyager/v2/yubikey"
+
+	"github.com/akerl/input/list"
+	"github.com/akerl/speculate/v2/creds"
+	"github.com/spf13/cobra"
+)
+
+var xargsCmd = &cobra.Command{
+	Use:   "xargs",
+	Short: "Run a command across many AWS accounts",
+	RunE:  xargsRunner,
+}
+
+func init() {
+	rootCmd.AddCommand(xargsCmd)
+	xargsCmd.Flags().StringP("role", "r", "", "Choose target role to use")
+	xargsCmd.Flags().String("profile", "", "Choose source profile to use")
+	xargsCmd.Flags().StringP("prompt", "p", "", "Choose prompt to use")
+	xargsCmd.Flags().BoolP("yubikey", "y", false, "Use Yubikey for MFA")
+	xargsCmd.Flags().StringP("command", "c", "", "Command to execute")
+}
+
+// revive:disable-next-line:cyclomatic
+func xargsRunner(cmd *cobra.Command, args []string) error {
+	flags := cmd.Flags()
+
+	flagRole, err := flags.GetString("role")
+	if err != nil {
+		return err
+	}
+
+	flagProfile, err := flags.GetString("profile")
+	if err != nil {
+		return err
+	}
+
+	promptFlag, err := flags.GetString("prompt")
+	if err != nil {
+		return err
+	}
+	promptGenerator, ok := list.Types[promptFlag]
+	if !ok {
+		return fmt.Errorf("prompt type not found: %s", promptFlag)
+	}
+	prompt := promptGenerator()
+
+	useYubikey, err := flags.GetBool("yubikey")
+	if err != nil {
+		return err
+	}
+
+	commandStr, err := flags.GetString("command")
+	if err != nil {
+		return err
+	}
+	if commandStr == "" {
+		return fmt.Errorf("command must be provided via --command / -c")
+	}
+
+	pack := cartogram.Pack{}
+	if err := pack.Load(); err != nil {
+		return err
+	}
+
+	grapher := travel.Grapher{
+		Prompt: prompt,
+		Pack:   pack,
+	}
+
+	opts := travel.DefaultTraverseOptions()
+	if useYubikey {
+		opts.MfaPrompt = &creds.MultiMfaPrompt{Backends: []creds.MfaPrompt{
+			yubikey.NewPrompt(),
+			&creds.DefaultMfaPrompt{},
+		}}
+	}
+
+	processor := multi.Processor{
+		Grapher:      grapher,
+		Options:      opts,
+		Args:         args,
+		RoleNames:    []string{flagRole},
+		ProfileNames: []string{flagProfile},
+	}
+
+	results, err := processor.ExecString(commandStr)
+	if err != nil {
+		return err
+	}
+
+	buffer, err := json.MarshalIndent(results, "", "  ")
+	if err != nil {
+		return err
+	}
+	fmt.Println(string(buffer))
+
+	return nil
+}

go.mod

@@ -4,10 +4,12 @@ go 1.12
 
 require (
 	github.com/99designs/keyring v1.1.1
+	github.com/BurntSushi/locker v0.0.0-20171006230638-a6e239ea1c69
 	github.com/akerl/input v0.0.3
-	github.com/akerl/speculate/v2 v2.0.8
+	github.com/akerl/speculate/v2 v2.1.0
 	github.com/akerl/timber/v2 v2.0.1
-	github.com/aws/aws-sdk-go v1.23.1
+	github.com/aws/aws-sdk-go v1.23.3
 	github.com/spf13/cobra v0.0.5
+	github.com/vbauerster/mpb/v4 v4.9.2
 	github.com/yawn/ykoath v1.0.2
 )

go.sum

@@ -1,16 +1,20 @@
 github.com/99designs/keyring v1.1.1 h1:3XSIPsh/MH5tHYT7ub1LoOkpGQuW/TdK//c04TnU2j4=
 github.com/99designs/keyring v1.1.1/go.mod h1:657DQuMrBZRtuL/voxVyiyb6zpMehlm5vLB9Qwrv904=
+github.com/BurntSushi/locker v0.0.0-20171006230638-a6e239ea1c69 h1:+tu3HOoMXB7RXEINRVIpxJCT+KdYiI7LAEAUrOw3dIU=
+github.com/BurntSushi/locker v0.0.0-20171006230638-a6e239ea1c69/go.mod h1:L1AbZdiDllfyYH5l5OkAaZtk7VkWe89bPJFmnDBNHxg=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/VividCortex/ewma v1.1.1 h1:MnEK4VOv6n0RSY4vtRe3h11qjxL3+t0B8yOL8iMXdcM=
+github.com/VividCortex/ewma v1.1.1/go.mod h1:2Tkkvm3sRDVXaiyucHiACn4cqf7DpdyLvmxzcbUokwA=
 github.com/akerl/input v0.0.3 h1:gy9u/L8WJWIJ9OB/5HSdOBIjrFfVkrRxmRlWd3I9+Co=
 github.com/akerl/input v0.0.3/go.mod h1:AVqHnOGmoQDEzX4jMX/ZZOC20+7MfodkmNbXhzqIJ2U=
-github.com/akerl/speculate/v2 v2.0.8 h1:NJ8xQGopA3OlFFOir7duni2K3aN3h75eolVC+NVK8nY=
-github.com/akerl/speculate/v2 v2.0.8/go.mod h1:qkhyP7kYtsPqTNpuE+6S4/gx5q2X1AmTcuDqp7Jmegw=
+github.com/akerl/speculate/v2 v2.1.0 h1:pmz1d5iIaGESJ/zh05P0eAURfx6o4+yy0Z99qnBWwwM=
+github.com/akerl/speculate/v2 v2.1.0/go.mod h1:qkhyP7kYtsPqTNpuE+6S4/gx5q2X1AmTcuDqp7Jmegw=
 github.com/akerl/timber/v2 v2.0.1 h1:hY4VCOJns7KsxwxP/ifSt3Rz9GZCfKewapaimObnA2E=
 github.com/akerl/timber/v2 v2.0.1/go.mod h1:jBjRGI2CWuvbZlrZkp1JO/X51pMlbg72NFy+Vnd59oI=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
 github.com/aws/aws-sdk-go v1.22.3/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
-github.com/aws/aws-sdk-go v1.23.1 h1:MXfB75PkuWJJmZMCQ40haFUuOIIGt1FuiWtPBXy8URA=
-github.com/aws/aws-sdk-go v1.23.1/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
+github.com/aws/aws-sdk-go v1.23.3 h1:Ty/4P6tOFJkDnKDrFJWnveznvESblf8QOheD1CwQPDU=
+github.com/aws/aws-sdk-go v1.23.3/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
 github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=
 github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
@@ -27,15 +31,14 @@ github.com/dvsekhvalnov/jose2go v0.0.0-20180829124132-7f401d37b68a/go.mod h1:7Bv
 github.com/ebfe/scard v0.0.0-20190212122703-c3d1b1916a95 h1:OM0MnUcXBysj7ZtXvThVWHMoahuKQ8FuwIdeSLcNdP4=
 github.com/ebfe/scard v0.0.0-20190212122703-c3d1b1916a95/go.mod h1:8hHvF8DlEq5kE3KWOsZQezdWq1OTOVxZArZMscS954E=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
-github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 h1:ZpnhV/YsD2/4cESfV5+Hoeu/iUR3ruzNvZ+yQfO03a0=
 github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2/go.mod h1:bBOAhwG1umN6/6ZUMtDFBMQR8jRg9O75tm9K00oMsK4=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
-github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c h1:6rhixN/i8ZofjG1Y75iExal34USq5p+wiN1tpie8IrU=
 github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c/go.mod h1:NMPJylDgVpX0MLRlPy15sqSwOFv/U1GZ2m21JhFfek0=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af h1:pmfjZENx5imkbgOkpRUYLnmbU7UEFbjtDA2hxJ1ichM=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
+github.com/keybase/go-keychain v0.0.0-20190712205309-48d3d31d256d h1:Z+RDyXzjKE0i2sTjZ/b1uxiGtPhFy34Ou/Tk0qwN0kM=
 github.com/keybase/go-keychain v0.0.0-20190712205309-48d3d31d256d/go.mod h1:JJNrCn9otv/2QP4D7SMJBgaleKpOf66PnW6F5WGNRIc=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -71,6 +74,8 @@ github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoH
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
+github.com/vbauerster/mpb/v4 v4.9.2 h1:/niWWc4MsgvOaCXJCdrWq2vOyN5inuQIq/Spn4Fi1XI=
+github.com/vbauerster/mpb/v4 v4.9.2/go.mod h1:svMK4M+/dWn+xLRlw1Hq4MzYq+5yX2moH6tktD+lIk0=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
 github.com/yawn/ykoath v1.0.2 h1:BYrUYmzu7nRjiqx5olhVhVZqh5NEUUMJf49ZqgrlEnc=
 github.com/yawn/ykoath v1.0.2/go.mod h1:dcXMmLrvt6WFkySkG2k8ZEqxiTbu/TWSI4+/Cb54+Lg=
@@ -83,8 +88,9 @@ golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20190712062909-fae7ac547cb7 h1:LepdCS8Gf/MVejFIt8lsiexZATdoGVyp5bcyS+rYoUI=
 golang.org/x/sys v0.0.0-20190712062909-fae7ac547cb7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190804053845-51ab0e2deafa h1:KIDDMLT1O0Nr7TSxp8xM5tJcdn8tgyAONntO829og1M=
+golang.org/x/sys v0.0.0-20190804053845-51ab0e2deafa/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

main.go

@@ -1,15 +1,16 @@
 package main
 
 import (
-	"fmt"
 	"os"
 
 	"github.com/akerl/voyager/v2/cmd"
+
+	"github.com/akerl/speculate/v2/helpers"
 )
 
 func main() {
 	if err := cmd.Execute(); err != nil {
-		fmt.Fprintln(os.Stderr, err)
+		helpers.PrintAwsError(err)
 		os.Exit(1)
 	}
 }

multi/main.go

@@ -0,0 +1,113 @@
+package multi
+
+import (
+	"os"
+	"strings"
+	"time"
+
+	"github.com/akerl/voyager/v2/travel"
+
+	"github.com/akerl/speculate/v2/creds"
+	"github.com/akerl/timber/v2/log"
+	"github.com/vbauerster/mpb/v4"
+	"github.com/vbauerster/mpb/v4/decor"
+)
+
+var logger = log.NewLogger("voyager")
+
+// Processor defines the settings for parallel processing
+type Processor struct {
+	Grapher      travel.Grapher
+	Options      travel.TraverseOptions
+	Args         []string
+	RoleNames    []string
+	ProfileNames []string
+}
+
+// ExecString runs a command string against a set of accounts
+func (p Processor) ExecString(cmd string) (map[string]creds.ExecResult, error) {
+	cmdSlice := strings.Split(cmd, " ")
+	return p.Exec(cmdSlice)
+}
+
+// Exec runs a command against a set of accounts
+func (p Processor) Exec(cmd []string) (map[string]creds.ExecResult, error) {
+	logger.InfoMsgf("processing command: %v", cmd)
+
+	paths, err := p.Grapher.ResolveAll(p.Args, p.RoleNames, p.ProfileNames)
+	if err != nil {
+		return map[string]creds.ExecResult{}, err
+	}
+
+	inputCh := make(chan workerInput, len(paths))
+	outputCh := make(chan workerOutput, len(paths))
+	refreshCh := make(chan time.Time)
+
+	for i := 1; i <= 10; i++ {
+		go execWorker(inputCh, outputCh)
+	}
+
+	for _, item := range paths {
+		inputCh <- workerInput{
+			Path:    item,
+			Options: p.Options,
+			Command: cmd,
+		}
+	}
+	close(inputCh)
+
+	progress := mpb.New(
+		mpb.WithOutput(os.Stderr),
+		mpb.WithManualRefresh(refreshCh),
+	)
+	bar := progress.AddBar(
+		int64(len(paths)),
+		mpb.AppendDecorators(
+			decor.Percentage(),
+		),
+		mpb.PrependDecorators(
+			decor.CountersNoUnit("%d / %d", decor.WCSyncWidth),
+		),
+	)
+
+	output := map[string]creds.ExecResult{}
+	for i := 1; i <= len(paths); i++ {
+		result := <-outputCh
+		output[result.AccountID] = result.ExecResult
+		bar.Increment()
+		refreshCh <- time.Now()
+	}
+	progress.Wait()
+
+	return output, nil
+}
+
+type workerInput struct {
+	Path    travel.Path
+	Options travel.TraverseOptions
+	Command []string
+}
+
+type workerOutput struct {
+	AccountID  string
+	ExecResult creds.ExecResult
+}
+
+func execWorker(inputCh <-chan workerInput, outputCh chan<- workerOutput) {
+	for item := range inputCh {
+		c, err := item.Path.TraverseWithOptions(item.Options)
+		if err != nil {
+			outputCh <- workerOutput{ExecResult: creds.ExecResult{Error: err}}
+			continue
+		}
+		accountID, err := c.AccountID()
+		if err != nil {
+			outputCh <- workerOutput{ExecResult: creds.ExecResult{Error: err}}
+			continue
+		}
+		outputCh <- workerOutput{
+			AccountID:  accountID,
+			ExecResult: c.Exec(item.Command),
+		}
+	}
+}

travel/cache.go

@@ -2,6 +2,7 @@ package travel
 
 import (
 	"fmt"
+
 	"github.com/akerl/speculate/v2/creds"
 	"github.com/aws/aws-sdk-go/service/sts"
 )
@@ -16,6 +17,7 @@ type Cache interface {
 // CheckCache returns credentials if they exist in the cache and are still valid
 // If the credentials exist but are invalid/expired, it removes them from the cache
 func CheckCache(c Cache, h Hop) (creds.Creds, bool) {
+	logger.DebugMsgf("checking cache for %+v", h)
 	cachedCreds, ok := c.Get(h)
 	if !ok {
 		return creds.Creds{}, false
@@ -57,6 +59,7 @@ type MapCache struct {
 // Put stores the credentials in the map
 func (mc *MapCache) Put(h Hop, c creds.Creds) error {
 	key := mc.hopToKey(h)
+	logger.DebugMsgf("mapcache: caching %s", key)
 	if mc.creds == nil {
 		mc.creds = map[string]creds.Creds{}
 	}
@@ -67,13 +70,15 @@ func (mc *MapCache) Put(h Hop, c creds.Creds) error {
 // Get returns credentials from the map, if they exist
 func (mc *MapCache) Get(h Hop) (creds.Creds, bool) {
 	key := mc.hopToKey(h)
+	logger.DebugMsgf("mapcache: getting %s", key)
 	creds, ok := mc.creds[key]
 	return creds, ok
 }
 
 // Delete removes credentials from the cache
 func (mc *MapCache) Delete(h Hop) error {
 	key := mc.hopToKey(h)
+	logger.DebugMsgf("mapcache: deleting %s", key)
 	delete(mc.creds, key)
 	return nil
 }