feat(misconf): use Terraform plan configuration to partially restore schema

pkg/iac/scanners/terraformplan/tfjson/parser/parser.go

@@ -9,6 +9,8 @@
 	"os"
 	"strings"
 
+	"github.com/samber/lo"
+
 	"github.com/aquasecurity/trivy/pkg/iac/terraform"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/mapfs"
@@ -51,21 +53,18 @@
 }
 
 func (p *PlanFile) ToFS() (fs.FS, error) {
-
-	rootFS := mapfs.New()
-
-	var fileResources []string
-
 	resources, err := getResources(p.PlannedValues.RootModule, p.ResourceChanges, p.Configuration)
 	if err != nil {
 		return nil, err
 	}
 
+	fileResources := make([]string, 0, len(resources))
 	for _, r := range resources {
 		fileResources = append(fileResources, r.ToHCL())
 	}
-
 	fileContent := strings.Join(fileResources, "\n\n")
+
+	rootFS := mapfs.New()
 	if err := rootFS.WriteVirtualFile("main.tf", []byte(fileContent), os.ModePerm); err != nil {
 		return nil, err
 	}
@@ -84,40 +83,19 @@
 			resourceName = fmt.Sprintf("%s_%s", r.Name, hash)
 		}
 
-		res := terraform.NewPlanBlock(r.Mode, r.Type, resourceName)
-
-		changes := getValues(r.Address, resourceChanges)
-		// process the changes to get the after state
-		for k, v := range changes.After {
-			switch t := v.(type) {
-			case []any:
-				if len(t) == 0 {
-					continue
-				}
-				val := t[0]
-				switch v := val.(type) {
-				// is it a HCL block?
-				case map[string]any:
-					res.Blocks[k] = v
-				// just a normal attribute then
-				default:
-					res.Attributes[k] = v
-				}
-			default:
-				res.Attributes[k] = v
-			}
-		}
-
 		resourceConfig := getConfiguration(r.Address, configuration.RootModule)
+		schema := make(BlockSchema)
 		if resourceConfig != nil {
-
-			for attr, val := range resourceConfig.Expressions {
-				if value, shouldReplace := unpackConfigurationValue(val, r); shouldReplace || !res.HasAttribute(attr) {
-					res.Attributes[attr] = value
-				}
-			}
+			schema = schemaForBlock(r, resourceConfig.Expressions)
 		}
-		resources = append(resources, *res)
+
+		changes := getValues(r.Address, resourceChanges)
+		resource := decodeBlock(schema, changes.After)
+		// fill top-level block fileds
+		resource.BlockType = lo.Ternary(r.Mode == "managed", "resource", r.Mode)
+		resource.Type = r.Type
+		resource.Name = resourceName
+		resources = append(resources, resource)
 	}
 
 	for _, m := range module.ChildModules {
@@ -131,6 +109,98 @@
 	return resources, nil
 }
 
