Add support for filtering list by tag values.

Author: itsdalmo

README.md

@@ -65,6 +65,7 @@ $ ssm-sh list --help
 
 ...
 [list command options]
+      -f, --filter= Filter the produced list by tag (key=value,..)
       -l, --limit=  Limit the number of instances printed (default: 50)
       -o, --output= Path to a file where the list of instances will be written as YAML.[list command options]
 ```
@@ -84,11 +85,11 @@ $ ssm-sh run --help
 ### Example
 
 ```bash
-$ vaulted -n lab-admin -- ssm-sh list -o example.json
+$ vaulted -n lab-admin -- ssm-sh list --filter Name="*itsdalmo" -o example.json
 
 Instance ID         | Name                             | State   | Image ID     | Platform     | Version | IP            | Status | Last pinged
-i-03762678c45546813 | ssm-manager-manual-test-kristian | running | ami-db1688a2 | Amazon Linux | 2.0     | 172.53.17.163 | Online | 2018-02-09 12:37
-i-0d04464ff18b5db7d | ssm-manager-manual-test-kristian | running | ami-db1688a2 | Amazon Linux | 2.0     | 172.53.20.172 | Online | 2018-02-09 12:39
+i-03762678c45546813 | ssm-manager-manual-test-itsdalmo | running | ami-db1688a2 | Amazon Linux | 2.0     | 172.53.17.163 | Online | 2018-02-09 12:37
+i-0d04464ff18b5db7d | ssm-manager-manual-test-itsdalmo | running | ami-db1688a2 | Amazon Linux | 2.0     | 172.53.20.172 | Online | 2018-02-09 12:39
 
 $ vaulted -n lab-admin -- ssm-sh shell --target-file example.json
 Initialized with targets: [i-03762678c45546813 i-0d04464ff18b5db7d]

command/list.go

@@ -6,9 +6,11 @@ import (
 	"github.com/pkg/errors"
 	"io/ioutil"
 	"os"
+	"strings"
 )
 
 type ListCommand struct {
+	Tags   []*tag `short:"f" long:"filter" description:"Filter the produced list by tag (key=value,..)"`
 	Limit  int64  `short:"l" long:"limit" description:"Limit the number of instances printed" default:"50"`
 	Output string `short:"o" long:"output" description:"Path to a file where the list of instances will be written as JSON."`
 }
@@ -20,7 +22,14 @@ func (command *ListCommand) Execute([]string) error {
 	}
 	m := manager.NewManager(sess, Command.AwsOpts.Region)
 
-	instances, err := m.ListInstances(command.Limit)
+	var filters []*manager.TagFilter
+	for _, tag := range command.Tags {
+		filters = append(filters, &manager.TagFilter{
+			Key:    tag.Key,
+			Values: tag.Values,
+		})
+	}
+	instances, err := m.ListInstances(command.Limit, filters)
 	if err != nil {
 		return errors.Wrap(err, "failed to list instances")
 	}
@@ -41,3 +50,22 @@ func (command *ListCommand) Execute([]string) error {
 
 	return nil
 }
+
+type tag manager.TagFilter
+
+func (t *tag) UnmarshalFlag(value string) error {
+	parts := strings.Split(value, "=")
+	if len(parts) != 2 {
+		return errors.New("expected a key and a value separated by =")
+	}
+
+	values := strings.Split(parts[1], ",")
+	if len(values) < 1 {
+		return errors.New("expected one or more values separated by ,")
+	}
+
+	t.Key = parts[0]
+	t.Values = values
+
+	return nil
+}

manager/manager.go

@@ -18,6 +18,28 @@ import (
 	"time"
 )
 
+// TagFilter represents a key=value pair for AWS EC2 tags.
+type TagFilter struct {
+	Key    string
+	Values []string
+}
+
+// Filter returns the ec2.Filter representation of the TagFilter.
+func (t *TagFilter) Filter() *ec2.Filter {
+	return &ec2.Filter{
+		Name:   aws.String(fmt.Sprintf("tag:%s", t.Key)),
+		Values: aws.StringSlice(t.Values),
+	}
+}
+
+// CommandOutput is the return type transmitted over a channel when fetching output.
+type CommandOutput struct {
+	InstanceID string
+	Status     string
+	Output     string
+	Error      error
+}
+
 // Manager handles the clients interfacing with AWS.
 type Manager struct {
 	ssmClient ssmiface.SSMAPI
@@ -48,7 +70,7 @@ func NewTestManager(ssm ssmiface.SSMAPI, s3 s3iface.S3API, ec2 ec2iface.EC2API)
 }
 
 // ListInstances fetches a list of instances managed by SSM. Paginates until all responses have been collected.
