-
Notifications
You must be signed in to change notification settings - Fork 24
/
main.go
396 lines (325 loc) · 12.8 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
package main
import (
"fmt"
"os"
"os/exec"
"os/signal"
"runtime"
"strings"
"sync"
"syscall"
"github.com/AlecAivazis/survey/v2"
"github.com/aws/aws-sdk-go/service/ssm"
"github.com/sirupsen/logrus"
log "github.com/sirupsen/logrus"
flag "github.com/spf13/pflag"
"github.com/disneystreaming/go-ssmhelpers/aws"
"github.com/disneystreaming/go-ssmhelpers/aws/session"
ssmhelpers "github.com/disneystreaming/go-ssmhelpers/ssm"
"github.com/disneystreaming/go-ssmhelpers/ssm/instance"
"github.com/disneystreaming/go-ssmhelpers/util"
"github.com/disneystreaming/gomux"
)
var myInstances ssmhelpers.CommaSlice
var myFilters ssmhelpers.CommaSlice
var myProfiles ssmhelpers.CommaSlice
var myRegions ssmhelpers.CommaSlice
var myTags ssmhelpers.ListSlice
var allProfilesFlag bool
var verboseFlag int
var dryRunFlag bool
var sessionName string
var limitFlag int
var totalInstances int
var versionFlag bool
var version = "devel"
var commit = "notpassed"
func main() {
// Get the number of cores available for parallelization
runtime.GOMAXPROCS(runtime.NumCPU())
// Set logs to go to stdout by default
log.SetOutput(os.Stdout)
log.SetFormatter(&log.TextFormatter{
// Disable level truncation, timestamp, and pad out the level text to even it up
DisableLevelTruncation: true,
DisableTimestamp: true,
})
errLog := logrus.New()
errLog.SetFormatter(&log.TextFormatter{
// Disable level truncation, timestamp, and pad out the level text to even it up
DisableLevelTruncation: true,
DisableTimestamp: true,
})
// Flag to indicate that the user wants to perform a dry run
flag.BoolVar(&dryRunFlag, "dry-run", false, "Retrieve the list of profiles, regions, and instances on which sessions would be started.")
// Flag to indicate that the user wants to execute a command against all of their configured profiles
flag.BoolVar(&allProfilesFlag, "all-profiles", false, "[USE WITH CAUTION] Parse through ~/.aws/config to target all profiles.")
// Flag to enable increasingly-verbose output
flag.IntVar(&verboseFlag, "log-level", 0, "Sets verbosity of output:\n0 = quiet, 1 = terse, 2 = warn, 3 = debug")
// Flag for instance selection
flag.VarP(&myInstances, "instances", "i", "Specify what instance IDs you want to target.\nMultiple allowed, delimited by commas (e.g. --instances i-12345,i-23456)")
// Flags for filters
flag.VarP(&myFilters, "filter", "f", "Filter instances based on tag value. Tags are evaluated with logical AND (instances must match all tags).\nMultiple allowed, delimited by commas (e.g. env=dev,foo=bar)")
flag.VarP(&myTags, "tag", "t", "Adds the specified tag as an additional column to be displayed during the instance selection prompt.")
// Flags for profiles/regions
flag.VarP(&myProfiles, "profiles", "p", "Specify a specific profile to use with your API calls.\nMultiple allowed, delimited by commas (e.g. --profiles dev,myaccount)")
flag.VarP(&myRegions, "regions", "r", "Specify a specific region to use with your API calls.\n"+
"This option will override any profile settings in your config file.\n"+
"Multiple allowed, delimited by commas (e.g. --regions us-east-1,us-west-2)\n\n"+
"[NOTE] Mixing --profiles and --regions will result in your command targeting every matching instance in the selected profiles and regions.\n"+
"e.g., \"--profiles foo,bar,baz --regions us-east-1,us-west-2,eu-east-1\" will target instances in each of the profile/region combinations:\n"+
"\t\"foo@us-east-1, foo@us-west-2, foo@eu-east-1\"\n"+
"\t\"bar@us-east-1, bar@us-west-2, bar@eu-east-1\"\n"+
"\t\"baz@us-east-1, baz@us-west-2, baz@eu-east-1\"\n"+
"Please be careful.")
// Flag to allow naming of tmux session
flag.StringVar(&sessionName, "session-name", "ssm-session", "Specify a name for the tmux session created when multiple instances are selected")
// Flag to set a limit to the number of instances returned by the SSM/EC2 API query
flag.IntVarP(&limitFlag, "limit", "l", 20, "Set a limit for the number of instance results returned per profile/region combination.")
// Flag to show the version number
flag.BoolVar(&versionFlag, "version", false, "Show version and quit")
flag.Parse()
if versionFlag {
fmt.Printf("Version: %s\tGit Commit Hash: %s\n", version, commit)
os.Exit(0)
}
if verboseFlag == 3 {
log.SetLevel(log.DebugLevel)
}
if myProfiles == nil {
env, exists := os.LookupEnv("AWS_PROFILE")
if exists {
myProfiles.Set(env)
} else {
myProfiles.Set("default")
}
}
if myRegions == nil {
env, exists := os.LookupEnv("AWS_REGION")
if exists {
myRegions.Set(env)
}
}
// If --all-profiles is set, we call getAWSProfiles() and iterate through the user's ~/.aws/config
if allProfilesFlag {
profiles, err := aws.GetAWSProfiles()
if profiles != nil && err == nil {
myProfiles = profiles
} else {
errLog.Error("Could not load profiles.", err)
return
}
}
// Set up our AWS session for each permutation of profile + region
sessionPool := session.NewPoolSafe(myProfiles, myRegions)
// Set up our filters
var filterMaps []map[string]string
// Convert the filter slice to a map
filterMap := make(map[string]string)
if len(myFilters) > 0 {
util.SliceToMap(myFilters, &filterMap)
filterMaps = append(filterMaps, filterMap)
}
var wg sync.WaitGroup
// Master list for later
instancePool := instance.InstanceInfoSafe{
AllInstances: make(map[string]instance.InstanceInfo),
}
// Iterate through our AWS sessions
wg.Add(len(sessionPool.Sessions))
for _, sess := range sessionPool.Sessions {
go func(sess *session.Pool, instancePool *instance.InstanceInfoSafe) {
defer wg.Done()
instanceChan := make(chan []*ssm.InstanceInformation)
errChan := make(chan error)
svc := ssm.New(sess.Session)
go ssmhelpers.GetInstanceList(svc, filterMaps, myInstances, false, instanceChan, errChan)
instanceList, err := <-instanceChan, <-errChan
if err != nil {
log.Debugf("AWS Session Parameters: %s, %s", *sess.Session.Config.Region, sess.ProfileName)
log.Error(err)
}
totalInstances = len(instanceList)
ssmhelpers.CheckInstanceReadiness(sess, svc, instanceList, instancePool, limitFlag)
}(sess, &instancePool)
}
wg.Wait()
if verboseFlag > 0 {
log.Infof("Found %d usable instances.", len(instancePool.AllInstances))
}
// No functional results, exit now
if len(instancePool.AllInstances) == 0 {
return
}
// If -i flag is set, don't prompt for instance selection
if !dryRunFlag {
// Single instance specified or found, starting session in current terminal (non-multiplexed)
if len(myInstances) == 1 {
for _, v := range instancePool.AllInstances {
if err := startSSMSession(v.Profile, v.Region, v.InstanceID); err != nil {
log.Errorf("Failed to start ssm-session for instance %s\n%s", v.InstanceID, err)
}
}
return
}
// Multiple instances specified or found, check to see if we're in a tmux session to avoid nesting
if len(myInstances) > 1 && len(instancePool.AllInstances) > 1 {
var instances []instance.InstanceInfo
for _, v := range instancePool.AllInstances {
instances = append(instances, v)
}
if err := configTmuxSession(sessionName, instances); err != nil {
log.Fatal(err)
}
} else {
// If -i was not specified, go to a selection prompt before starting sessions
selectedInstances, err := startSelectionPrompt(&instancePool, totalInstances, myTags)
if err != nil {
log.Fatalf("Error during instance selection\n%s", err)
}
// If only one instance was selected, don't bother with a tmux session
if len(selectedInstances) == 1 {
for _, v := range selectedInstances {
if err := startSSMSession(v.Profile, v.Region, v.InstanceID); err != nil {
log.Errorf("Failed to start ssm-session for instance %s\n%s", v.InstanceID, err)
}
}
return
}
if err = configTmuxSession(sessionName, selectedInstances); err != nil {
log.Fatal(err)
}
}
// Make sure we aren't going to nest tmux sessions
currentTmuxSocket := os.Getenv("TMUX")
if len(currentTmuxSocket) == 0 {
if err := attachTmuxSession(sessionName); err != nil {
log.Errorf("Could not attach to tmux session '%s'\n%s", sessionName, err)
}
} else {
log.Info("To force nested Tmux sessions unset $TMUX")
log.Infof("Attach to the session with `tmux attach -t %s`", sessionName)
}
}
}
func configTmuxSession(sessionName string, selectedInstances []instance.InstanceInfo) (err error) {
// Initialize our tmux session
tmuxSession, err := gomux.NewSession(sessionName)
if err != nil {
return fmt.Errorf("Failed to create tmux session\n%s", err)
}
// Create the window in which our ssm session panes will live
tmuxWindow, err := tmuxSession.AddWindow("ssm")
if err != nil {
return fmt.Errorf("Failed to create tmux window\n%s", err)
}
// Configure our session-specific settings
configList := []string{
"set-option -t " + sessionName + " pane-border-status top",
"set-option -t " + sessionName + " mouse on",
}
for _, v := range configList {
if err = tmuxSession.Windows[0].SetConfig(v); err != nil {
return fmt.Errorf("Failed to set tmux configuration for window\n%s", err)
}
}
// Multiple instances specified or found, starting tmux session and attaching current terminal to it
for _, v := range selectedInstances {
// Add a window for our instance to our tmux session
if err = addInstanceToTmuxWindow(tmuxWindow, v.Profile, v.Region, v.InstanceID); err != nil {
return fmt.Errorf("Failed to add instance %s to tmux session\n%s", v.InstanceID, err)
}
// Re-tile our layout after each window to avoid the "pane too small" error
if err = tmuxWindow.SetConfig("select-layout -t " + sessionName + " tiled"); err != nil {
return fmt.Errorf("Failed to re-tile panes\n%s", err)
}
}
// Don't kill window 0 of our current session if we're already in a session
currentTmuxSocket := os.Getenv("TMUX")
if len(currentTmuxSocket) == 0 {
if err = tmuxWindow.KillPane(0); err != nil {
return fmt.Errorf("Failed to remove empty pane at index 0\n%s", err)
}
}
// Re-tile our layout one last time since we removed the empty pane
if err = tmuxWindow.SetConfig("select-layout -t " + sessionName + " tiled"); err != nil {
return fmt.Errorf("Failed to re-tile panes\n%s", err)
}
return
}
func setSessionStatus(sessionName string, status string) (err error) {
rawCmd := exec.Command("tmux", "set-option", "-t", sessionName, "status-left", status)
return rawCmd.Run()
}
func startSSMSession(profile string, region string, instanceID string) error {
rawCmd := exec.Command("aws", "ssm", "start-session", "--profile", profile, "--region", region, "--target", instanceID)
rawCmd.Stdin = os.Stdin
rawCmd.Stdout = os.Stdout
rawCmd.Stderr = os.Stderr
err := rawCmd.Start()
if err != nil {
return err
}
// Set up to capture Ctrl+C
sigChan := make(chan os.Signal, 2)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
doneChan := make(chan struct{}, 2)
// Run Wait() in its own chan so we don't block
go func() {
err = rawCmd.Wait()
doneChan <- struct{}{}
}()
// Here we block until command is done
for {
select {
case s := <-sigChan:
// user typed Ctrl-C, most likley meant for ssm-session pass through
rawCmd.Process.Signal(s)
case <-doneChan:
// command is done
return err
}
}
return err
}
func attachTmuxSession(sessionName string) (err error) {
// If we don't redirect these, our console will detach when ssm-session finishes executing.
rawCmd := exec.Command("tmux", "attach", "-t", sessionName)
rawCmd.Stdin = os.Stdin
rawCmd.Stdout = os.Stdout
rawCmd.Stderr = os.Stderr
return rawCmd.Run()
}
func addInstanceToTmuxWindow(tmuxWindow *gomux.Window, profile string, region string, instanceID string) (err error) {
tPane, err := tmuxWindow.Pane(0).Split()
if err != nil {
return err
}
if err = tPane.SetName(instanceID); err != nil {
return err
}
return tPane.Exec(fmt.Sprintf("aws ssm start-session --profile %s --region %s --target %s", profile, region, instanceID))
}
func startSelectionPrompt(instances *instance.InstanceInfoSafe, totalInstances int, tags ssmhelpers.ListSlice) (selectedInstances []instance.InstanceInfo, err error) {
instanceIDList := []string{}
promptList := instances.FormatStringSlice([]string(tags)...)
fmt.Println(" ", promptList[0])
prompt := &survey.MultiSelect{
Message: fmt.Sprintf("Showing %d/%d instances. Make a Selection:", len(instances.AllInstances), totalInstances),
Options: promptList[1 : len(promptList)-1],
}
if err := survey.AskOne(prompt, &instanceIDList, survey.WithPageSize(25)); err != nil {
return nil, err
}
if len(instanceIDList) == 0 {
return nil, fmt.Errorf("No instances selected")
}
// This is clunky, but currently necessary in order to finish creating the tmux sessions
for _, v := range instanceIDList {
id := strings.Split(v, " ")[0]
if id != "" {
selectedInstances = append(selectedInstances, instances.AllInstances[id])
}
}
return selectedInstances, nil
}