+func decodeBlock(schema BlockSchema, rawBlock map[string]any) terraform.PlanBlock {
+	block := terraform.PlanBlock{
+		Attributes: make(map[string]any),
+	}
+
+	for k, child := range rawBlock {
+		childSchema := schema[k]
+		switch t := child.(type) {
+		case []any:
+			if childSchema != nil {
+				switch childSchema.Type {
+				case Attribute:
+					block.Attributes[k] = decodeAttribute(childSchema, child)
+				case Block:
+					nestedBlocks := decodeNestedBlocks(childSchema, k, t)
+					block.Blocks = append(block.Blocks, nestedBlocks...)
+				}
+			} else {
+				// just attribute
+				block.Attributes[k] = t
+			}
+		default:
+			if childSchema != nil {
+				switch childSchema.Type {
+				case Attribute:
+					block.Attributes[k] = decodeAttribute(childSchema, child)
+				case Block:
+					nestedBlocks := decodeNestedBlocks(childSchema, k, []any{child})
+					block.Blocks = append(block.Blocks, nestedBlocks...)
+				}
+			} else {
+				block.Attributes[k] = child
+			}
+		}
+	}
+	return block
+}
+
+func decodeNestedBlocks(schema *SchemaNode, name string, v []any) []terraform.PlanBlock {
+	nestedBlocks := make([]terraform.PlanBlock, 0, len(v))
+	for i, el := range v {
+		m, ok := el.(map[string]any)
+		if !ok {
+			continue
+		}
+		nestedBlockSchema := make(BlockSchema)
+		if i < len(schema.Children) {
+			nestedBlockSchema = schema.Children[i]
+		}
+		nestedBlock := decodeBlock(nestedBlockSchema, m)
+		nestedBlock.Name = name
+		nestedBlocks = append(nestedBlocks, nestedBlock)
+	}
+	return nestedBlocks
+}
+
+func decodeAttribute(schema *SchemaNode, rawAttr any) any {
+	if schema.Value == nil {
+		return rawAttr
+	}
+
+	return rawAttr
+	// TODO: For attributes of type object or map, the schema does not include field names and looks like:
+	// "list_attr": { "references": ["local.foo"] },
+	// Therefore, we cannot determine which specific fields are unknown.
+	// return resolveAttribute(rawAttr, schema.Value)
+}
+
+func resolveAttribute(known, config any) any {
+	switch v := known.(type) {
+	case []any:
+		return v
+	case map[string]any:
+		cm, ok := config.(map[string]any)
+		if !ok {
+			return v
+		}
+
+		result := make(map[string]any)
+		for k, cv := range cm {
+			if vv, exists := result[k]; exists {
+				result[k] = resolveAttribute(vv, cv)
+			} else {
+				result[k] = cv
+			}
+		}
+		return result
+	default:
+		return known
+	}
+}
+
 func unpackConfigurationValue(val any, r Resource) (any, bool) {
 	if t, ok := val.(map[string]any); ok {
 		for k, v := range t {
@@ -204,3 +274,62 @@
 	}
 	return nil
 }
+
+type NodeType = int
+
+const (
+	Attribute NodeType = iota
+	Block
+)
+
+type BlockSchema = map[string]*SchemaNode
+
+type SchemaNode struct {
+	Type     NodeType
+	Children []BlockSchema // only for blocks
+	Value    any           // only for attributes, maybe null
+}
+
+func schemaForBlock(r Resource, expressions map[string]any) BlockSchema {
+	schema := make(map[string]*SchemaNode)
+	for n, expr := range expressions {
+		nodeSchema := schemaForExpression(r, expr)
+		if nodeSchema != nil {
+			schema[n] = nodeSchema
+		}
+	}
+	return schema
+}
+
+func schemaForExpression(r Resource, expr any) *SchemaNode {
+	switch v := expr.(type) {
+	case map[string]any:
+		attrKeys := []string{"constant_value", "references"}
+		for _, k := range attrKeys {
+			if _, exists := v[k]; exists {
+				attrVal, _ := unpackConfigurationValue(v, r)
+				return &SchemaNode{
+					Type:  Attribute,
+					Value: attrVal,
+				}
+			}
+		}
+		return &SchemaNode{
+			Type:     Block,
+			Children: []BlockSchema{schemaForBlock(r, v)},
+		}
+	case []any:
+		children := make([]BlockSchema, 0, len(v))
+		for _, el := range v {
+			if m, ok := el.(map[string]any); ok {
+				children = append(children, schemaForBlock(r, m))
+			}
+		}
+		return &SchemaNode{
+			Type:     Block,
+			Children: children,
+		}
+	}
+
+	return nil
+}

pkg/iac/terraform/resource_block.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"strings"
 	"text/template"
+
+	"github.com/samber/lo"
 )
 
 type PlanReference struct {
@@ -15,10 +17,11 @@ type PlanBlock struct {
 	Type       string
 	Name       string
 	BlockType  string
-	Blocks     map[string]map[string]any
+	Blocks     []PlanBlock
 	Attributes map[string]any
 }
 
+// TODO: unused
 func NewPlanBlock(blockType, resourceType, resourceName string) *PlanBlock {
 	if blockType == "managed" {
 		blockType = "resource"
@@ -28,32 +31,26 @@ func NewPlanBlock(blockType, resourceType, resourceName string) *PlanBlock {
 		Type:       resourceType,
 		Name:       resourceName,
 		BlockType:  blockType,
-		Blocks:     make(map[string]map[string]any),
+		Blocks:     []PlanBlock{},
 		Attributes: make(map[string]any),
 	}
 }
 
 func (rb *PlanBlock) HasAttribute(attribute string) bool {
-	for k := range rb.Attributes {
-		if k == attribute {
-			return true
-		}
-	}
-	return false
+	_, exists := rb.Attributes[attribute]
+	return exists
 }
 
 func (rb *PlanBlock) ToHCL() string {
+	tmpl := template.New("resource").Funcs(template.FuncMap{
+		"RenderValue": renderTemplateValue,
+	})
 
-	resourceTmpl, err := template.New("resource").Funcs(template.FuncMap{
-		"RenderValue":     renderTemplateValue,
-		"RenderPrimitive": renderPrimitive,
-	}).Parse(resourceTemplate)
-	if err != nil {
-		panic(err)
-	}
+	tmpl = lo.Must(tmpl.Parse(resourceTemplate))
+	tmpl = lo.Must(tmpl.Parse(blocksTemplate))
 
 	var res bytes.Buffer
-	if err := resourceTmpl.Execute(&res, map[string]any{
+	if err := tmpl.Execute(&res, map[string]any{
 		"BlockType":  rb.BlockType,
 		"Type":       rb.Type,
 		"Name":       rb.Name,
@@ -66,11 +63,28 @@ func (rb *PlanBlock) ToHCL() string {
 }
 
 var resourceTemplate = `{{ .BlockType }} "{{ .Type }}" "{{ .Name }}" {
-	{{ range $name, $value := .Attributes }}{{ if $value }}{{ $name }} {{ RenderValue $value }}
-	{{end}}{{ end }}{{  range $name, $block := .Blocks }}{{ $name }} {
-	{{ range $name, $value := $block }}{{ if $value }}{{ $name }} {{ RenderValue $value }}
-	{{end}}{{ end }}}
-{{end}}}`
+{{- range $name, $value := .Attributes }}
+  {{- if $value }}
+  {{ $name }} {{ RenderValue $value }}
+  {{- end }}
+{{- end }}
+
+{{- template "renderBlocks" .Blocks }}
+}`
+
+var blocksTemplate = `{{ define "renderBlocks" }}
+{{- range . }}
+{{ .Name }} {
+{{- range $name, $value := .Attributes }}
+  {{- if $value }}
+  {{ $name }} {{ RenderValue $value }}
+  {{- end }}
+{{- end }}
+
+{{- template "renderBlocks" .Blocks }}
+}
+{{- end }}
+{{- end }}`
 
 func renderTemplateValue(val any) string {
 	switch t := val.(type) {