From 6c51ab7609a4b204cd9d8aab2a4f4453d38253cb Mon Sep 17 00:00:00 2001 From: Ezbon Jacob Date: Thu, 25 May 2023 17:06:32 -0500 Subject: [PATCH 1/2] feat: add option to filter instance by specific tag key in ask --- cmd/cmd.go | 5 ++++- cmd/fwd.go | 7 +++++-- cmd/fwdrem.go | 7 +++++-- cmd/session.go | 9 +++++++-- cmd/ssh.go | 2 +- internal/ssm.go | 21 +++++++++++++++------ internal/ssm_test.go | 2 +- 7 files changed, 38 insertions(+), 15 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index a6c8a73..5783598 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -34,8 +34,9 @@ var ( // get targets argTarget := strings.TrimSpace(viper.GetString("cmd-target")) + argTag := strings.TrimSpace(viper.GetString("cmd-filter")) if argTarget != "" { - table, err := internal.FindInstances(ctx, *_credential.awsConfig) + table, err := internal.FindInstances(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -86,9 +87,11 @@ var ( func init() { cmdCommand.Flags().StringP("exec", "e", "", "[required] execute command") cmdCommand.Flags().StringP("target", "t", "", "[optional] it is ec2 instanceId.") + cmdCommand.Flags().StringP("filter", "f", "", "[optional] which tag key to add to in table filtering.") viper.BindPFlag("cmd-exec", cmdCommand.Flags().Lookup("exec")) viper.BindPFlag("cmd-target", cmdCommand.Flags().Lookup("target")) + viper.BindPFlag("cmd-filter", startSessionCommand.Flags().Lookup("filter")) rootCmd.AddCommand(cmdCommand) } diff --git a/cmd/fwd.go b/cmd/fwd.go index bdb474b..c7dcd5f 100644 --- a/cmd/fwd.go +++ b/cmd/fwd.go @@ -30,8 +30,9 @@ var ( // get target argTarget := strings.TrimSpace(viper.GetString("fwd-target")) + argTag := strings.TrimSpace(viper.GetString("fwd-filter")) if argTarget != "" { - table, err := internal.FindInstances(ctx, *_credential.awsConfig) + table, err := internal.FindInstances(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -43,7 +44,7 @@ var ( } } if target == nil { - target, err = internal.AskTarget(ctx, *_credential.awsConfig) + target, err = internal.AskTarget(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -115,11 +116,13 @@ func init() { fwdCommand.Flags().StringP("remote", "z", "", "[optional] remote port to forward to, ex) 8080") fwdCommand.Flags().StringP("local", "l", "", "[optional] local port to use, ex) 1234") fwdCommand.Flags().StringP("target", "t", "", "[optional] it is ec2 instanceId.") + fwdCommand.Flags().StringP("filter", "f", "", "[optional] which tag key to add to in table filtering.") // mapping viper viper.BindPFlag("fwd-remote-port", fwdCommand.Flags().Lookup("remote")) viper.BindPFlag("fwd-local-port", fwdCommand.Flags().Lookup("local")) viper.BindPFlag("fwd-target", fwdCommand.Flags().Lookup("target")) + viper.BindPFlag("fwd-filter", fwdCommand.Flags().Lookup("filter")) rootCmd.AddCommand(fwdCommand) } diff --git a/cmd/fwdrem.go b/cmd/fwdrem.go index dc08c4a..cbef8f9 100644 --- a/cmd/fwdrem.go +++ b/cmd/fwdrem.go @@ -31,8 +31,9 @@ var ( // get target argTarget := strings.TrimSpace(viper.GetString("fwd-target")) + argTag := strings.TrimSpace(viper.GetString("fwd-filter")) if argTarget != "" { - table, err := internal.FindInstances(ctx, *_credential.awsConfig) + table, err := internal.FindInstances(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -44,7 +45,7 @@ var ( } } if target == nil { - target, err = internal.AskTarget(ctx, *_credential.awsConfig) + target, err = internal.AskTarget(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -131,12 +132,14 @@ func init() { fwdremCommand.Flags().StringP("local", "l", "", "[optional] local port to use, ex) 1234") fwdremCommand.Flags().StringP("target", "t", "", "[optional] it is ec2 instanceId to proxy through.") fwdremCommand.Flags().StringP("host", "a", "", "[optional] it is remote host address to proxy to.") + fwdremCommand.Flags().StringP("filter", "f", "", "[optional] which tag key to add to in table filtering.") // mapping viper viper.BindPFlag("fwd-remote-port", fwdremCommand.Flags().Lookup("remote")) viper.BindPFlag("fwd-local-port", fwdremCommand.Flags().Lookup("local")) viper.BindPFlag("fwd-target", fwdremCommand.Flags().Lookup("target")) viper.BindPFlag("fwd-host", fwdremCommand.Flags().Lookup("host")) + viper.BindPFlag("fwd-filter", fwdremCommand.Flags().Lookup("filter")) rootCmd.AddCommand(fwdremCommand) } diff --git a/cmd/session.go b/cmd/session.go index 77dabb2..063287b 100644 --- a/cmd/session.go +++ b/cmd/session.go @@ -27,8 +27,10 @@ var ( // get target argTarget := strings.TrimSpace(viper.GetString("start-session-target")) + argTag := strings.TrimSpace(viper.GetString("start-session-filter")) + if argTarget != "" { - table, err := internal.FindInstances(ctx, *_credential.awsConfig) + table, err := internal.FindInstances(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -40,7 +42,7 @@ var ( } } if target == nil { - target, err = internal.AskTarget(ctx, *_credential.awsConfig) + target, err = internal.AskTarget(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -80,7 +82,10 @@ var ( func init() { startSessionCommand.Flags().StringP("target", "t", "", "[optional] it is ec2 instanceId.") + startSessionCommand.Flags().StringP("filter", "f", "", "[optional] which tag key to add to in table filtering.") + viper.BindPFlag("start-session-target", startSessionCommand.Flags().Lookup("target")) + viper.BindPFlag("start-session-filter", startSessionCommand.Flags().Lookup("filter")) // add sub command rootCmd.AddCommand(startSessionCommand) diff --git a/cmd/ssh.go b/cmd/ssh.go index 652032d..364b8c0 100644 --- a/cmd/ssh.go +++ b/cmd/ssh.go @@ -32,7 +32,7 @@ var ( var sshCommand string var targetName string if exec == "" { - target, err := internal.AskTarget(ctx, *_credential.awsConfig) + target, err := internal.AskTarget(ctx, *_credential.awsConfig, "") if err != nil { panicRed(err) } diff --git a/internal/ssm.go b/internal/ssm.go index dba8457..e1a3585 100644 --- a/internal/ssm.go +++ b/internal/ssm.go @@ -110,8 +110,8 @@ func AskRegion(ctx context.Context, cfg aws.Config) (*Region, error) { } // AskTarget asks you which selects an instance. -func AskTarget(ctx context.Context, cfg aws.Config) (*Target, error) { - table, err := FindInstances(ctx, cfg) +func AskTarget(ctx context.Context, cfg aws.Config, tag string) (*Target, error) { + table, err := FindInstances(ctx, cfg, tag) if err != nil { return nil, err } @@ -142,7 +142,7 @@ func AskTarget(ctx context.Context, cfg aws.Config) (*Target, error) { // AskMultiTarget asks you which selects multi targets. func AskMultiTarget(ctx context.Context, cfg aws.Config) ([]*Target, error) { - table, err := FindInstances(ctx, cfg) + table, err := FindInstances(ctx, cfg, "") if err != nil { return nil, err } @@ -207,7 +207,7 @@ func AskPorts() (port *Port, retErr error) { } // FindInstances returns all of instances-map with running state. -func FindInstances(ctx context.Context, cfg aws.Config) (map[string]*Target, error) { +func FindInstances(ctx context.Context, cfg aws.Config, filterTag string) (map[string]*Target, error) { var ( client = ec2.NewFromConfig(cfg) table = make(map[string]*Target) @@ -215,13 +215,22 @@ func FindInstances(ctx context.Context, cfg aws.Config) (map[string]*Target, err for _, rv := range output.Reservations { for _, inst := range rv.Instances { name := "" + tags := "" for _, tag := range inst.Tags { if aws.ToString(tag.Key) == "Name" { name = aws.ToString(tag.Value) - break + } else { + if len(filterTag) > 0 { + if aws.ToString(tag.Key) == filterTag { + tags += fmt.Sprintf("%s:%s, ", aws.ToString(tag.Key), aws.ToString(tag.Value)) + } + } } } - table[fmt.Sprintf("%s\t(%s)", name, *inst.InstanceId)] = &Target{ + if len(tags) > 0 { + tags = strings.TrimSuffix(tags, ", ") + } + table[fmt.Sprintf("%s\t(%s)\t(%s)", name, *inst.InstanceId, tags)] = &Target{ Name: aws.ToString(inst.InstanceId), PublicDomain: aws.ToString(inst.PublicDnsName), PrivateDomain: aws.ToString(inst.PrivateDnsName), diff --git a/internal/ssm_test.go b/internal/ssm_test.go index 89ab584..0e4ac71 100644 --- a/internal/ssm_test.go +++ b/internal/ssm_test.go @@ -28,7 +28,7 @@ func TestFindInstances(t *testing.T) { } for _, t := range tests { - result, err := FindInstances(t.ctx, t.cfg) + result, err := FindInstances(t.ctx, t.cfg, "") assert.Equal(t.isErr, err != nil) fmt.Println(len(result)) } From 290aca388e1755cc17b1f1908ce0d617bee3ebc5 Mon Sep 17 00:00:00 2001 From: Ezbon Jacob Date: Thu, 25 May 2023 19:05:36 -0500 Subject: [PATCH 2/2] add tag support for cmd interactive table --- cmd/cmd.go | 5 +++-- internal/ssm.go | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 5783598..c94297b 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -35,6 +35,7 @@ var ( // get targets argTarget := strings.TrimSpace(viper.GetString("cmd-target")) argTag := strings.TrimSpace(viper.GetString("cmd-filter")) + if argTarget != "" { table, err := internal.FindInstances(ctx, *_credential.awsConfig, argTag) if err != nil { @@ -49,7 +50,7 @@ var ( } if len(targets) == 0 { - targets, err = internal.AskMultiTarget(ctx, *_credential.awsConfig) + targets, err = internal.AskMultiTarget(ctx, *_credential.awsConfig, argTag) if err != nil { panicRed(err) } @@ -91,7 +92,7 @@ func init() { viper.BindPFlag("cmd-exec", cmdCommand.Flags().Lookup("exec")) viper.BindPFlag("cmd-target", cmdCommand.Flags().Lookup("target")) - viper.BindPFlag("cmd-filter", startSessionCommand.Flags().Lookup("filter")) + viper.BindPFlag("cmd-filter", cmdCommand.Flags().Lookup("filter")) rootCmd.AddCommand(cmdCommand) } diff --git a/internal/ssm.go b/internal/ssm.go index e1a3585..92a59d7 100644 --- a/internal/ssm.go +++ b/internal/ssm.go @@ -141,8 +141,8 @@ func AskTarget(ctx context.Context, cfg aws.Config, tag string) (*Target, error) } // AskMultiTarget asks you which selects multi targets. -func AskMultiTarget(ctx context.Context, cfg aws.Config) ([]*Target, error) { - table, err := FindInstances(ctx, cfg, "") +func AskMultiTarget(ctx context.Context, cfg aws.Config, tag string) ([]*Target, error) { + table, err := FindInstances(ctx, cfg, tag) if err != nil { return nil, err }