allow-list and deny-list implemented (#11)

* allow-list and deny-list impl

* update usage with allow-list and deny-list

* improved test failure text

* fixed test failure text

Author: bwagner5

README.md

@@ -134,12 +134,14 @@ ec2-instance-selector --vcpus 4 --region us-east-2 --availability-zones us-east-
 ec2-instance-selector --memory-min 4096 --memory-max 8192 --vcpus-min 4 --vcpus-max 8 --region us-east-2
 
 Filter Flags:
+      --allow-list string                 List of allowed instance types to select from w/ regex syntax (Example: m[3-5]\.*)
       --availability-zone string          [DEPRECATED] use --availability-zones instead
   -z, --availability-zones strings        Availability zones or zone ids to check EC2 capacity offered in specific AZs
       --baremetal                         Bare Metal instance types (.metal instances)
   -b, --burst-support                     Burstable instance types
   -a, --cpu-architecture string           CPU architecture [x86_64, i386, or arm64]
       --current-generation                Current generation instance types (explicitly set this to false to not return current generation instance types)
+      --deny-list string                  List of instance types which should be excluded w/ regex syntax (Example: m[1-2]\.*)
   -e, --ena-support                       Instance types where ENA is supported or required
   -f, --fpga-support                      FPGA instance types
       --gpu-memory-total int              Number of GPUs' total memory in MiB (Example: 4096) (sets --gpu-memory-total-min and -max to the same value)

cmd/main.go

@@ -67,6 +67,8 @@ const (
 	currentGeneration      = "current-generation"
 	networkInterfaces      = "network-interfaces"
 	networkPerformance     = "network-performance"
+	allowList              = "allow-list"
+	denyList               = "deny-list"
 )
 
 // Aggregate Filter Flags
@@ -133,6 +135,8 @@ Full docs can be found at github.com/aws/amazon-` + binName
 	cli.BoolFlag(currentGeneration, nil, nil, "Current generation instance types (explicitly set this to false to not return current generation instance types)")
 	cli.IntMinMaxRangeFlags(networkInterfaces, nil, nil, "Number of network interfaces (ENIs) that can be attached to the instance")
 	cli.IntMinMaxRangeFlags(networkPerformance, nil, nil, "Bandwidth in Gib/s of network performance (Example: 100)")
+	cli.RegexFlag(allowList, nil, nil, "List of allowed instance types to select from w/ regex syntax (Example: m[3-5]\\.*)")
+	cli.RegexFlag(denyList, nil, nil, "List of instance types which should be excluded w/ regex syntax (Example: m[1-2]\\.*)")
 
 	// Suite Flags - higher level aggregate filters that return opinionated result
 
@@ -204,6 +208,8 @@ Full docs can be found at github.com/aws/amazon-` + binName
 		MaxResults:             cli.IntMe(flags[maxResults]),
 		NetworkInterfaces:      cli.IntRangeMe(flags[networkInterfaces]),
 		NetworkPerformance:     cli.IntRangeMe(flags[networkPerformance]),
+		AllowList:              cli.RegexMe(flags[allowList]),
+		DenyList:               cli.RegexMe(flags[denyList]),
 		InstanceTypeBase:       cli.StringMe(flags[instanceTypeBase]),
 	}
 

pkg/cli/cli_test.go

@@ -293,3 +293,23 @@ func TestParseAndValidateFlags(t *testing.T) {
 	_, err := cli.ParseAndValidateFlags()
 	h.Nok(t, err)
 }
+
+func TestParseAndValidateRegexFlag(t *testing.T) {
+	flagName := "test-regex-flag"
+	flagArg := fmt.Sprintf("--%s", flagName)
+
+	cli := getTestCLI()
+	cli.RegexFlag(flagName, nil, nil, "Test with validation")
+	os.Args = []string{"ec2-instance-selector", flagArg, "c4.*"}
+	flags, err := cli.ParseAndValidateFlags()
+	h.Ok(t, err)
+	h.Assert(t, len(flags) == 1, "1 flag should have been processed")
+	_, err = cli.ParseAndValidateFlags()
+	h.Ok(t, err)
+
+	cli = getTestCLI()
+	cli.RegexFlag(flagName, nil, nil, "Test with validation")
+	os.Args = []string{"ec2-instance-selector", flagArg, "(("}
+	_, err = cli.ParseAndValidateFlags()
+	h.Nok(t, err)
+}

