enable KMS encryption to test in dev env

uc-cdis · Jan 23, 2025 · 212181e · 212181e
1 parent 6f15910
commit 212181e
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 115 deletions.
diff --git a/gen3workflow/aws_utils.py b/gen3workflow/aws_utils.py
@@ -107,62 +107,62 @@ def create_user_bucket(user_id: str) -> Tuple[str, str, str]:
         logger.debug(f"Created KMS key alias '{kms_key_alias}'")
 
     # TODO enable KMS encryption when Funnel workers can push with KMS key or use our S3 endpoint
-    # logger.debug(f"Setting KMS encryption on bucket '{user_bucket_name}'")
-    # s3_client.put_bucket_encryption(
-    #     Bucket=user_bucket_name,
-    #     ServerSideEncryptionConfiguration={
-    #         "Rules": [
-    #             {
-    #                 "ApplyServerSideEncryptionByDefault": {
-    #                     "SSEAlgorithm": "aws:kms",
-    #                     "KMSMasterKeyID": kms_key_arn,
-    #                 },
-    #                 "BucketKeyEnabled": True,
-    #             },
-    #         ],
-    #     },
-    # )
-
-    # logger.debug("Enforcing KMS encryption through bucket policy")
-    # s3_client.put_bucket_policy(
-    #     Bucket=user_bucket_name,
-    #     # using 2 statements here, because for some reason the condition below allows using a
-    #     # different key as long as "s3:x-amz-server-side-encryption: aws:kms" is specified:
-    #     # "StringNotEquals": {
-    #     #     "s3:x-amz-server-side-encryption": "aws:kms",
-    #     #     "s3:x-amz-server-side-encryption-aws-kms-key-id": "{kms_key_arn}"
-    #     # }
-    #     Policy=f"""{{
-    #         "Version": "2012-10-17",
-    #         "Statement": [
-    #             {{
-    #                 "Sid": "RequireKMSEncryption",
-    #                 "Effect": "Deny",
-    #                 "Principal": "*",
-    #                 "Action": "s3:PutObject",
-    #                 "Resource": "arn:aws:s3:::{user_bucket_name}/*",
-    #                 "Condition": {{
-    #                     "StringNotEquals": {{
-    #                         "s3:x-amz-server-side-encryption": "aws:kms"
-    #                     }}
-    #                 }}
-    #             }},
-    #             {{
-    #                 "Sid": "RequireSpecificKMSKey",
-    #                 "Effect": "Deny",
-    #                 "Principal": "*",
-    #                 "Action": "s3:PutObject",
-    #                 "Resource": "arn:aws:s3:::{user_bucket_name}/*",
-    #                 "Condition": {{
-    #                     "StringNotEquals": {{
-    #                         "s3:x-amz-server-side-encryption-aws-kms-key-id": "{kms_key_arn}"
-    #                     }}
-    #                 }}
-    #             }}
-    #         ]
-    #     }}
-    #     """,
-    # )
+    logger.debug(f"Setting KMS encryption on bucket '{user_bucket_name}'")
+    s3_client.put_bucket_encryption(
+        Bucket=user_bucket_name,
+        ServerSideEncryptionConfiguration={
+            "Rules": [
+                {
+                    "ApplyServerSideEncryptionByDefault": {
+                        "SSEAlgorithm": "aws:kms",
+                        "KMSMasterKeyID": kms_key_arn,
+                    },
+                    "BucketKeyEnabled": True,
+                },
+            ],
+        },
+    )
+
+    logger.debug("Enforcing KMS encryption through bucket policy")
+    s3_client.put_bucket_policy(
+        Bucket=user_bucket_name,
+        # using 2 statements here, because for some reason the condition below allows using a
+        # different key as long as "s3:x-amz-server-side-encryption: aws:kms" is specified:
+        # "StringNotEquals": {
+        #     "s3:x-amz-server-side-encryption": "aws:kms",
+        #     "s3:x-amz-server-side-encryption-aws-kms-key-id": "{kms_key_arn}"
+        # }
+        Policy=f"""{{
+            "Version": "2012-10-17",
+            "Statement": [
+                {{
+                    "Sid": "RequireKMSEncryption",
+                    "Effect": "Deny",
+                    "Principal": "*",
+                    "Action": "s3:PutObject",
+                    "Resource": "arn:aws:s3:::{user_bucket_name}/*",
+                    "Condition": {{
+                        "StringNotEquals": {{
+                            "s3:x-amz-server-side-encryption": "aws:kms"
+                        }}
+                    }}
+                }},
+                {{
+                    "Sid": "RequireSpecificKMSKey",
+                    "Effect": "Deny",
+                    "Principal": "*",
+                    "Action": "s3:PutObject",
+                    "Resource": "arn:aws:s3:::{user_bucket_name}/*",
+                    "Condition": {{
+                        "StringNotEquals": {{
+                            "s3:x-amz-server-side-encryption-aws-kms-key-id": "{kms_key_arn}"
+                        }}
+                    }}
+                }}
+            ]
+        }}
+        """,
+    )
 
     expiration_days = config["S3_OBJECTS_EXPIRATION_DAYS"]
     logger.debug(f"Setting bucket objects expiration to {expiration_days} days")

diff --git a/gen3workflow/routes/s3.py b/gen3workflow/routes/s3.py
@@ -136,11 +136,11 @@ async def s3_endpoint(path: str, request: Request):
         headers["x-amz-security-token"] = credentials.token
 
     # TODO enable KMS encryption when Funnel workers can push with KMS key or use our S3 endpoint
