Skip to content

Commit

Permalink
Refactor command methods.
Browse files Browse the repository at this point in the history
  • Loading branch information
@gjbae1212
gjbae1212 committed Aug 31, 2021
1 parent 7169dab commit 1d13c81
Show file tree
Hide file tree
Showing 12 changed files with 518 additions and 1,039 deletions.
79 changes: 52 additions & 27 deletions cmd/cmd.go
View file
Original file line number Diff line number Diff line change
@@ -1,69 +1,94 @@
package cmd

import (
"context"
"fmt"
"strings"
"time"

"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/ssm"
"github.com/fatih/color"

"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ssm"
"github.com/gjbae1212/gossm/internal"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)

var (
// cmdCommand is a AWS Systems Manager Run Command.
// cmdCommand is AWS Systems Manager Run Command.
cmdCommand = &cobra.Command{
Use: "cmd",
Short: "Exec `run command` under AWS SSM with interactive CLI",
Long: "Exec `run command` under AWS SSM with interactive CLI",
PreRun: func(cmd *cobra.Command, args []string) {
initCredential()
executor.execCommand = viper.GetString("cmd-exec")
Run: func(cmd *cobra.Command, args []string) {
var (
targets []*internal.Target
err error
)
ctx := context.Background()

// set region
if err := setRegion(credential); err != nil {
panicRed(err)
exec := strings.TrimSpace(viper.GetString("cmd-exec"))
if exec == "" {
panicRed(fmt.Errorf("[err] not found exec command"))
}

// set multi target
if err := setMultiTarget(credential, executor); err != nil {
panicRed(err)
// get targets
argTarget := strings.TrimSpace(viper.GetString("cmd-target"))
if argTarget != "" {
table, err := internal.FindInstances(ctx, *_credential.awsConfig)
if err != nil {
panicRed(err)
}
for _, t := range table {
if t.Name == argTarget {
targets = append(targets, t)
break
}
}
}
},
Run: func(cmd *cobra.Command, args []string) {
// send command
cmdOutput, err := sendCommand(credential.awsSession, credential.awsRegion, executor.multiTarget, executor.execCommand)

if len(targets) == 0 {
targets, err = internal.AskMultiTarget(ctx, *_credential.awsConfig)
if err != nil {
panicRed(err)
}
}

var targetName string
for _, t := range targets {
targetName += " " + t.Name + " "
}

internal.PrintReady(exec, _credential.awsConfig.Region, targetName)

sendOutput, err := internal.SendCommand(ctx, *_credential.awsConfig, targets, exec)
if err != nil {
panicRed(err)
}

commandId := *cmdOutput.Command.CommandId
instanceIds := aws.StringValueSlice(cmdOutput.Command.InstanceIds)
fmt.Printf("command: %s, commandId: %s, targets: %s\n",
color.GreenString(executor.execCommand), color.GreenString(commandId),
color.GreenString(fmt.Sprintf("%v", instanceIds)))
fmt.Printf("%s\n", color.YellowString("Waiting Response ..."))

// wait 3 seconds
time.Sleep(time.Second * 3)

// show result
var inputs []*ssm.GetCommandInvocationInput
for _, inst := range cmdOutput.Command.InstanceIds {
for _, inst := range sendOutput.Command.InstanceIds {
inputs = append(inputs, &ssm.GetCommandInvocationInput{
CommandId: cmdOutput.Command.CommandId,
InstanceId: inst,
CommandId: sendOutput.Command.CommandId,
InstanceId: aws.String(inst),
})
}
printCommandInvocation(credential.awsSession, credential.awsRegion, inputs)
internal.PrintCommandInvocation(ctx, *_credential.awsConfig, inputs)
},
}
)

func init() {
cmdCommand.Flags().StringP("exec", "e", "", "[required] execute command")
cmdCommand.Flags().StringP("target", "t", "", "[optional] it is ec2 instanceId.")

viper.BindPFlag("cmd-exec", cmdCommand.Flags().Lookup("exec"))
viper.BindPFlag("cmd-target", cmdCommand.Flags().Lookup("target"))

rootCmd.AddCommand(cmdCommand)
}
108 changes: 71 additions & 37 deletions cmd/fwd.go
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,15 @@
package cmd

import (
"context"
"encoding/json"
"fmt"
"strings"

"github.com/aws/aws-sdk-go/service/ssm"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/ssm"
"github.com/fatih/color"
"github.com/gjbae1212/gossm/internal"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)
Expand All @@ -14,47 +19,72 @@ var (
Use: "fwd",
Short: "Exec `fwd` under AWS SSM with interactive CLI",
Long: "Exec `fwd` under AWS SSM with interactive CLI",
PreRun: func(cmd *cobra.Command, args []string) {
initCredential()
// set remote and local port values from cli flags (if present)
executor.remotePort = viper.GetString("remote-port")
executor.localPort = viper.GetString("local-port")
Run: func(cmd *cobra.Command, args []string) {
ctx := context.Background()
var (
target *internal.Target
remotePort string
localPort string
err error
)

// set region
if err := setRegion(credential); err != nil {
panicRed(err)
// get target
argTarget := strings.TrimSpace(viper.GetString("fwd-target"))
if argTarget != "" {
table, err := internal.FindInstances(ctx, *_credential.awsConfig)
if err != nil {
panicRed(err)
}
for _, t := range table {
if t.Name == argTarget {
target = t
break
}
}
}

// set target
if err := setTarget(credential, executor); err != nil {
panicRed(err)
if target == nil {
target, err = internal.AskTarget(ctx, *_credential.awsConfig)
if err != nil {
panicRed(err)
}
}

// set ports
if err := setFwdPorts(credential, executor); err != nil {
panicRed(err)
// get port
argRemotePort := strings.TrimSpace(viper.GetString("fwd-remote-port"))
argLocalPort := strings.TrimSpace(viper.GetString("fwd-local-port"))
if argRemotePort == "" {
askPort, err := internal.AskPorts()
if err != nil {
panicRed(err)
}
remotePort = askPort.Remote
localPort = askPort.Local
} else {
remotePort = argRemotePort
localPort = argLocalPort
if localPort == "" {
localPort = remotePort
}
}
internal.PrintReady(fmt.Sprintf("start-port-forwarding %s -> %s", localPort, remotePort), _credential.awsConfig.Region, target.Name)

printReady("start-port-forwarding", credential, executor)
},
Run: func(cmd *cobra.Command, args []string) {
docName := "AWS-StartPortForwardingSession" // https://us-east-1.console.aws.amazon.com/systems-manager/documents/AWS-StartPortForwardingSession/description?region=us-east-1
input := &ssm.StartSessionInput{
DocumentName: &docName,
Parameters: map[string][]*string{
"portNumber": []*string{&executor.remotePort},
"localPortNumber": []*string{&executor.localPort},
Parameters: map[string][]string{
"portNumber": []string{remotePort},
"localPortNumber": []string{localPort},
},
Target: &executor.target,
Target: aws.String(target.Name),
}

// create session
sess, endpoint, err := createStartSession(credential, input)
// start session
session, err := internal.CreateStartSession(ctx, *_credential.awsConfig, input)
if err != nil {
panicRed(err)
}

sessJson, err := json.Marshal(sess)
sessJson, err := json.Marshal(session)
if err != nil {
panicRed(err)
}
Expand All @@ -64,28 +94,32 @@ var (
panicRed(err)
}

// call session-manager-plugin
if err := callSubprocess(viper.GetString("plugin"), string(sessJson),
credential.awsRegion, "StartSession", credential.awsProfile, string(paramsJson), endpoint); err != nil {
color.Red("%v", err)
if err := internal.CallProcess(_credential.ssmPluginPath, string(sessJson),
_credential.awsConfig.Region, "StartSession",
_credential.awsProfile, string(paramsJson)); err != nil {
color.Red("[err] %v", err.Error())
}

// delete Session
if err := deleteStartSession(credential, *sess.SessionId); err != nil {
color.Red("%v", err)
// delete session
if err := internal.DeleteStartSession(ctx, *_credential.awsConfig, &ssm.TerminateSessionInput{
SessionId: session.SessionId,
}); err != nil {
panicRed(err)
}
},
}
)

func init() {
// add sub command
fwdCommand.Flags().StringP("remote", "z", "", "[optional] remote port to forward to, ex) - 8080")
fwdCommand.Flags().StringP("local", "l", "", "[optional] local port to use, ex) \"-l 1234\"")
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.")

// mapping viper
viper.BindPFlag("remote-port", fwdCommand.Flags().Lookup("remote"))
viper.BindPFlag("local-port", fwdCommand.Flags().Lookup("local"))
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"))

rootCmd.AddCommand(fwdCommand)
}

0 comments on commit 1d13c81

Please sign in to comment.