pkg/cli/flags.go

@@ -21,9 +21,10 @@ func (cl *CommandLineInterface) RatioFlag(name string, shorthand *string, defaul
 	}
 	if shorthand != nil {
 		cl.Flags[name] = cl.Command.Flags().StringP(name, string(*shorthand), *defaultValue, description)
-		return nil
+	} else {
+		cl.Flags[name] = cl.Command.Flags().String(name, *defaultValue, description)
 	}
-	cl.Flags[name] = cl.Command.Flags().String(name, *defaultValue, description)
+
 	cl.validators[name] = func(val interface{}) error {
 		if val == nil {
 			return nil
@@ -62,11 +63,15 @@ func (cl *CommandLineInterface) StringFlag(name string, shorthand *string, defau
 }
 
 // StringSliceFlag creates and registers a flag accepting a list of strings.
-// Suite flags will be grouped in the middle of the output --help
 func (cl *CommandLineInterface) StringSliceFlag(name string, shorthand *string, defaultValue []string, description string) {
 	cl.StringSliceFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description)
 }
 
+// RegexFlag creates and registers a flag accepting a string and validates that it is a valid regex.
+func (cl *CommandLineInterface) RegexFlag(name string, shorthand *string, defaultValue *string, description string) {
+	cl.RegexFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description)
+}
+
 // StringOptionsFlag creates and registers a flag accepting a string and valid options for use in validation.
 func (cl *CommandLineInterface) StringOptionsFlag(name string, shorthand *string, defaultValue *string, description string, validOpts []string) {
 	cl.StringOptionsFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description, validOpts)
@@ -188,9 +193,9 @@ func (cl *CommandLineInterface) StringFlagOnFlagSet(flagSet *pflag.FlagSet, name
 	if shorthand != nil {
 		cl.Flags[name] = flagSet.StringP(name, string(*shorthand), *defaultValue, description)
 		cl.validators[name] = validationFn
-		return
+	} else {
+		cl.Flags[name] = flagSet.String(name, *defaultValue, description)
 	}
-	cl.Flags[name] = flagSet.String(name, *defaultValue, description)
 	cl.validators[name] = validationFn
 }
 
@@ -223,3 +228,36 @@ func (cl *CommandLineInterface) StringSliceFlagOnFlagSet(flagSet *pflag.FlagSet,
 	}
 	cl.Flags[name] = flagSet.StringSlice(name, defaultValue, description)
 }
+
+// RegexFlagOnFlagSet creates and registers a flag accepting a string slice of regular expressions.
+func (cl *CommandLineInterface) RegexFlagOnFlagSet(flagSet *pflag.FlagSet, name string, shorthand *string, defaultValue *string, description string) {
+	if defaultValue == nil {
+		cl.nilDefaults[name] = true
+		defaultValue = cl.StringMe("")
+	}
+	if shorthand != nil {
+		cl.Flags[name] = flagSet.StringP(name, string(*shorthand), *defaultValue, description)
+	} else {
+		cl.Flags[name] = flagSet.String(name, *defaultValue, description)
+	}
+	cl.validators[name] = func(val interface{}) error {
+		if val == nil {
+			return nil
+		}
+		regexStringVal := ""
+		switch v := val.(type) {
+		case *string:
+			regexStringVal = *v
+		case *regexp.Regexp:
+			return nil
+		default:
+			return fmt.Errorf("Invalid regex input for --%s", name)
+		}
+		regexVal, err := regexp.Compile(regexStringVal)
+		if err != nil {
+			return fmt.Errorf("Invalid regex input for --%s", name)
+		}
+		cl.Flags[name] = regexVal
+		return nil
+	}
+}

pkg/cli/flags_test.go

@@ -150,3 +150,19 @@ func TestIntMinMaxRangeFlags(t *testing.T) {
 	h.Assert(t, len(cli.Flags) == 3, "Should contain 3 flags w/ no shorthand")
 	h.Assert(t, ok, "Should contain %s flag w/ no shorthand", flagName)
 }