-    # # if this is a PUT request, we need the KMS key ID to use for encryption
-    # if request.method == "PUT":
-    #     _, kms_key_arn = aws_utils.get_existing_kms_key_for_bucket(user_bucket, user_id)
-    #     headers["x-amz-server-side-encryption"] = "aws:kms"
-    #     headers["x-amz-server-side-encryption-aws-kms-key-id"] = kms_key_arn
+    # if this is a PUT request, we need the KMS key ID to use for encryption
+    if request.method == "PUT":
+        _, kms_key_arn = aws_utils.get_existing_kms_key_for_bucket(user_bucket, user_id)
+        headers["x-amz-server-side-encryption"] = "aws:kms"
+        headers["x-amz-server-side-encryption-aws-kms-key-id"] = kms_key_arn
 
     # construct the canonical request
     canonical_headers = "".join(

diff --git a/tests/test_misc.py b/tests/test_misc.py
@@ -65,59 +65,55 @@ async def test_storage_info(client, access_token_patcher, mock_aws_services):
     }
 
     # TODO enable when KMS encryption is enabled
-    # # check that the bucket is setup with KMS encryption
-    # kms_key = aws_utils.kms_client.describe_key(
-    #     KeyId=f"alias/key-{expected_bucket_name}"
-    # )
-    # kms_key_arn = kms_key["KeyMetadata"]["Arn"]
-    # bucket_encryption = aws_utils.s3_client.get_bucket_encryption(
-    #     Bucket=expected_bucket_name
-    # )
-    # assert bucket_encryption.get("ServerSideEncryptionConfiguration") == {
-    #     "Rules": [
-    #         {
-    #             "ApplyServerSideEncryptionByDefault": {
-    #                 "SSEAlgorithm": "aws:kms",
-    #                 "KMSMasterKeyID": kms_key_arn,
-    #             },
-    #             "BucketKeyEnabled": True,
-    #         }
-    #     ]
-    # }
-
-    # # check the bucket policy, which should enforce KMS encryption
-    # bucket_policy = aws_utils.s3_client.get_bucket_policy(
-    #     Bucket=expected_bucket_name
-    # )
-    # assert json.loads(bucket_policy.get("Policy", "{}")) == {
-    #     "Version": "2012-10-17",
-    #     "Statement": [
-    #         {
-    #             "Sid": "RequireKMSEncryption",
-    #             "Effect": "Deny",
-    #             "Principal": "*",
-    #             "Action": "s3:PutObject",
-    #             "Resource": "arn:aws:s3:::gen3wf-localhost-64/*",
-    #             "Condition": {
-    #                 "StringNotEquals": {
-    #                     's3:x-amz-server-side-encryption': 'aws:kms'
-    #                 }
-    #             },
-    #         },
-    #         {
-    #             "Sid": "RequireSpecificKMSKey",
-    #             "Effect": "Deny",
-    #             "Principal": "*",
-    #             "Action": "s3:PutObject",
-    #             "Resource": "arn:aws:s3:::gen3wf-localhost-64/*",
-    #             "Condition": {
-    #                 "StringNotEquals": {
-    #                     "s3:x-amz-server-side-encryption-aws-kms-key-id": kms_key_arn
-    #                 }
-    #             },
-    #         }
-    #     ],
-    # }
+    # check that the bucket is setup with KMS encryption
+    kms_key = aws_utils.kms_client.describe_key(
+        KeyId=f"alias/key-{expected_bucket_name}"
+    )
+    kms_key_arn = kms_key["KeyMetadata"]["Arn"]
+    bucket_encryption = aws_utils.s3_client.get_bucket_encryption(
+        Bucket=expected_bucket_name
+    )
+    assert bucket_encryption.get("ServerSideEncryptionConfiguration") == {
+        "Rules": [
+            {
+                "ApplyServerSideEncryptionByDefault": {
+                    "SSEAlgorithm": "aws:kms",
+                    "KMSMasterKeyID": kms_key_arn,
+                },
+                "BucketKeyEnabled": True,
+            }
+        ]
+    }
+
+    # check the bucket policy, which should enforce KMS encryption
+    bucket_policy = aws_utils.s3_client.get_bucket_policy(Bucket=expected_bucket_name)
+    assert json.loads(bucket_policy.get("Policy", "{}")) == {
+        "Version": "2012-10-17",
+        "Statement": [
+            {
+                "Sid": "RequireKMSEncryption",
+                "Effect": "Deny",
+                "Principal": "*",
+                "Action": "s3:PutObject",
+                "Resource": "arn:aws:s3:::gen3wf-localhost-64/*",
+                "Condition": {
+                    "StringNotEquals": {"s3:x-amz-server-side-encryption": "aws:kms"}
+                },
+            },
+            {
+                "Sid": "RequireSpecificKMSKey",
+                "Effect": "Deny",
+                "Principal": "*",
+                "Action": "s3:PutObject",
+                "Resource": "arn:aws:s3:::gen3wf-localhost-64/*",
+                "Condition": {
+                    "StringNotEquals": {
+                        "s3:x-amz-server-side-encryption-aws-kms-key-id": kms_key_arn
+                    }
+                },
+            },
+        ],
+    }
 
     # check the bucket's lifecycle configuration
     lifecycle_config = aws_utils.s3_client.get_bucket_lifecycle_configuration(
@@ -133,7 +129,7 @@ async def test_storage_info(client, access_token_patcher, mock_aws_services):
     ]
 
 
-@pytest.mark.skip(reason="TODO enable when KMS encryption is enabled")
+# @pytest.mark.skip(reason="TODO enable when KMS encryption is enabled")
 @pytest.mark.asyncio
 async def test_bucket_enforces_encryption(
     client, access_token_patcher, mock_aws_services