fix(terraform): resolve counted IAM policy documents

assets/queries/terraform/aws/s3_bucket_policy_accepts_http_requests/test/negative9.tf

@@ -0,0 +1,38 @@
+variable "ssmlogs_s3_bucket" {
+  description = "SessionManager log bucket"
+  type        = string
+  default     = ""
+}
+
+resource "aws_s3_bucket_policy" "ssmlogs-s3-bucket" {
+  count  = var.ssmlogs_s3_bucket != "" ? 1 : 0
+  bucket = aws_s3_bucket.ssmlogs-s3-bucket[0].id
+  policy = data.aws_iam_policy_document.ssmlogs-s3-bucket[0].json
+}
+
+data "aws_iam_policy_document" "ssmlogs-s3-bucket" {
+  count = var.ssmlogs_s3_bucket != "" ? 1 : 0
+
+  statement {
+    sid    = "EnforceSSLOnlyRequests"
+    effect = "Deny"
+
+    principals {
+      type        = "*"
+      identifiers = ["*"]
+    }
+
+    actions = ["s3:*"]
+
+    resources = [
+      "arn:aws:s3:::${aws_s3_bucket.ssmlogs-s3-bucket[0].bucket}",
+      "arn:aws:s3:::${aws_s3_bucket.ssmlogs-s3-bucket[0].bucket}/*"
+    ]
+
+    condition {
+      test     = "Bool"
+      variable = "aws:SecureTransport"
+      values   = ["false"]
+    }
+  }
+}

pkg/parser/terraform/data_source.go

@@ -14,7 +14,6 @@ import (
 	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/rs/zerolog/log"
 	"github.com/zclconf/go-cty/cty"
-	"github.com/zclconf/go-cty/cty/gocty"
 	ctyjson "github.com/zclconf/go-cty/cty/json"
 )
 
