From fb3ba0a56d6177f569aacac930eba70b12479e58 Mon Sep 17 00:00:00 2001
From: James Kwon <96548424+hongil0316@users.noreply.github.com>
Date: Wed, 1 Jan 2025 13:21:13 -0500
Subject: [PATCH] add get bucket ownership control

---
 modules/aws/s3.go      | 60 +++++++++++++++++++-----------------------
 modules/aws/s3_test.go | 58 ++++++++++++++++++++++------------------
 2 files changed, 59 insertions(+), 59 deletions(-)

diff --git a/modules/aws/s3.go b/modules/aws/s3.go
index 6fc21ad20..42a93e731 100644
--- a/modules/aws/s3.go
+++ b/modules/aws/s3.go
@@ -420,6 +420,33 @@ func GetS3BucketPolicyE(t testing.TestingT, awsRegion string, bucket string) (st
 	return aws.ToString(res.Policy), nil
 }
 
+func GetS3BucketOwnershipControls(t testing.TestingT, awsRegion, bucket string) []string {
+	rules, err := GetS3BucketOwnershipControlsE(t, awsRegion, bucket)
+	require.NoError(t, err)
+
+	return rules
+}
+
+func GetS3BucketOwnershipControlsE(t testing.TestingT, awsRegion, bucket string) ([]string, error) {
+	s3Client, err := NewS3ClientE(t, awsRegion)
+	if err != nil {
+		return nil, err
+	}
+
+	out, err := s3Client.GetBucketOwnershipControls(context.Background(), &s3.GetBucketOwnershipControlsInput{
+		Bucket: &bucket,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	rules := make([]string, 0, len(out.OwnershipControls.Rules))
+	for _, rule := range out.OwnershipControls.Rules {
+		rules = append(rules, string(rule.ObjectOwnership))
+	}
+	return rules, nil
+}
+
 // AssertS3BucketExists checks if the given S3 bucket exists in the given region and fail the test if it does not.
 func AssertS3BucketExists(t testing.TestingT, region string, name string) {
 	err := AssertS3BucketExistsE(t, region, name)
@@ -478,39 +505,6 @@ func AssertS3BucketPolicyExistsE(t testing.TestingT, region string, bucketName s
 	return nil
 }
 
-// AssertS3BucketServerSideEncryption checks if the given S3 bucket has a server side encryption configured using the given algorithm and fail the test if it does not
-func AssertS3BucketServerSideEncryption(t testing.TestingT, region string, bucketName string, algorithm types.ServerSideEncryption) {
-	err := AssertS3BucketServerSideEncryptionE(t, region, bucketName, algorithm)
-	require.NoError(t, err)
-}
-
-// AssertS3BucketServerSideEncryptionE checks if the given S3 bucket has a server side encryption configured using the given algorithm and returns an error if it does not
-func AssertS3BucketServerSideEncryptionE(t testing.TestingT, region string, bucketName string, algorithm types.ServerSideEncryption) (err error) {
-	s3Client, err := NewS3ClientE(t, region)
-	if err != nil {
-		return err
-	}
-	input := &s3.GetBucketEncryptionInput{
-		Bucket: aws.String(bucketName),
-	}
-	c, err := s3Client.GetBucketEncryption(context.Background(), input)
-	if err != nil {
-		return err
-	}
-
-	err = fmt.Errorf("SSE is not enabled for bucket %s in region %s", bucketName, region)
-	for _, rule := range c.ServerSideEncryptionConfiguration.Rules {
-		if rule.ApplyServerSideEncryptionByDefault == nil {
-			continue
-		}
-		if rule.ApplyServerSideEncryptionByDefault.SSEAlgorithm == algorithm {
-			return nil
-		}
-	}
-	return
-
-}
-
 // NewS3Client creates an S3 client.
 func NewS3Client(t testing.TestingT, region string) *s3.Client {
 	client, err := NewS3ClientE(t, region)
diff --git a/modules/aws/s3_test.go b/modules/aws/s3_test.go
index c375e8a90..63a37f4a5 100644
--- a/modules/aws/s3_test.go
+++ b/modules/aws/s3_test.go
@@ -268,41 +268,47 @@ func testEmptyBucket(t *testing.T, s3Client *s3.Client, region string, s3BucketN
 	require.Equal(t, 0, len((*bucketObjects).Contents))
 }
 
-func TestAssertS3BucketServerSideEncryptionE(t *testing.T) {
+func TestGetS3BucketOwnershipControls(t *testing.T) {
 	t.Parallel()
 
 	region := GetRandomStableRegion(t, nil, nil)
-	s3client := NewS3Client(t, region)
-
 	id := random.UniqueId()
 	logger.Default.Logf(t, "Random values selected. Region = %s, Id = %s\n", region, id)
 
-	table := []types.ServerSideEncryption{
-		types.ServerSideEncryptionAes256,
-		types.ServerSideEncryptionAwsKms,
-	}
-	for i, tt := range table {
-		t.Run(fmt.Sprintf("%s", tt), func(t *testing.T) {
-			s3BucketName := fmt.Sprintf("gruntwork-terratest-sse-%d-%s", i, strings.ToLower(id))
-			CreateS3Bucket(t, region, s3BucketName)
-			t.Cleanup(func() { DeleteS3Bucket(t, region, s3BucketName) })
+	s3BucketName := "gruntwork-terratest-" + strings.ToLower(id)
+	CreateS3Bucket(t, region, s3BucketName)
+	t.Cleanup(func() {
+		DeleteS3Bucket(t, region, s3BucketName)
+	})
 
-			input := &s3.PutBucketEncryptionInput{
-				Bucket: aws.String(s3BucketName),
-				ServerSideEncryptionConfiguration: &types.ServerSideEncryptionConfiguration{
-					Rules: []types.ServerSideEncryptionRule{
-						{
-							ApplyServerSideEncryptionByDefault: &types.ServerSideEncryptionByDefault{
-								SSEAlgorithm: tt,
-							},
-						},
+	t.Run("Exist", func(t *testing.T) {
+		s3Client, err := NewS3ClientE(t, region)
+		require.NoError(t, err)
+		_, err = s3Client.PutBucketOwnershipControls(context.Background(), &s3.PutBucketOwnershipControlsInput{
+			Bucket: &s3BucketName,
+			OwnershipControls: &types.OwnershipControls{
+				Rules: []types.OwnershipControlsRule{
+					{
+						ObjectOwnership: types.ObjectOwnershipBucketOwnerEnforced,
 					},
 				},
-			}
-			_, err := s3client.PutBucketEncryption(context.Background(), input)
+			},
+		})
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			_, err := s3Client.DeleteBucketOwnershipControls(context.Background(), &s3.DeleteBucketOwnershipControlsInput{
+				Bucket: &s3BucketName,
+			})
 			require.NoError(t, err)
-
-			AssertS3BucketServerSideEncryption(t, region, s3BucketName, tt)
 		})
-	}
+
+		controls := GetS3BucketOwnershipControls(t, region, s3BucketName)
+		assert.Equal(t, 1, len(controls))
+		assert.Equal(t, string(types.ObjectOwnershipBucketOwnerEnforced), controls[0])
+	})
+
+	t.Run("NotExist", func(t *testing.T) {
+		_, err := GetS3BucketOwnershipControlsE(t, region, s3BucketName)
+		assert.Error(t, err)
+	})
 }