feat(cli): include queries filter (#3431)

Signed-off-by: Rogério Peixoto <[email protected]>
Co-authored-by: João Reigota <[email protected]>

e2e/fixtures/E2E_CLI_002

@@ -24,6 +24,10 @@ Flags:
       --ignore-on-exit string        defines which kind of non-zero exits code should be ignored
                                      accepts: all, results, errors, none
                                      example: if 'results' is set, only engine errors will make KICS exit code different from 0 (default "none")
+  -i, --include-queries strings      include queries by providing the query ID
+                                     takes precedence over excluded queries
+                                     can be provided multiple times or as a comma separated string
+                                     example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'
       --minimal-ui                   simplified version of CLI output
       --no-progress                  hides the progress bar
   -o, --output-path string           directory path to store reports

e2e/fixtures/E2E_CLI_003

@@ -23,6 +23,10 @@ Flags:
       --ignore-on-exit string        defines which kind of non-zero exits code should be ignored
                                      accepts: all, results, errors, none
                                      example: if 'results' is set, only engine errors will make KICS exit code different from 0 (default "none")
+  -i, --include-queries strings      include queries by providing the query ID
+                                     takes precedence over excluded queries
+                                     can be provided multiple times or as a comma separated string
+                                     example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'
       --minimal-ui                   simplified version of CLI output
       --no-progress                  hides the progress bar
   -o, --output-path string           directory path to store reports

e2e/fixtures/E2E_CLI_004

@@ -23,6 +23,10 @@ Flags:
       --ignore-on-exit string        defines which kind of non-zero exits code should be ignored
                                      accepts: all, results, errors, none
                                      example: if 'results' is set, only engine errors will make KICS exit code different from 0 (default "none")
+  -i, --include-queries strings      include queries by providing the query ID
+                                     takes precedence over excluded queries
+                                     can be provided multiple times or as a comma separated string
+                                     example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'
       --minimal-ui                   simplified version of CLI output
       --no-progress                  hides the progress bar
   -o, --output-path string           directory path to store reports

e2e/fixtures/E2E_CLI_016_INVALID_SCAN_FLAG

@@ -23,6 +23,10 @@ Flags:
       --ignore-on-exit string        defines which kind of non-zero exits code should be ignored
                                      accepts: all, results, errors, none
                                      example: if 'results' is set, only engine errors will make KICS exit code different from 0 (default "none")
+  -i, --include-queries strings      include queries by providing the query ID
+                                     takes precedence over excluded queries
+                                     can be provided multiple times or as a comma separated string
+                                     example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'
       --minimal-ui                   simplified version of CLI output
       --no-progress                  hides the progress bar
   -o, --output-path string           directory path to store reports

internal/console/scan.go

@@ -47,6 +47,7 @@ var (
 	excludeIDs        []string
 	excludePath       []string
 	excludeResults    []string
+	includeIDs        []string
 	failOn            []string
 	ignoreOnExit      string
 	min               bool
@@ -69,6 +70,8 @@ const (
 	excludeQueriesFlag      = "exclude-queries"
 	excludeResultsFlag      = "exclude-results"
 	excludeResutlsShorthand = "x"
+	includeQueriesFlag      = "include-queries"
+	inlcudeQueriesShorthand = "i"
 	failOnFlag              = "fail-on"
 	ignoreOnExitFlag        = "ignore-on-exit"
 	minimalUIFlag           = "minimal-ui"
@@ -134,6 +137,10 @@ func preRun(cmd *cobra.Command) error {
 	if err != nil {
 		return errors.New(initError + err.Error())
 	}
+	err = validateQuerySelectionFlags()
+	if err != nil {
+		return err
+	}
 	err = internalPrinter.SetupPrinter(cmd.InheritedFlags())
 	if err != nil {
 		return errors.New(initError + err.Error())
@@ -145,6 +152,22 @@ func preRun(cmd *cobra.Command) error {
 	return nil
 }
 
+func formatNewError(flag1, flag2 string) error {
+	return errors.Errorf("can't provide '%s' and '%s' flags simultaneously",
+		flag1,
+		flag2)
+}
+
+func validateQuerySelectionFlags() error {
+	if len(includeIDs) > 0 && len(excludeIDs) > 0 {
+		return formatNewError(includeQueriesFlag, excludeQueriesFlag)
+	}
+	if len(includeIDs) > 0 && len(excludeCategories) > 0 {
+		return formatNewError(includeQueriesFlag, excludeCategoriesFlag)
+	}
+	return nil
+}
+
 func initializeConfig(cmd *cobra.Command) error {
 	log.Debug().Msg("console.initializeConfig()")
 
@@ -314,6 +337,15 @@ func initScanFlags(scanCmd *cobra.Command) {
 			msg+
 			"example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'",
 	)
+	scanCmd.Flags().StringSliceVarP(&includeIDs,
+		includeQueriesFlag,
+		inlcudeQueriesShorthand,
+		[]string{},
+		"include queries by providing the query ID\n"+
+			"takes precedence over excluded queries\n"+
+			msg+
+			"example: 'e69890e6-fce5-461d-98ad-cb98318dfc96,4728cd65-a20c-49da-8b31-9c08b423e4db'",
+	)
 	scanCmd.Flags().StringSliceVarP(&excludeResults,
 		excludeResultsFlag,
 		excludeResutlsShorthand,
@@ -398,9 +430,22 @@ func createInspector(t engine.Tracker, querySource source.QueriesSource) (*engin
 		ByCategories: excludeCategories,
 	}
 
+	includeQueries := source.IncludeQueries{
+		ByIDs: includeIDs,
+	}
+
+	queryFilter := source.QuerySelectionFilter{
+		IncludeQueries: includeQueries,
+		ExcludeQueries: excludeQueries,
+	}
+
 	inspector, err := engine.NewInspector(ctx,
-		querySource, engine.DefaultVulnerabilityBuilder,
-		t, excludeQueries, excludeResultsMap, queryExecTimeout)
+		querySource,
+		engine.DefaultVulnerabilityBuilder,
+		t,
+		queryFilter,
+		excludeResultsMap,
+		queryExecTimeout)
 	if err != nil {
 		return nil, err
 	}

pkg/engine/inspector.go

@@ -101,13 +101,13 @@ func NewInspector(
 	queriesSource source.QueriesSource,
 	vb VulnerabilityBuilder,
 	tracker Tracker,
-	excludeQueries source.ExcludeQueries,
+	queryFilter source.QuerySelectionFilter,
 	excludeResults map[string]bool,
 	queryTimeout int) (*Inspector, error) {
 	log.Debug().Msg("engine.NewInspector()")
 
 	metrics.Metric.Start("get_queries")
-	queries, err := queriesSource.GetQueries(excludeQueries)
+	queries, err := queriesSource.GetQueries(queryFilter)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get queries")
 	}

pkg/engine/inspector_test.go

@@ -357,7 +357,7 @@ func TestNewInspector(t *testing.T) { // nolint
 		source           source.QueriesSource
 		vb               VulnerabilityBuilder
 		tracker          Tracker
-		excludeQueries   source.ExcludeQueries
+		queryFilter      source.QuerySelectionFilter
 		excludeResults   map[string]bool
 		queryExecTimeout int
 	}
@@ -374,9 +374,14 @@ func TestNewInspector(t *testing.T) { // nolint
 				vb:      vbs,
 				tracker: track,
 				source:  sources,
-				excludeQueries: source.ExcludeQueries{
-					ByIDs:        []string{},
-					ByCategories: []string{},
+				queryFilter: source.QuerySelectionFilter{
+					IncludeQueries: source.IncludeQueries{
+						ByIDs: []string{},
+					},
+					ExcludeQueries: source.ExcludeQueries{
+						ByIDs:        []string{},
+						ByCategories: []string{},
+					},
 				},
 				excludeResults:   map[string]bool{},
 				queryExecTimeout: 60,
@@ -395,7 +400,7 @@ func TestNewInspector(t *testing.T) { // nolint
 				tt.args.source,
 				tt.args.vb,
 				tt.args.tracker,
-				tt.args.excludeQueries,
+				tt.args.queryFilter,
 				tt.args.excludeResults,
 				tt.args.queryExecTimeout)
 
@@ -556,7 +561,7 @@ func newInspectorInstance(t *testing.T, queryPath string) *Inspector {
 		detector *detector.DetectLine) (model.Vulnerability, error) {
 		return model.Vulnerability{}, nil
 	}
-	ins, err := NewInspector(context.Background(), querySource, vb, &tracker.CITracker{}, source.ExcludeQueries{}, map[string]bool{}, 60)
+	ins, err := NewInspector(context.Background(), querySource, vb, &tracker.CITracker{}, source.QuerySelectionFilter{}, map[string]bool{}, 60)
 	require.NoError(t, err)
 	return ins
 }
@@ -566,10 +571,10 @@ type mockSource struct {
 	Types  []string
 }
 
-func (m *mockSource) GetQueries(excludeQueries source.ExcludeQueries) ([]model.QueryMetadata, error) {
+func (m *mockSource) GetQueries(queryFilter source.QuerySelectionFilter) ([]model.QueryMetadata, error) {
 	sources := source.NewFilesystemSource(m.Source, []string{""})
 
-	return sources.GetQueries(excludeQueries)
+	return sources.GetQueries(queryFilter)
 }
 
 func (m *mockSource) GetQueryLibrary(platform string) (string, error) {

pkg/engine/source/filesystem.go

@@ -114,6 +114,21 @@ func (s *FilesystemSource) CheckType(queryPlatform interface{}) bool {
 	return true
 }
 
+func checkQueryInclude(id interface{}, includedQueries []string) bool {
+	queryMetadataKey, ok := id.(string)
+	if !ok {
+		log.Warn().
+			Msgf("Can't cast query metadata key = %v", id)
+		return false
+	}
+	for _, includedQuery := range includedQueries {
+		if queryMetadataKey == includedQuery {
+			return true
+		}
+	}
+	return false
+}
+
 func checkQueryExclude(id interface{}, excludeQueries []string) bool {
 	queryMetadataKey, ok := id.(string)
 	if !ok {
@@ -131,7 +146,7 @@ func checkQueryExclude(id interface{}, excludeQueries []string) bool {
 
 // GetQueries walks a given filesource path returns all queries found in an array of
 // QueryMetadata struct
-func (s *FilesystemSource) GetQueries(excludeQueries ExcludeQueries) ([]model.QueryMetadata, error) {
+func (s *FilesystemSource) GetQueries(queryFilter QuerySelectionFilter) ([]model.QueryMetadata, error) {
 	queryDirs := make([]string, 0)
 	err := filepath.Walk(s.Source,
 		func(p string, f os.FileInfo, err error) error {
@@ -164,14 +179,21 @@ func (s *FilesystemSource) GetQueries(excludeQueries ExcludeQueries) ([]model.Qu
 		if !s.CheckType(query.Metadata["platform"]) {
 			continue
 		}
-		if checkQueryExclude(query.Metadata["id"], excludeQueries.ByIDs) ||
-			checkQueryExclude(query.Metadata["category"], excludeQueries.ByCategories) {
-			log.Debug().
-				Msgf("Excluding query ID: %s category: %s", query.Metadata["id"], query.Metadata["category"])
-			continue
-		}
 
-		queries = append(queries, query)
+		if len(queryFilter.IncludeQueries.ByIDs) > 0 {
+			if checkQueryInclude(query.Metadata["id"], queryFilter.IncludeQueries.ByIDs) {
+				queries = append(queries, query)
+			}
+		} else {
+			if checkQueryExclude(query.Metadata["id"], queryFilter.ExcludeQueries.ByIDs) ||
+				checkQueryExclude(query.Metadata["category"], queryFilter.ExcludeQueries.ByCategories) {
+				log.Debug().
+					Msgf("Excluding query ID: %s category: %s", query.Metadata["id"], query.Metadata["category"])
+				continue
+			}
+
+			queries = append(queries, query)
+		}
 	}
 
 	return queries, err