✨ add filters to aws discovery (#2182)

providers/aws/config/config.go

@@ -84,6 +84,12 @@ var Config = plugin.Provider{
 					Default: "",
 					Desc:    "Override option for EBS scanning that tells it to not create the snapshot or volume",
 				},
+				{
+					Long:    "filters",
+					Type:    plugin.FlagType_KeyValue,
+					Default: "",
+					Desc:    "Filter options",
+				},
 			},
 		},
 	},

providers/aws/connection/connection.go

@@ -6,6 +6,7 @@ package connection
 import (
 	"context"
 	"errors"
+	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -29,6 +30,34 @@ type AwsConnection struct {
 	profile           string
 	PlatformOverride  string
 	connectionOptions map[string]string
+	Filters           DiscoveryFilters
+	RegionLimits      []string
+}
+
+type DiscoveryFilters struct {
+	Ec2DiscoveryFilters     Ec2DiscoveryFilters
+	EcrDiscoveryFilters     EcrDiscoveryFilters
+	EcsDiscoveryFilters     EcsDiscoveryFilters
+	GeneralDiscoveryFilters GeneralResourceDiscoveryFilters
+}
+
+type GeneralResourceDiscoveryFilters struct {
+	Tags    map[string]string
+	Regions []string
+}
+
+type Ec2DiscoveryFilters struct {
+	Regions     []string
+	Tags        map[string]string
+	InstanceIds []string
+}
+type EcrDiscoveryFilters struct {
+	Tags []string
+}
+type EcsDiscoveryFilters struct {
+	OnlyRunningContainers bool
+	DiscoverImages        bool
+	DiscoverInstances     bool
 }
 
 func NewMockConnection(id uint32, asset *inventory.Asset, conf *inventory.Config) *AwsConnection {
@@ -77,9 +106,40 @@ func NewAwsConnection(id uint32, asset *inventory.Asset, conf *inventory.Config)
 	c.accountId = *identity.Account
 	c.profile = asset.Options["profile"]
 	c.connectionOptions = asset.Options
+	c.Filters = parseOptsToFilters(conf.Discover.Filter)
+	c.RegionLimits = c.Filters.GeneralDiscoveryFilters.Regions
 	return c, nil
 }
 
+func parseOptsToFilters(opts map[string]string) DiscoveryFilters {
+	d := DiscoveryFilters{
+		Ec2DiscoveryFilters:     Ec2DiscoveryFilters{Tags: map[string]string{}},
+		EcsDiscoveryFilters:     EcsDiscoveryFilters{},
+		EcrDiscoveryFilters:     EcrDiscoveryFilters{Tags: []string{}},
+		GeneralDiscoveryFilters: GeneralResourceDiscoveryFilters{Tags: map[string]string{}},
+	}
+	for k, v := range opts {
+		switch {
+		case strings.HasPrefix(k, "ec2:tag:"):
+			d.Ec2DiscoveryFilters.Tags[strings.TrimPrefix("ec2:tag:", k)] = v
+		case k == "ec2:region":
+			d.Ec2DiscoveryFilters.Regions = append(d.Ec2DiscoveryFilters.Regions, v)
+		case k == "all:region":
+			d.GeneralDiscoveryFilters.Regions = append(d.GeneralDiscoveryFilters.Regions, v)
+		case k == "region":
+			d.GeneralDiscoveryFilters.Regions = append(d.GeneralDiscoveryFilters.Regions, v)
+			d.Ec2DiscoveryFilters.Regions = append(d.Ec2DiscoveryFilters.Regions, v)
+		case k == "instance-id":
+			d.Ec2DiscoveryFilters.InstanceIds = append(d.Ec2DiscoveryFilters.InstanceIds, v)
+		case strings.HasPrefix(k, "all:tag:"):
+			d.GeneralDiscoveryFilters.Tags[strings.TrimPrefix("all:tag:", k)] = v
+		case k == "ecr:tag":
+			d.EcrDiscoveryFilters.Tags = append(d.EcrDiscoveryFilters.Tags, v)
+		}
+	}
+	return d
+}
+
 func parseFlagsForConnectionOptions(m map[string]string) []ConnectionOption {
 	o := make([]ConnectionOption, 0)
 	if apiEndpoint, ok := m["endpoint-url"]; ok {
@@ -229,6 +289,12 @@ func (h *AwsConnection) Regions() ([]string, error) {
 	}
 	log.Debug().Msg("no region cache found. fetching regions")
 
+	if len(h.RegionLimits) > 0 {
+		log.Debug().Interface("regions", h.RegionLimits).Msg("using region limits")
+		// cache the regions as part of the provider instance
+		h.clientcache.Store("_regions", &CacheEntry{Data: h.RegionLimits})
+		return h.RegionLimits, nil
+	}
 	// if no cache, get regions using ec2 client (using the ssm list global regions does not give the same list)
 	regions := []string{}
 	svc := h.Ec2("us-east-1")

providers/aws/provider/provider.go

@@ -6,6 +6,7 @@ package provider
 import (
 	"errors"
 	"strconv"
+	"strings"
 
 	"go.mondoo.com/cnquery/v9/llx"
 	"go.mondoo.com/cnquery/v9/providers-sdk/v1/inventory"
@@ -57,8 +58,9 @@ func (s *Service) ParseCLI(req *plugin.ParseCLIReq) (*plugin.ParseCLIRes, error)
 			discoverTargets = append(discoverTargets, entry)
 		}
 	}
+	filterOpts := parseFlagsToFiltersOpts(flags)
 
-	inventoryConfig.Discover = &inventory.Discovery{Targets: discoverTargets}
+	inventoryConfig.Discover = &inventory.Discovery{Targets: discoverTargets, Filter: filterOpts}
 	asset := inventory.Asset{
 		Connections: []*inventory.Config{inventoryConfig},
 		Options:     opts,
@@ -79,6 +81,26 @@ func handleAwsEc2Subcommands(args []string, opts map[string]string) *inventory.A
 	return asset
 }
 
+func parseFlagsToFiltersOpts(m map[string]*llx.Primitive) map[string]string {
+	o := make(map[string]string, 0)
+
+	if x, ok := m["filters"]; ok && len(x.Map) != 0 {
+		for k, v := range x.Map {
+			if strings.Contains(k, "tag:") {
+				o[k] = string(v.Value)
+			}
+			if k == "instance-id" {
+				o[k] = string(v.Value)
+			}
+			if strings.Contains(k, ":region") {
+				o[k] = string(v.Value)
+			}
+		}
+	}
+
+	return o
+}
+
 func parseFlagsToOptions(m map[string]*llx.Primitive) map[string]string {
 	o := make(map[string]string, 0)
 	for k, v := range m {
@@ -300,5 +322,5 @@ func (s *Service) discover(conn *connection.AwsConnection) (*inventory.Inventory
 		return nil, errors.New("connection " + strconv.FormatUint(uint64(conn.ID()), 10) + " not found")
 	}
 
-	return resources.Discover(runtime)
+	return resources.Discover(runtime, conn.Filters)
 }

providers/aws/resources/discovery.go

@@ -89,7 +89,93 @@ var AllAPIResources = []string{
 	DiscoverySagemakerNotebookInstances,
 }
 
-func Discover(runtime *plugin.Runtime) (*inventory.Inventory, error) {
+func contains(sl []string, s string) bool {
+	for i := range sl {
+		if sl[i] == s {
+			return true
+		}
+	}
+	return false
+}
+
+func containsInterfaceSlice(sl []interface{}, s string) bool {
+	for i := range sl {
+		if sl[i].(string) == s {
+			return true
+		}
+	}
+	return false
+}
+
+func instanceMatchesFilters(instance *mqlAwsEc2Instance, filters connection.DiscoveryFilters) bool {
+	matches := true
+	f := filters.Ec2DiscoveryFilters
+	if len(f.Regions) > 0 {
+		if !contains(f.Regions, instance.Region.Data) {
+			matches = false
+		}
+	}
+	if len(f.InstanceIds) > 0 {
+		if !contains(f.InstanceIds, instance.InstanceId.Data) {
+			matches = false
+		}
+	}
+	if len(f.Tags) > 0 {
+		for k, v := range f.Tags {
+			if instance.Tags.Data[k] == nil {
+				return false
+			}
+			if instance.Tags.Data[k].(string) != v {
+				return false
+			}
+		}
+	}
+	return matches
+}
+
+func imageMatchesFilters(image *mqlAwsEcrImage, filters connection.DiscoveryFilters) bool {
+	f := filters.EcrDiscoveryFilters
+	if len(f.Tags) > 0 {
+		for i := range f.Tags {
+			t := f.Tags[i]
+			if !containsInterfaceSlice(image.Tags.Data, t) {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func containerMatchesFilters(container *mqlAwsEcsContainer, filters connection.DiscoveryFilters) bool {
+	f := filters.EcsDiscoveryFilters
+	if f.OnlyRunningContainers {
+		if container.Status.Data != "RUNNING" {
+			return false
+		}
+	}
+	return true
+}
+
+func shouldScanEcsContainerInstances(filters connection.DiscoveryFilters) bool {
+	return filters.EcsDiscoveryFilters.DiscoverInstances
+}
+
+func shouldScanEcsContainerImages(filters connection.DiscoveryFilters) bool {
+	return filters.EcsDiscoveryFilters.DiscoverImages
+}
+
+func discoveredAssetMatchesGeneralFilters(asset *inventory.Asset, filters connection.GeneralResourceDiscoveryFilters) bool {
+	if len(filters.Tags) > 0 {
+		for k, v := range filters.Tags {
+			if asset.Labels[k] != v {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func Discover(runtime *plugin.Runtime, filters connection.DiscoveryFilters) (*inventory.Inventory, error) {
 	conn := runtime.Connection.(*connection.AwsConnection)
 
 	in := &inventory.Inventory{Spec: &inventory.InventorySpec{
@@ -106,11 +192,20 @@ func Discover(runtime *plugin.Runtime) (*inventory.Inventory, error) {
 	targets := handleTargets(conn.Conf.Discover.Targets)
 	for i := range targets {
 		target := targets[i]
-		list, err := discover(runtime, awsAccount, target)
+		list, err := discover(runtime, awsAccount, target, filters)
 		if err != nil {
 			log.Error().Err(err).Msg("error during discovery")
 			continue
 		}
+		if len(filters.GeneralDiscoveryFilters.Tags) > 0 {
+			newList := []*inventory.Asset{}
+			for i := range list {
+				if discoveredAssetMatchesGeneralFilters(list[i], filters.GeneralDiscoveryFilters) {
+					newList = append(newList, list[i])
+				}
+			}
+			list = newList
+		}
 		in.Spec.Assets = append(in.Spec.Assets, list...)
 	}
 
@@ -133,7 +228,10 @@ func handleTargets(targets []string) []string {
 	return targets
 }
 
-func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string) ([]*inventory.Asset, error) {
+// for now we have to post proces the filters
+// more ideally, we should pass the filters in when discovering
+// so that we dont unnecesarily discover assets we will later discard
+func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string, filters connection.DiscoveryFilters) ([]*inventory.Asset, error) {
 	conn := runtime.Connection.(*connection.AwsConnection)
 	accountId := trimAwsAccountIdToJustId(awsAccount.Id.Data)
 	assetList := []*inventory.Asset{}
@@ -156,7 +254,11 @@ func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string)
 
 		for i := range ins.Data {
 			instance := ins.Data[i].(*mqlAwsEc2Instance)
+			if !instanceMatchesFilters(instance, filters) {
+				continue
+			}
 			assetList = append(assetList, addConnectionInfoToEc2Asset(instance, accountId, conn))
+
 		}
 	case DiscoverySSMInstances:
 		res, err := NewResource(runtime, "aws.ec2", map[string]*llx.RawData{})
@@ -173,6 +275,9 @@ func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string)
 
 		for i := range ins.Data {
 			instance := ins.Data[i].(*mqlAwsEc2Instance)
+			if !instanceMatchesFilters(instance, filters) {
+				continue
+			}
 			if instance.GetSsm() != nil {
 				if s := instance.GetSsm().Data.(map[string]interface{})["PingStatus"]; s != nil && s == "Online" {
 					assetList = append(assetList, addSSMConnectionInfoToEc2Asset(instance, accountId, conn))
@@ -210,6 +315,9 @@ func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string)
 
 		for i := range images.Data {
 			a := images.Data[i].(*mqlAwsEcrImage)
+			if !imageMatchesFilters(a, filters) {
+				continue
+			}
 			assetList = append(assetList, addConnectionInfoToEcrAsset(a, conn))
 		}
 	case DiscoveryECS:
@@ -227,18 +335,23 @@ func discover(runtime *plugin.Runtime, awsAccount *mqlAwsAccount, target string)
 
 		for i := range containers.Data {
 			c := containers.Data[i].(*mqlAwsEcsContainer)
+			if !containerMatchesFilters(c, filters) {
+				continue
+			}
 			assetList = append(assetList, addConnectionInfoToECSContainerAsset(c, accountId, conn))
 		}
-		containerInst := ecs.GetContainerInstances()
-		if containerInst == nil {
-			return assetList, nil
-		}
+		if shouldScanEcsContainerInstances(filters) {
+			containerInst := ecs.GetContainerInstances()
+			if containerInst == nil {
+				return assetList, nil
+			}
 
-		for i := range containerInst.Data {
-			if a, ok := containerInst.Data[i].(*mqlAwsEc2Instance); ok {
-				assetList = append(assetList, addConnectionInfoToEc2Asset(a, accountId, conn))
-			} else if b, ok := containerInst.Data[i].(*mqlAwsEcsInstance); ok {
-				assetList = append(assetList, addConnectionInfoToECSContainerInstanceAsset(b, accountId, conn))
+			for i := range containerInst.Data {
+				if a, ok := containerInst.Data[i].(*mqlAwsEc2Instance); ok {
+					assetList = append(assetList, addConnectionInfoToEc2Asset(a, accountId, conn))
+				} else if b, ok := containerInst.Data[i].(*mqlAwsEcsInstance); ok {
+					assetList = append(assetList, addConnectionInfoToECSContainerInstanceAsset(b, accountId, conn))
+				}
 			}
 		}
 	// case DiscoveryECSContainersAPI: