From 45592f225f7eaf530c558369a136225d213fd862 Mon Sep 17 00:00:00 2001
From: James Kwon <96548424+hongil0316@users.noreply.github.com>
Date: Thu, 19 Dec 2024 22:53:16 -0500
Subject: [PATCH] assert SSE

---
 modules/aws/s3.go      | 33 +++++++++++++++++++++++++++++++++
 modules/aws/s3_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/modules/aws/s3.go b/modules/aws/s3.go
index 518d141fe..6fc21ad20 100644
--- a/modules/aws/s3.go
+++ b/modules/aws/s3.go
@@ -478,6 +478,39 @@ 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 fc00ef87f..c375e8a90 100644
--- a/modules/aws/s3_test.go
+++ b/modules/aws/s3_test.go
@@ -267,3 +267,42 @@ func testEmptyBucket(t *testing.T, s3Client *s3.Client, region string, s3BucketN
 	}
 	require.Equal(t, 0, len((*bucketObjects).Contents))
 }
+
+func TestAssertS3BucketServerSideEncryptionE(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) })
+
+			input := &s3.PutBucketEncryptionInput{
+				Bucket: aws.String(s3BucketName),
+				ServerSideEncryptionConfiguration: &types.ServerSideEncryptionConfiguration{
+					Rules: []types.ServerSideEncryptionRule{
+						{
+							ApplyServerSideEncryptionByDefault: &types.ServerSideEncryptionByDefault{
+								SSEAlgorithm: tt,
+							},
+						},
+					},
+				},
+			}
+			_, err := s3client.PutBucketEncryption(context.Background(), input)
+			require.NoError(t, err)
+
+			AssertS3BucketServerSideEncryption(t, region, s3BucketName, tt)
+		})
+	}
+}