+
+func TestRegexFlag(t *testing.T) {
+	cli := getTestCLI()
+	for _, flagFn := range []func(string, *string, *string, string){cli.RegexFlag} {
+		flagName := "test-regex"
+		flagFn(flagName, cli.StringMe("t"), nil, "Test Regex")
+		_, ok := cli.Flags[flagName]
+		h.Assert(t, len(cli.Flags) == 1, "Should contain 1 flag")
+		h.Assert(t, ok, "Should contain %s flag", flagName)
+
+		cli = getTestCLI()
+		flagFn(flagName, nil, nil, "Test Regex")
+		h.Assert(t, len(cli.Flags) == 1, "Should contain 1 flag w/ no shorthand")
+		h.Assert(t, ok, "Should contain %s flag w/ no shorthand", flagName)
+	}
+}

pkg/cli/types.go

@@ -16,6 +16,7 @@ package cli
 
 import (
 	"log"
+	"regexp"
 
 	"github.com/aws/amazon-ec2-instance-selector/pkg/selector"
 	"github.com/spf13/cobra"
@@ -166,3 +167,20 @@ func (*CommandLineInterface) StringSliceMe(i interface{}) *[]string {
 		return nil
 	}
 }
+
+// RegexMe takes an interface and returns a pointer to a regex
+// If the underlying interface kind is not regexp.Regexp or *regexp.Regexp then nil is returned
+func (*CommandLineInterface) RegexMe(i interface{}) *regexp.Regexp {
+	if i == nil {
+		return nil
+	}
+	switch v := i.(type) {
+	case *regexp.Regexp:
+		return v
+	case regexp.Regexp:
+		return &v
+	default:
+		log.Printf("%s cannot be converted to a regexp", i)
+		return nil
+	}
+}

pkg/cli/types_test.go

