From bc3741ae2c68cdd00fc0aef7e51985568b2eb78a Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Tue, 18 Jun 2024 05:20:38 +0700
Subject: [PATCH] feat(misconf): support of selectors for all providers for
 Rego (#6905)

Signed-off-by: nikpivkin <nikita.pivkin@smartforce.io>
---
 pkg/iac/providers/provider.go |  7 +++++++
 pkg/iac/rego/scanner.go       | 38 ++++++++++++++++++++++++-----------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/pkg/iac/providers/provider.go b/pkg/iac/providers/provider.go
index cef13ee8f205..46dbf19ec43c 100755
--- a/pkg/iac/providers/provider.go
+++ b/pkg/iac/providers/provider.go
@@ -26,6 +26,13 @@ const (
 	CloudStackProvider   Provider = "cloudstack"
 )
 
+func AllProviders() []Provider {
+	return []Provider{
+		AWSProvider, AzureProvider, DigitalOceanProvider, GitHubProvider, GoogleProvider,
+		KubernetesProvider, OracleProvider, OpenStackProvider, NifcloudProvider, CloudStackProvider,
+	}
+}
+
 func RuleProviderToString(provider Provider) string {
 	return strings.ToUpper(string(provider))
 }
diff --git a/pkg/iac/rego/scanner.go b/pkg/iac/rego/scanner.go
index 723f4c02181b..2e0516761a02 100644
--- a/pkg/iac/rego/scanner.go
+++ b/pkg/iac/rego/scanner.go
@@ -17,12 +17,30 @@ import (
 
 	"github.com/aquasecurity/trivy/pkg/iac/debug"
 	"github.com/aquasecurity/trivy/pkg/iac/framework"
+	"github.com/aquasecurity/trivy/pkg/iac/providers"
 	"github.com/aquasecurity/trivy/pkg/iac/rego/schemas"
 	"github.com/aquasecurity/trivy/pkg/iac/scan"
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/options"
 	"github.com/aquasecurity/trivy/pkg/iac/types"
 )
 
+var checkTypesWithSubtype = map[types.Source]struct{}{
+	types.SourceCloud:      {},
+	types.SourceDefsec:     {},
+	types.SourceKubernetes: {},
+}
+
+var supportedProviders = makeSupportedProviders()
+
+func makeSupportedProviders() map[string]struct{} {
+	m := make(map[string]struct{})
+	for _, p := range providers.AllProviders() {
+		m[string(p)] = struct{}{}
+	}
+	m["kind"] = struct{}{} // kubernetes
+	return m
+}
+
 var _ options.ConfigurableScanner = (*Scanner)(nil)
 
 type Scanner struct {
@@ -295,12 +313,8 @@ func (s *Scanner) ScanInput(ctx context.Context, inputs ...Input) (scan.Results,
 }
 
 func isPolicyWithSubtype(sourceType types.Source) bool {
-	for _, s := range []types.Source{types.SourceCloud, types.SourceDefsec, types.SourceKubernetes} {
-		if sourceType == s {
-			return true
-		}
-	}
-	return false
+	_, exists := checkTypesWithSubtype[sourceType]
+	return exists
 }
 
 func checkSubtype(ii map[string]any, provider string, subTypes []SubType) bool {
@@ -311,10 +325,11 @@ func checkSubtype(ii map[string]any, provider string, subTypes []SubType) bool {
 	for _, st := range subTypes {
 		switch services := ii[provider].(type) {
 		case map[string]any:
-			for service := range services {
-				if (service == st.Service) && (st.Provider == provider) {
-					return true
-				}
+			if st.Provider != provider {
+				continue
+			}
+			if _, exists := services[st.Service]; exists {
+				return true
 			}
 		case string: // k8s - logic can be improved
 			if strings.EqualFold(services, st.Group) ||
@@ -331,8 +346,7 @@ func isPolicyApplicable(staticMetadata *StaticMetadata, inputs ...Input) bool {
 	for _, input := range inputs {
 		if ii, ok := input.Contents.(map[string]any); ok {
 			for provider := range ii {
-				// TODO(simar): Add other providers
-				if !strings.Contains(strings.Join([]string{"kind", "aws", "azure"}, ","), provider) {
+				if _, exists := supportedProviders[provider]; !exists {
 					continue
 				}