Merge pull request #12 from mumoshu/optional-s3-output

Support S3 bucket and key for run/shell output

Author: Kristian

command/list.go

@@ -2,11 +2,12 @@ package command
 
 import (
 	"encoding/json"
-	"github.com/itsdalmo/ssm-sh/manager"
-	"github.com/pkg/errors"
 	"io/ioutil"
 	"os"
 	"strings"
+
+	"github.com/itsdalmo/ssm-sh/manager"
+	"github.com/pkg/errors"
 )
 
 type ListCommand struct {
@@ -20,7 +21,7 @@ func (command *ListCommand) Execute([]string) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to create new session")
 	}
-	m := manager.NewManager(sess, Command.AwsOpts.Region)
+	m := manager.NewManager(sess, Command.AwsOpts.Region, manager.Opts{})
 
 	var filters []*manager.TagFilter
 	for _, tag := range command.Tags {

command/run.go

@@ -3,15 +3,17 @@ package command
 import (
 	"context"
 	"fmt"
-	"github.com/itsdalmo/ssm-sh/manager"
-	"github.com/pkg/errors"
 	"os"
 	"strings"
 	"time"
+
+	"github.com/itsdalmo/ssm-sh/manager"
+	"github.com/pkg/errors"
 )
 
 type RunCommand struct {
-	Timeout    int `short:"i" long:"timeout" description:"Seconds to wait for command result before timing out." default:"30"`
+	Timeout    int        `short:"i" long:"timeout" description:"Seconds to wait for command result before timing out." default:"30"`
+	SSMOpts    SSMOptions `group:"SSM options"`
 	TargetOpts TargetOptions
 }
 
@@ -21,7 +23,11 @@ func (command *RunCommand) Execute(args []string) error {
 		return errors.Wrap(err, "failed to create new aws session")
 	}
 
-	m := manager.NewManager(sess, Command.AwsOpts.Region)
+	opts, err := command.SSMOpts.Parse()
+	if err != nil {
+		return err
+	}
+	m := manager.NewManager(sess, Command.AwsOpts.Region, *opts)
 	targets, err := setTargets(command.TargetOpts)
 	if err != nil {
 		return errors.Wrap(err, "failed to set targets")

command/shell.go

@@ -3,15 +3,17 @@ package command
 import (
 	"context"
 	"fmt"
-	"github.com/chzyer/readline"
-	"github.com/itsdalmo/ssm-sh/manager"
-	"github.com/pkg/errors"
 	"io"
 	"os"
 	"strings"
+
+	"github.com/chzyer/readline"
+	"github.com/itsdalmo/ssm-sh/manager"
+	"github.com/pkg/errors"
 )
 
 type ShellCommand struct {
+	SSMOpts    SSMOptions `group:"SSM options"`
 	TargetOpts TargetOptions
 }
 
@@ -21,7 +23,11 @@ func (command *ShellCommand) Execute([]string) error {
 		return errors.Wrap(err, "failed to create new aws session")
 	}
 
-	m := manager.NewManager(sess, Command.AwsOpts.Region)
+	opts, err := command.SSMOpts.Parse()
+	if err != nil {
+		return err
+	}
+	m := manager.NewManager(sess, Command.AwsOpts.Region, *opts)
 	targets, err := setTargets(command.TargetOpts)
 	if err != nil {
 		return errors.Wrap(err, "failed to set targets")

command/ssm_opts.go

@@ -0,0 +1,32 @@
+package command
+
+import (
+	"fmt"
+
+	"github.com/itsdalmo/ssm-sh/manager"
+)
+
+type SSMOptions struct {
+	ExtendOutput bool   `short:"x" long:"extend-output" description:"Extend truncated command outputs by fetching S3 objects containing full ones"`
+	S3Bucket     string `short:"b" long:"s3-bucket" description:"S3 bucket in which S3 objects containing full command outputs are stored. Required when --extend-output is provided." default:""`
+	S3KeyPrefix  string `short:"k" long:"s3-key-prefix" description:"Key prefix of S3 objects containing full command outputs." default:""`
+}
+
+func (o SSMOptions) Validate() error {
+	if o.ExtendOutput && o.S3Bucket == "" {
+		return fmt.Errorf("--s3-bucket must be a non-empty string when --extend-output is provided")
+	}
+	return nil
+}
+
+func (o SSMOptions) Parse() (*manager.Opts, error) {
+	err := o.Validate()
+	if err != nil {
+		return nil, err
+	}
+	return &manager.Opts{
+		ExtendOutput: o.ExtendOutput,
+		S3Bucket:     o.S3Bucket,
+		S3KeyPrefix:  o.S3KeyPrefix,
+	}, nil
+}

command/utils.go

@@ -4,16 +4,17 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/fatih/color"
-	"github.com/itsdalmo/ssm-sh/manager"
 	"io"
 	"io/ioutil"
 	"os"
 	"os/signal"
 	"strings"
 	"text/tabwriter"
 	"time"
+
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/fatih/color"
+	"github.com/itsdalmo/ssm-sh/manager"
 )
 
 // Create a new AWS session
@@ -103,6 +104,11 @@ func PrintCommandOutput(wrt io.Writer, output *manager.CommandOutput) error {
 	if _, err := fmt.Fprintf(wrt, "%s\n", output.Output); err != nil {
 		return err
 	}
+	if output.OutputUrl != "" {
+		if _, err := fmt.Fprintf(wrt, "(Output URL: %s)\n", output.OutputUrl); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

command/utils_test.go

@@ -3,12 +3,13 @@ package command_test
 import (
 	"bytes"
 	"errors"
-	"github.com/itsdalmo/ssm-sh/command"
-	"github.com/itsdalmo/ssm-sh/manager"
-	"github.com/stretchr/testify/assert"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/itsdalmo/ssm-sh/command"
+	"github.com/itsdalmo/ssm-sh/manager"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPrintInstances(t *testing.T) {
@@ -61,6 +62,13 @@ func TestPrintCommandOutput(t *testing.T) {
 			Output:     "Standard output",
 			Error:      nil,
 		},
+		{
+			InstanceID: "i-00000000000000001",
+			Status:     "Success",
+			Output:     "Extended standard output",
+			OutputUrl:  "https://s3-ap-northeast-1.amazonaws.com/mybucket/foobar/c0896747-af2b-4359-bc34-0f951ce02007/i-00000000000000001/awsrunShellScript/0.awsrunShellScript/stdout",
+			Error:      nil,
+		},
 		{
 			InstanceID: "i-00000000000000002",
 			Status:     "Failed",
@@ -80,6 +88,10 @@ func TestPrintCommandOutput(t *testing.T) {
 i-00000000000000001 - Success:
 Standard output
 
+i-00000000000000001 - Success:
+Extended standard output
+(Output URL: https://s3-ap-northeast-1.amazonaws.com/mybucket/foobar/c0896747-af2b-4359-bc34-0f951ce02007/i-00000000000000001/awsrunShellScript/0.awsrunShellScript/stdout)
+
 i-00000000000000002 - Failed:
 Standard error
 

main.go

@@ -1,9 +1,10 @@
 package main
 
 import (
+	"os"
+
 	"github.com/itsdalmo/ssm-sh/command"
 	"github.com/jessevdk/go-flags"
-	"os"
 )
 
 // Version is set on build by the Git release tag.

manager/manager.go

@@ -4,6 +4,12 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
+	"regexp"
+	"strings"
+	"sync"
+	"time"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -13,9 +19,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
 	"github.com/pkg/errors"
-	"io"
-	"sync"
-	"time"
 )
 
 // TagFilter represents a key=value pair for AWS EC2 tags.
@@ -37,26 +40,40 @@ type CommandOutput struct {
 	InstanceID string
 	Status     string
 	Output     string
+	OutputUrl  string
 	Error      error
 }
 
 // Manager handles the clients interfacing with AWS.
 type Manager struct {
-	ssmClient ssmiface.SSMAPI
-	s3Client  s3iface.S3API
-	ec2Client ec2iface.EC2API
-	region    string
+	ssmClient    ssmiface.SSMAPI
+	s3Client     s3iface.S3API
+	ec2Client    ec2iface.EC2API
+	extendOutput bool
+	region       string
+	s3Bucket     string
+	s3KeyPrefix  string
+}
+
+type Opts struct {
+	ExtendOutput bool
+	S3Bucket     string
+	S3KeyPrefix  string
 }
 
 // NewManager creates a new Manager from an AWS session and region.
-func NewManager(sess *session.Session, region string) *Manager {
-	config := &aws.Config{Region: aws.String(region)}
-	return &Manager{
-		ssmClient: ssm.New(sess, config),
-		s3Client:  s3.New(sess, config),
-		ec2Client: ec2.New(sess, config),
+func NewManager(sess *session.Session, region string, opts Opts) *Manager {
+	awsCfg := &aws.Config{Region: aws.String(region)}
+	m := &Manager{
+		ssmClient: ssm.New(sess, awsCfg),
+		s3Client:  s3.New(sess, awsCfg),
+		ec2Client: ec2.New(sess, awsCfg),
 		region:    region,
 	}
+	m.extendOutput = opts.ExtendOutput
+	m.s3Bucket = opts.S3Bucket
+	m.s3KeyPrefix = opts.S3KeyPrefix
+	return m
 }
 
 // NewTestManager creates a new manager for testing purposes.
@@ -154,7 +171,12 @@ func (m *Manager) RunCommand(instanceIds []string, command string) (string, erro
 		Comment:      aws.String("Interactive command."),
 		Parameters:   map[string][]*string{"commands": {aws.String(command)}},
 	}
-
+	if m.s3Bucket != "" {
+		input.OutputS3BucketName = aws.String(m.s3Bucket)
+	}
+	if m.s3KeyPrefix != "" {
+		input.OutputS3KeyPrefix = aws.String(m.s3KeyPrefix)
+	}
 	res, err := m.ssmClient.SendCommand(input)
 	if err != nil {
 		return "", err
@@ -206,15 +228,15 @@ func (m *Manager) pollInstanceOutput(ctx context.Context, instanceID string, com
 				CommandId:  aws.String(commandID),
 				InstanceId: aws.String(instanceID),
 			})
-			if out, ok := newCommandOutput(result, err); ok {
+			if out, ok := m.newCommandOutput(result, err); ok {
 				c <- out
 				return
 			}
 		}
 	}
 }
 
-func newCommandOutput(result *ssm.GetCommandInvocationOutput, err error) (*CommandOutput, bool) {
+func (m *Manager) newCommandOutput(result *ssm.GetCommandInvocationOutput, err error) (*CommandOutput, bool) {
 	out := &CommandOutput{
 		InstanceID: aws.StringValue(result.InstanceId),
 		Status:     aws.StringValue(result.StatusDetails),
@@ -234,17 +256,53 @@ func newCommandOutput(result *ssm.GetCommandInvocationOutput, err error) (*Comma
 		return out, true
 	case "Success":
 		out.Output = aws.StringValue(result.StandardOutputContent)
+		out.OutputUrl = aws.StringValue(result.StandardOutputUrl)
+		if m.extendOutput {
+			return m.extendTruncatedOutput(*out), true
+		}
 		return out, true
 	case "Failed":
 		out.Output = aws.StringValue(result.StandardErrorContent)
+		out.OutputUrl = aws.StringValue(result.StandardErrorUrl)
+		if m.extendOutput {
+			return m.extendTruncatedOutput(*out), true
+		}
 		return out, true
 	default:
 		out.Error = fmt.Errorf("Unrecoverable status: %s", out.Status)
 		return out, true
 	}
 }
 
-func (m *Manager) readS3Output(bucket, key string) (string, error) {
+func (m *Manager) extendTruncatedOutput(out CommandOutput) *CommandOutput {
+	const truncationMarker = "--output truncated--"
+	if strings.Contains(out.Output, truncationMarker) {
+		s3out, err := m.readOutput(out.OutputUrl)
+		if err != nil {
+			out.Error = errors.Wrap(err, "failed to fetch extended output")
+		}
+		out.Output = s3out
+		return &out
+	}
+	return &out
+}
+
+func (m *Manager) readOutput(url string) (string, error) {
+	regex := regexp.MustCompile(`://s3[\-a-z0-9]*\.amazonaws.com/([^/]+)/(.+)|://([^.]+)\.s3\.amazonaws\.com/(.+)`)
+	matches := regex.FindStringSubmatch(url)
+	if len(matches) == 0 {
+		return "", errors.Errorf("failed due to unexpected s3 url pattern: %s", url)
+	}
+	bucket := matches[1]
+	key := matches[2]
+	out, err := m.readS3Object(bucket, key)
+	if err != nil {
+		return "", errors.Wrapf(err, "failed to fetch s3 object: %s", url)
+	}
+	return out, nil
+}
+
+func (m *Manager) readS3Object(bucket, key string) (string, error) {
 	output, err := m.s3Client.GetObject(&s3.GetObjectInput{
 		Bucket: aws.String(bucket),
 		Key:    aws.String(key),