@@ -15,6 +15,7 @@ package cli_test
 
 import (
 	"reflect"
+	"regexp"
 	"testing"
 
 	"github.com/aws/amazon-ec2-instance-selector/pkg/selector"
@@ -105,3 +106,17 @@ func TestIntRangeMe(t *testing.T) {
 	val = cli.IntRangeMe(nil)
 	h.Assert(t, val == nil, "Should return nil if nil is passed in")
 }
+
+func TestRegexMe(t *testing.T) {
+	cli := getTestCLI()
+	regexVal, err := regexp.Compile("c4.*")
+	h.Ok(t, err)
+	val := cli.RegexMe(*regexVal)
+	h.Assert(t, val.String() == regexVal.String(), "Should return %s from passed in regex value", regexVal)
+	val = cli.RegexMe(regexVal)
+	h.Assert(t, val.String() == regexVal.String(), "Should return %s from passed in regex pointer", regexVal)
+	val = cli.RegexMe(true)
+	h.Assert(t, val == nil, "Should return nil from other data type passed in")
+	val = cli.RegexMe(nil)
+	h.Assert(t, val == nil, "Should return nil if nil is passed in")
+}

pkg/selector/selector.go

@@ -17,6 +17,7 @@ package selector
 import (
 	"fmt"
 	"reflect"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -59,6 +60,8 @@ const (
 	currentGeneration      = "currentGeneration"
 	networkInterfaces      = "networkInterfaces"
 	networkPerformance     = "networkPerformance"
+	allowList              = "allowList"
+	denyList               = "denyList"
 
 	// AggregateLowPercentile is the default lower percentile for resource ranges on similar instance type comparisons
 	AggregateLowPercentile = 0.8
@@ -223,6 +226,10 @@ func (itf Selector) rawFilter(filters Filters) ([]*ec2.InstanceTypeInfo, error)
 				networkPerformance:     {filters.NetworkPerformance, getNetworkPerformance(instanceTypeInfo.NetworkInfo.NetworkPerformance)},
 			}
 
+			if isInDenyList(filters.DenyList, instanceTypeName) || !isInAllowList(filters.AllowList, instanceTypeName) {
+				delete(instanceTypeCandidates, instanceTypeName)
+			}
+
 			if !isSupportedInLocation(locationInstanceOfferings, instanceTypeName) {
 				delete(instanceTypeCandidates, instanceTypeName)
 			}
@@ -403,3 +410,17 @@ func isSupportedInLocation(instanceOfferings map[string]string, instanceType str
 	_, ok := instanceOfferings[instanceType]
 	return ok
 }
+
+func isInDenyList(denyRegex *regexp.Regexp, instanceTypeName string) bool {
+	if denyRegex == nil {
+		return false
+	}
+	return denyRegex.MatchString(instanceTypeName)
+}
+
+func isInAllowList(allowRegex *regexp.Regexp, instanceTypeName string) bool {
+	if allowRegex == nil {
+		return true
+	}
+	return allowRegex.MatchString(instanceTypeName)
+}

pkg/selector/selector_test.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"regexp"
 	"strconv"
 	"testing"
 
@@ -496,3 +497,60 @@ func TestRetrieveInstanceTypesSupportedInAZs_DescribeAZErr(t *testing.T) {
 	_, err := itf.RetrieveInstanceTypesSupportedInLocations([]string{"us-east-2a"})
 	h.Nok(t, err)
 }
+
+func TestFilter_AllowList(t *testing.T) {
+	ec2Mock := mockedEC2{
+		DescribeInstanceTypesPagesResp:    setupMock(t, describeInstanceTypesPages, "25_instances.json").DescribeInstanceTypesPagesResp,
+		DescribeInstanceTypeOfferingsResp: setupMock(t, describeInstanceTypeOfferings, "us-east-2a.json").DescribeInstanceTypeOfferingsResp,
+	}
+	itf := selector.Selector{
+		EC2: ec2Mock,
+	}
+	allowRegex, err := regexp.Compile("c4.large")
+	h.Ok(t, err)
+	filters := selector.Filters{
+		AllowList: allowRegex,
+	}
+	results, err := itf.Filter(filters)
+	h.Ok(t, err)
+	h.Assert(t, len(results) == 1, "Allow List Regex: 'c4.large' should return 1 instance type")
+}
+
+func TestFilter_DenyList(t *testing.T) {
+	ec2Mock := mockedEC2{
+		DescribeInstanceTypesPagesResp:    setupMock(t, describeInstanceTypesPages, "25_instances.json").DescribeInstanceTypesPagesResp,
+		DescribeInstanceTypeOfferingsResp: setupMock(t, describeInstanceTypeOfferings, "us-east-2a.json").DescribeInstanceTypeOfferingsResp,
+	}
+	itf := selector.Selector{
+		EC2: ec2Mock,
+	}
+	denyRegex, err := regexp.Compile("c4.large")
+	h.Ok(t, err)
+	filters := selector.Filters{
+		DenyList: denyRegex,
+	}
+	results, err := itf.Filter(filters)
+	h.Ok(t, err)
+	h.Assert(t, len(results) == 24, "Deny List Regex: 'c4.large' should return 24 instance type matching regex but returned %d", len(results))
+}
+
+func TestFilter_AllowAndDenyList(t *testing.T) {
+	ec2Mock := mockedEC2{
+		DescribeInstanceTypesPagesResp:    setupMock(t, describeInstanceTypesPages, "25_instances.json").DescribeInstanceTypesPagesResp,
+		DescribeInstanceTypeOfferingsResp: setupMock(t, describeInstanceTypeOfferings, "us-east-2a.json").DescribeInstanceTypeOfferingsResp,
+	}
+	itf := selector.Selector{
+		EC2: ec2Mock,
+	}
+	allowRegex, err := regexp.Compile("c4.*")
+	h.Ok(t, err)
+	denyRegex, err := regexp.Compile("c4.large")
+	h.Ok(t, err)
+	filters := selector.Filters{
+		AllowList: allowRegex,
+		DenyList:  denyRegex,
+	}
+	results, err := itf.Filter(filters)
+	h.Ok(t, err)
+	h.Assert(t, len(results) == 4, "Allow/Deny List Regex: 'c4.large' should return 4 instance types matching the regex but returned %d", len(results))
+}

pkg/selector/types.go

@@ -14,6 +14,8 @@
 package selector
 
 import (
+	"regexp"
+
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 )
@@ -133,6 +135,12 @@ type Filters struct {
 	// VcpusToMemoryRatio is a ratio of vcpus to memory expressed as a floating point
 	VCpusToMemoryRatio *float64
 
+	// AllowList is a regex of allowed instance types
+	AllowList *regexp.Regexp
+
+	// DenyList is a regex of excluded instance types
+	DenyList *regexp.Regexp
+
 	// InstanceTypeBase is a base instance type which is used to retrieve similarly spec'd instance types
 	InstanceTypeBase *string
 }