-func (m *Manager) ListInstances(limit int64) ([]*Instance, error) {
+func (m *Manager) ListInstances(limit int64, tagFilters []*TagFilter) ([]*Instance, error) {
 	var out []*Instance
 
 	input := &ssm.DescribeInstanceInformationInput{
@@ -60,7 +82,7 @@ func (m *Manager) ListInstances(limit int64) ([]*Instance, error) {
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to describe instance information")
 		}
-		ssmInstances, ec2Instances, err := m.describeInstances(response.InstanceInformationList)
+		ssmInstances, ec2Instances, err := m.describeInstances(response.InstanceInformationList, tagFilters)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to retrieve ec2 instance information")
 		}
@@ -79,25 +101,29 @@ func (m *Manager) ListInstances(limit int64) ([]*Instance, error) {
 }
 
 // describeInstances retrieves additional information about SSM managed instances from EC2.
-func (m *Manager) describeInstances(instances []*ssm.InstanceInformation) (map[string]*ssm.InstanceInformation, map[string]*ec2.Instance, error) {
+func (m *Manager) describeInstances(instances []*ssm.InstanceInformation, tagFilters []*TagFilter) (map[string]*ssm.InstanceInformation, map[string]*ec2.Instance, error) {
 	var ids []*string
+	var filters []*ec2.Filter
 
 	org := make(map[string]*ssm.InstanceInformation)
 	out := make(map[string]*ec2.Instance)
 
 	for _, instance := range instances {
-		id := aws.StringValue(instance.InstanceId)
-		org[id] = instance
+		org[aws.StringValue(instance.InstanceId)] = instance
 		ids = append(ids, instance.InstanceId)
 	}
 
+	filters = append(filters, &ec2.Filter{
+		Name:   aws.String("instance-id"),
+		Values: ids,
+	})
+
+	for _, f := range tagFilters {
+		filters = append(filters, f.Filter())
+	}
+
 	input := &ec2.DescribeInstancesInput{
-		Filters: []*ec2.Filter{
-			{
-				Name:   aws.String("instance-id"),
-				Values: ids,
-			},
-		},
+		Filters: filters,
 	}
 
 	for {
@@ -149,14 +175,6 @@ func (m *Manager) AbortCommand(instanceIds []string, commandID string) error {
 	return nil
 }
 
-// CommandOutput is the return type transmitted over a channel when fetching output.
-type CommandOutput struct {
-	InstanceID string
-	Status     string
-	Output     string
-	Error      error
-}
-
 // GetCommandOutput fetches the results from a command invocation for all specified instanceIds and
 // closes the receiving channel before exiting.
 func (m *Manager) GetCommandOutput(ctx context.Context, instanceIds []string, commandID string, out chan<- *CommandOutput) {

manager/manager_test.go

@@ -105,15 +105,15 @@ func TestList(t *testing.T) {
 
 	t.Run("Get managed instances works", func(t *testing.T) {
 		expected := outputInstances
-		actual, err := m.ListInstances(50)
+		actual, err := m.ListInstances(50, nil)
 		assert.Nil(t, err)
 		assert.NotNil(t, actual)
 		assert.Equal(t, expected, actual)
 	})
 
 	t.Run("Limit number of instances works", func(t *testing.T) {
 		expected := outputInstances[:1]
-		actual, err := m.ListInstances(1)
+		actual, err := m.ListInstances(1, nil)
 		assert.Nil(t, err)
 		assert.NotNil(t, actual)
 		assert.Equal(t, expected, actual)
@@ -126,7 +126,22 @@ func TestList(t *testing.T) {
 		}()
 
 		expected := outputInstances
-		actual, err := m.ListInstances(50)
+		actual, err := m.ListInstances(50, nil)
+		assert.Nil(t, err)
+		assert.NotNil(t, actual)
+		assert.Equal(t, expected, actual)
+	})
+
+	t.Run("TagFilter works", func(t *testing.T) {
+		expected := outputInstances[:1]
+		actual, err := m.ListInstances(50, []*manager.TagFilter{
+			{
+				Key: "Name",
+				Values: []string{
+					"1",
+				},
+			},
+		})
 		assert.Nil(t, err)
 		assert.NotNil(t, actual)
 		assert.Equal(t, expected, actual)
@@ -138,7 +153,7 @@ func TestList(t *testing.T) {
 			ssmMock.Error = false
 		}()
 
-		actual, err := m.ListInstances(50)
+		actual, err := m.ListInstances(50, nil)
 		assert.NotNil(t, err)
 		assert.EqualError(t, err, "failed to describe instance information: expected")
 		assert.Nil(t, actual)

manager/testing.go

@@ -26,20 +26,59 @@ func (mock *MockEC2) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.
 		return nil, errors.New("expected")
 	}
 
+	var out []*ec2.Instance
+	var tmp []*ec2.Instance
 	var ids []string
+	var nameFilters []string
+
 	for _, filter := range input.Filters {
-		if aws.StringValue(filter.Name) == "instance-id" {
+		key := aws.StringValue(filter.Name)
+		if key == "instance-id" {
 			ids = aws.StringValueSlice(filter.Values)
+		} else if key == "tag:Name" {
+			nameFilters = aws.StringValueSlice(filter.Values)
 		}
 	}
 
-	if ids == nil {
-		return nil, errors.New("missing instance ids in input")
+	// Filter instance ids if a list was provided. If not, we
+	// provide the entire list of instances.
+	if ids != nil {
+		for _, id := range ids {
+			instance, ok := mock.Instances[id]
+			if !ok {
+				return nil, errors.New("instance id does not exist")
+			}
+			tmp = append(tmp, instance)
+		}
+
+	} else {
+		for _, instance := range mock.Instances {
+			tmp = append(tmp, instance)
+		}
 	}
 
-	var out []*ec2.Instance
-	for _, id := range ids {
-		out = append(out, mock.Instances[id])
+	// If a tag filter was supplied (only Name is supported for testing),
+	// filter instances which don't match.
+	if nameFilters != nil && len(nameFilters) > 0 {
+		for _, instance := range tmp {
+			for _, tag := range instance.Tags {
+				// Look for Name tag.
+				if aws.StringValue(tag.Key) != "Name" {
+					continue
+				}
+				// Once it is found, check whether it contains
+				// any of the name filters (simple contains).
+				name := aws.StringValue(tag.Value)
+				for _, filter := range nameFilters {
+					if strings.Contains(name, filter) {
+						out = append(out, instance)
+					}
+				}
+			}
+		}
+
+	} else {
+		out = tmp
 	}
 
 	// NOTE: It should not matter if we have multiple reservations