@@ -74,6 +73,13 @@ type convertedPolicy struct {
 
 var mutexData = &sync.Mutex{}
 
+var dataSourceMetaArguments = map[string]struct{}{
+	"count":      {},
+	"depends_on": {},
+	"for_each":   {},
+	"provider":   {},
+}
+
 func getDataSourcePolicy(currentPath string) {
 	tfFiles, err := filepath.Glob(filepath.Join(currentPath, "*.tf"))
 	if err != nil {
@@ -83,7 +89,7 @@ func getDataSourcePolicy(currentPath string) {
 	if len(tfFiles) == 0 {
 		return
 	}
-	jsonMap := make(map[string]map[string]string)
+	jsonMap := make(map[string]cty.Value)
 	for _, tfFile := range tfFiles {
 		parsedFile, parseErr := parseFile(tfFile, true)
 		if parseErr != nil {
@@ -97,26 +103,37 @@ func getDataSourcePolicy(currentPath string) {
 		for _, block := range body.Blocks {
 			if block.Type == "data" && block.Labels[0] == "aws_iam_policy_document" && len(block.Labels) > 1 {
 				policyJSON := parseDataSourceBody(block.Body)
-				jsonMap[block.Labels[1]] = map[string]string{
-					"json": policyJSON,
+				policyValue := cty.ObjectVal(map[string]cty.Value{
+					"json": cty.StringVal(policyJSON),
+				})
+				if hasCountMetaArgument(block.Body) {
+					policyValue = cty.TupleVal([]cty.Value{policyValue})
 				}
+				jsonMap[block.Labels[1]] = policyValue
 			}
 		}
 	}
-	policyResource := map[string]map[string]map[string]string{
-		"aws_iam_policy_document": jsonMap,
-	}
-	data, err := gocty.ToCtyValue(policyResource, cty.Map(cty.Map(cty.Map(cty.String))))
-	if err != nil {
-		log.Error().Msgf("Error trying to convert policy to cty value: %s", err)
-		return
-	}
+	data := cty.ObjectVal(map[string]cty.Value{
+		"aws_iam_policy_document": objectVal(jsonMap),
+	})
 
 	mutexData.Lock()
 	inputVariableMap["data"] = data
 	mutexData.Unlock()
 }
 
+func hasCountMetaArgument(body *hclsyntax.Body) bool {
+	_, ok := body.Attributes["count"]
+	return ok
+}
+
+func objectVal(values map[string]cty.Value) cty.Value {
+	if len(values) == 0 {
+		return cty.EmptyObjectVal
+	}
+	return cty.ObjectVal(values)
+}
+
 func decodeDataSourcePolicy(value cty.Value) dataSourcePolicy {
 	jsonUnified, err := ctyjson.Marshal(value, cty.DynamicPseudoType)
 	if err != nil {
@@ -233,6 +250,9 @@ func parseDataSourceBody(body *hclsyntax.Body) string {
 	}
 
 	resolveDataResources(body)
+	removedMetaArguments := removeDataSourceMetaArguments(body)
+	defer restoreDataSourceMetaArguments(body, removedMetaArguments)
+
 	paths := extractVariablePathsFromBody(body)
 	grouped := groupPathsByRoot(paths)
 
@@ -307,6 +327,23 @@ func parseDataSourceBody(body *hclsyntax.Body) string {
 	return buffer.String()
 }
 
+func removeDataSourceMetaArguments(body *hclsyntax.Body) map[string]*hclsyntax.Attribute {
+	removed := make(map[string]*hclsyntax.Attribute)
+	for name := range dataSourceMetaArguments {
+		if attr, ok := body.Attributes[name]; ok {
+			removed[name] = attr
+			delete(body.Attributes, name)
+		}
+	}
+	return removed
+}
+
+func restoreDataSourceMetaArguments(body *hclsyntax.Body, removed map[string]*hclsyntax.Attribute) {
+	for name, attr := range removed {
+		body.Attributes[name] = attr
+	}
+}
+
 // groupsPathsByRoot groups paths by their root variable (first element of the path)
 // allowing us to group paths that share the same root for further processing.
 func groupPathsByRoot(paths [][]string) map[string][][]string {

pkg/parser/terraform/data_source_test.go

@@ -6,7 +6,7 @@ import (
 
 	"github.com/Checkmarx/kics/v2/pkg/parser/terraform/converter"
 	"github.com/stretchr/testify/require"
-	"github.com/zclconf/go-cty/cty/gocty"
+	"github.com/zclconf/go-cty/cty"
 )
 
 func Test_getDataSourcePolicy(t *testing.T) {
@@ -44,6 +44,15 @@ func Test_getDataSourcePolicy(t *testing.T) {
 				resourceName: "support_site_bucket_policy_document",
 			},
 			want: `{"Statement":[{"Actions":["s3:GetObject"],"Condition":{"Bool":{"aws:SecureTransport":["true"]}},"Effect":"Allow","Principals":{"AWS":["support_site_origin_access_identity.iam_arn"]},"Resources":["arn:aws:s3:::local.support_site_bucket_name/*"],"Sid":"AllowCloudFrontAccess"}],"Version":"2012-10-17"}
+`,
+		},
+		{
+			name: "should load data source with count as json without errors",
+			args: args{
+				currentPath:  filepath.Join("..", "..", "..", "test", "fixtures", "test_terraform_data_source"),
+				resourceName: "ssmlogs-s3-bucket",
+			},
+			want: `{"Statement":[{"Actions":["s3:*"],"Condition":{"Bool":{"aws:SecureTransport":["false"]}},"Effect":"Deny","Principals":{"*":["*"]},"Resources":["arn:aws:s3:::aws_s3_bucket.ssmlogs-s3-bucket.0.bucket","arn:aws:s3:::aws_s3_bucket.ssmlogs-s3-bucket.0.bucket/*"],"Sid":"EnforceSSLOnlyRequests"}]}
 `,
 		},
 	}
@@ -54,15 +63,7 @@ func Test_getDataSourcePolicy(t *testing.T) {
 			if !ok {
 				t.FailNow()
 			}
-			var awsPolicyMap map[string]map[string]map[string]string
-			err := gocty.FromCtyValue(data, &awsPolicyMap)
-			if err != nil {
-				t.Errorf("getDataSourcePolicy() error = %v", err)
-			}
-			got, ok := awsPolicyMap["aws_iam_policy_document"][tt.args.resourceName]["json"]
-			if !ok {
-				t.FailNow()
-			}
+			got := getDataSourcePolicyJSON(t, data, tt.args.resourceName)
 			require.Equal(t, tt.want, got)
 		})
 	}
@@ -71,3 +72,17 @@ func Test_getDataSourcePolicy(t *testing.T) {
 		inputVariableMap = make(converter.VariableMap)
 	})
 }
+
+func getDataSourcePolicyJSON(t *testing.T, data cty.Value, resourceName string) string {
+	t.Helper()
+
+	policies := data.GetAttr("aws_iam_policy_document").AsValueMap()
+	policy, ok := policies[resourceName]
+	require.True(t, ok)
+
+	if policy.Type().IsTupleType() {
+		policy = policy.Index(cty.NumberIntVal(0))
+	}
+
+	return policy.GetAttr("json").AsString()
+}

test/fixtures/test_terraform_data_source/data_source_4.tf

@@ -0,0 +1,38 @@
+variable "ssmlogs_s3_bucket" {
+  description = "SessionManager log bucket"
+  type        = string
+  default     = ""
+}
+
+resource "aws_s3_bucket_policy" "ssmlogs-s3-bucket" {
+  count  = var.ssmlogs_s3_bucket != "" ? 1 : 0
+  bucket = aws_s3_bucket.ssmlogs-s3-bucket[0].id
+  policy = data.aws_iam_policy_document.ssmlogs-s3-bucket[0].json
+}
+
+data "aws_iam_policy_document" "ssmlogs-s3-bucket" {
+  count = var.ssmlogs_s3_bucket != "" ? 1 : 0
+
+  statement {
+    sid    = "EnforceSSLOnlyRequests"
+    effect = "Deny"
+
+    principals {
+      type        = "*"
+      identifiers = ["*"]
+    }
+
+    actions = ["s3:*"]
+
+    resources = [
+      "arn:aws:s3:::${aws_s3_bucket.ssmlogs-s3-bucket[0].bucket}",
+      "arn:aws:s3:::${aws_s3_bucket.ssmlogs-s3-bucket[0].bucket}/*"
+    ]
+
+    condition {
+      test     = "Bool"
+      variable = "aws:SecureTransport"
+      values   = ["false"]
+    }
+  }
+}