Merge pull request #206 from DataDog/gaspard.barbier/iac-add-arguments-to-config

Parse a new "arguments" keyword in new IaC v1.4 config

Co-authored-by: barbiergaspard <gaspard.barbier@datadoghq.com>

Author: gh-worker-dd-mergequeue-cf854d[bot]

pkg/config/config.go

@@ -41,9 +41,10 @@ type iacGlobalConfig struct {
 
 // iacRuleConfigYaml is the YAML representation of a per-rule override block.
 type iacRuleConfigYaml struct {
-	IgnorePaths []string     `yaml:"ignore-paths,omitempty"`
-	OnlyPaths   []string     `yaml:"only-paths,omitempty"`
-	Severity    *iacSeverity `yaml:"severity,omitempty"`
+	IgnorePaths []string       `yaml:"ignore-paths,omitempty"`
+	OnlyPaths   []string       `yaml:"only-paths,omitempty"`
+	Arguments   map[string]any `yaml:"arguments,omitempty"`
+	Severity    *iacSeverity   `yaml:"severity,omitempty"`
 }
 
 // ParseConfig turns a YAML configuration file into a parsed configuration
@@ -104,6 +105,7 @@ func parseRuleConfigs(in map[string]iacRuleConfigYaml) map[string]IacRuleConfig
 		out[ruleID] = IacRuleConfig{
 			IgnorePaths: rc.IgnorePaths,
 			OnlyPaths:   rc.OnlyPaths,
+			Arguments:   rc.Arguments,
 			Severity:    (*string)(rc.Severity),
 		}
 	}
@@ -119,6 +121,7 @@ func unparseRuleConfigs(in map[string]IacRuleConfig) map[string]iacRuleConfigYam
 		out[ruleID] = iacRuleConfigYaml{
 			IgnorePaths: rc.IgnorePaths,
 			OnlyPaths:   rc.OnlyPaths,
+			Arguments:   rc.Arguments,
 			Severity:    (*iacSeverity)(rc.Severity),
 		}
 	}
@@ -195,7 +198,7 @@ var (
 	// minIacVersion is the minimum schema version supporting the iac: configuration section.
 	minIacVersion = schemaVersion{1, 2}
 	// currentVersion is the current schema version emitted by UnparseConfig.
-	currentVersion = schemaVersion{1, 3}
+	currentVersion = schemaVersion{1, 4}
 	// minRequiredVersion is the minimum schema version supported by this scanner.
 	// It's different from minIacVersion for user-friendliness in case people add an `iac`
 	// section to an old configuration file and forget to upgrade the version number.

pkg/config/config_test.go

@@ -64,13 +64,13 @@ func TestUnparseConfig(t *testing.T) {
 	assert.Equal(t, cfgFile, string(b))
 }
 
-func TestUnparseConfigV13Fields(t *testing.T) {
+func TestUnparseConfigV14Fields(t *testing.T) {
 	t.Run("config with only-platforms serializes field", func(t *testing.T) {
 		cfg := parsedCfgFile
 		cfg.OnlyPlatforms = []string{"Terraform"}
 		b, err := UnparseConfig(&cfg)
 		require.NoError(t, err)
-		assert.Contains(t, string(b), "schema-version: v1.3")
+		assert.Contains(t, string(b), "schema-version: v1.4")
 		assert.Contains(t, string(b), "only-platforms:")
 	})
 
@@ -95,8 +95,8 @@ func TestUnparseConfigV13Fields(t *testing.T) {
 	})
 }
 
-func TestParseUnparseRoundTripV13(t *testing.T) {
-	input := `schema-version: v1.3
+func TestParseUnparseRoundTripV14(t *testing.T) {
+	input := `schema-version: v1.4
 iac:
   ignore-rules:
     - query1
@@ -111,6 +111,9 @@ iac:
       ignore-paths:
         - test/
       severity: low
+      arguments:
+        required_tags:
+          - Env
 `
 	cfg, err := ParseConfig([]byte(input))
 	require.NoError(t, err)
@@ -123,4 +126,6 @@ iac:
 	assert.Equal(t, []string{"test/"}, rc.IgnorePaths)
 	require.NotNil(t, rc.Severity)
 	assert.Equal(t, "low", *rc.Severity)
+	require.NotNil(t, rc.Arguments)
+	assert.Equal(t, []interface {}([]interface {}{"Env"}), rc.Arguments["required_tags"])
 }

pkg/config/config_validation_test.go

@@ -141,7 +141,7 @@ func TestParsePlatformValues(t *testing.T) {
 
 func TestParseRuleConfigs(t *testing.T) {
 	t.Run("full rule config", func(t *testing.T) {
-		input := `schema-version: v1.3
+		input := `schema-version: v1.4
 iac:
   rule-configs:
     terraform-aws-s3-unencrypted:
@@ -150,6 +150,8 @@ iac:
       only-paths:
         - src/
       severity: low
+      arguments:
+        required: enabled
 `
 		cfg, err := ParseConfig([]byte(input))
 		require.NoError(t, err)
@@ -160,10 +162,12 @@ iac:
 		assert.Equal(t, []string{"src/"}, rc.OnlyPaths)
 		require.NotNil(t, rc.Severity)
 		assert.Equal(t, "low", *rc.Severity)
+		require.NotNil(t, rc.Arguments)
+		assert.Equal(t, "enabled", (rc.Arguments)["required"])
 	})
 
 	t.Run("severity override only", func(t *testing.T) {
-		input := `schema-version: v1.3
+		input := `schema-version: v1.4
 iac:
   rule-configs:
     some-rule:
@@ -180,7 +184,7 @@ iac:
 	})
 
 	t.Run("invalid severity in rule-config", func(t *testing.T) {
-		input := `schema-version: v1.3
+		input := `schema-version: v1.4
 iac:
   rule-configs:
     some-rule:
@@ -192,7 +196,7 @@ iac:
 	})
 
 	t.Run("unknown field in rule-config", func(t *testing.T) {
-		input := `schema-version: v1.3
+		input := `schema-version: v1.4
 iac:
   rule-configs:
     some-rule:
@@ -202,6 +206,20 @@ iac:
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "iac.rule-configs.<rule-id>")
 	})
+
+	t.Run("arguments in rule-config", func(t *testing.T) {
+		input := `schema-version: v1.4
+iac:
+  rule-configs:
+    some-rule:
+      arguments:
+        required: enabled
+`
+		cfg, err := ParseConfig([]byte(input))
+		require.NoError(t, err)
+		require.NotNil(t, cfg.RuleConfigs["some-rule"].Arguments)
+		assert.Equal(t, "enabled", (cfg.RuleConfigs["some-rule"].Arguments)["required"])
+	})
 }
 
 func TestRewriteDecodeErrors(t *testing.T) {

pkg/config/reader_test.go

@@ -11,7 +11,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const cfgFile = `schema-version: v1.3
+const cfgFile = `schema-version: v1.4
 iac:
   ignore-rules:
     - query1
@@ -65,7 +65,7 @@ var parsedLegacyCfg = IacConfig{
 	IgnoreCategories: []string{"Access Control", "Availability"},
 }
 
-const convertedLegacyCfg = `schema-version: v1.3
+const convertedLegacyCfg = `schema-version: v1.4
 iac:
   ignore-rules:
     - query1
@@ -196,7 +196,7 @@ func TestReadLegacyOnly(t *testing.T) {
 func TestReadLegacyWithIncludeRulesOnly(t *testing.T) {
 	tmp := t.TempDir()
 	myCfgFile := legacyCfg + "include-queries: [query3, query4]\n"
-	includeOnly := "schema-version: v1.3\niac:\n  use-rules:\n    - query3\n    - query4\n"
+	includeOnly := "schema-version: v1.4\niac:\n  use-rules:\n    - query3\n    - query4\n"
 	require.NoError(t, os.WriteFile(filepath.Join(tmp, LegacyConfigFileName), []byte(myCfgFile), 0644))
 
 	cfg, b, err := ReadConfiguration(t.Context(), tmp)

pkg/config/types.go

@@ -20,5 +20,6 @@ type IacConfig struct {
 type IacRuleConfig struct {
 	IgnorePaths []string
 	OnlyPaths   []string
+	Arguments   map[string]any `yaml:"arguments,omitempty"`
 	Severity    *string
 }

pkg/engine/inspector.go

@@ -607,6 +607,45 @@ func (c *Inspector) GetFailedQueries() map[string]error {
 	return maps.Clone(c.failedQueries)
 }
 
+func ruleArgumentsValue(rc config.IacRuleConfig) (ast.Value, bool, error) {
+	if rc.Arguments == nil {
+		return nil, false, nil
+	}
+
+	value, err := ast.InterfaceToValue(rc.Arguments)
+	if err != nil {
+		return nil, false, err
+	}
+	return value, true, nil
+}
+
+func withRuleArguments(payload, args ast.Value) (ast.Value, error) {
+	obj, ok := payload.(ast.Object)
+	if !ok {
+		return nil, fmt.Errorf("expected OPA input payload object, got %T", payload)
+	}
+
+	out := ast.NewObject()
+	_ = obj.Iter(func(k, v *ast.Term) error {
+		out.Insert(k, v)
+		return nil
+	})
+	out.Insert(ast.StringTerm("arguments"), ast.NewTerm(args))
+	return out, nil
+}
+
+func queryIDsFromMetadata(ctx context.Context, metadata *model.QueryMetadata) (queryID, legacyQueryID string) {
+	queryID = DefaultQueryID
+	legacyQueryID = DefaultQueryID
+	if id, err := mapKeyToString(ctx, metadata.Metadata, "id", false); err == nil && id != nil {
+		queryID = *id
+	}
+	if legacyID, err := mapKeyToString(ctx, metadata.Metadata, "legacyId", true); err == nil && legacyID != nil {
+		legacyQueryID = *legacyID
+	}
+	return queryID, legacyQueryID
+}
+
 func (c *Inspector) doRun(ctx context.Context, qCtx *QueryContext) (vulns []model.Vulnerability, err error) {
 	contextLogger := logger.FromContext(ctx)
 	defer func() {
@@ -617,7 +656,19 @@ func (c *Inspector) doRun(ctx context.Context, qCtx *QueryContext) (vulns []mode
 		}
 	}()
 
-	options := []rego.EvalOption{rego.EvalParsedInput(*qCtx.payload)}
+	payload := *qCtx.payload
+	queryID, legacyQueryID := queryIDsFromMetadata(ctx, &qCtx.Query.Metadata)
+	if rc, found := lookupRuleConfig(c.ruleConfigs, queryID, legacyQueryID); found {
+		if args, ok, err := ruleArgumentsValue(rc); err != nil {
+			return nil, errors.Wrap(err, "Failed to prepare rule arguments for query "+queryID)
+		} else if ok {
+			payload, err = withRuleArguments(payload, args)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	options := []rego.EvalOption{rego.EvalParsedInput(payload)}
 
 	var cov *cover.Cover
 	if c.enableCoverageReport {