KMS_ENCRYPTION_ENABLED setting

uc-cdis · Jan 27, 2025 · 8d9cc5b · 8d9cc5b
1 parent ad8c4f9
commit 8d9cc5b
Show file tree

Hide file tree

Showing 6 changed files with 147 additions and 141 deletions.
diff --git a/.secrets.baseline b/.secrets.baseline
@@ -169,7 +169,7 @@
         "filename": "gen3workflow/config-default.yaml",
         "hashed_secret": "afc848c316af1a89d49826c5ae9d00ed769415f3",
         "is_verified": false,
-        "line_number": 39
+        "line_number": 41
       }
     ],
     "migrations/versions/e1886270d9d2_create_system_key_table.py": [
@@ -209,5 +209,5 @@
       }
     ]
   },
-  "generated_at": "2025-01-24T23:29:16Z"
+  "generated_at": "2025-01-27T18:55:41Z"
 }
diff --git a/gen3workflow/aws_utils.py b/gen3workflow/aws_utils.py
@@ -60,6 +60,86 @@ def get_existing_kms_key_for_bucket(bucket_name):
         raise
 
 
+def setup_kms_encryption_on_bucket(bucket_name):
+    # set up KMS encryption on the bucket.
+    # the only way to check if the KMS key has already been created is to use an alias
+    kms_key_alias, kms_key_arn = get_existing_kms_key_for_bucket(bucket_name)
+    if kms_key_arn:
+        logger.debug(f"Existing KMS key '{kms_key_alias}' - '{kms_key_arn}'")
+    else:
+        # the KMS key doesn't exist: create it
+        output = kms_client.create_key(
+            Tags=[
+                {
+                    "TagKey": "Name",
+                    "TagValue": get_safe_name_from_hostname(user_id=None),
+                }
+            ]
+        )
+        kms_key_arn = output["KeyMetadata"]["Arn"]
+        logger.debug(f"Created KMS key '{kms_key_arn}'")
+
+        kms_client.create_alias(AliasName=kms_key_alias, TargetKeyId=kms_key_arn)
+        logger.debug(f"Created KMS key alias '{kms_key_alias}'")
+
+    logger.debug(f"Setting KMS encryption on bucket '{bucket_name}'")
+    s3_client.put_bucket_encryption(
+        Bucket=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=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:::{bucket_name}/*",
+                    "Condition": {{
+                        "StringNotEquals": {{
+                            "s3:x-amz-server-side-encryption": "aws:kms"
+                        }}
+                    }}
+                }},
+                {{
+                    "Sid": "RequireSpecificKMSKey",
+                    "Effect": "Deny",
+                    "Principal": "*",
+                    "Action": "s3:PutObject",
+                    "Resource": "arn:aws:s3:::{bucket_name}/*",
+                    "Condition": {{
+                        "StringNotEquals": {{
+                            "s3:x-amz-server-side-encryption-aws-kms-key-id": "{kms_key_arn}"
+                        }}
+                    }}
+                }}
+            ]
+        }}
+        """,
+    )
+
+
 def create_user_bucket(user_id: str) -> Tuple[str, str, str]:
     """
     Create an S3 bucket for the specified user and return information about the bucket.
@@ -83,84 +163,10 @@ def create_user_bucket(user_id: str) -> Tuple[str, str, str]:
         )
     logger.info(f"Created S3 bucket '{user_bucket_name}' for user '{user_id}'")
 
-    # TODO enable KMS encryption when Funnel workers can push with KMS key or use our S3 endpoint
-    # # set up KMS encryption on the bucket.
-    # # the only way to check if the KMS key has already been created is to use an alias
-    # kms_key_alias, kms_key_arn = get_existing_kms_key_for_bucket(user_bucket_name)
-    # if kms_key_arn:
-    #     logger.debug(f"Existing KMS key '{kms_key_alias}' - '{kms_key_arn}'")
-    # else:
-    #     # the KMS key doesn't exist: create it
-    #     output = kms_client.create_key(
-    #         Tags=[
-    #             {
-    #                 "TagKey": "Name",
-    #                 "TagValue": get_safe_name_from_hostname(user_id=None),
-    #             }
-    #         ]
-    #     )
-    #     kms_key_arn = output["KeyMetadata"]["Arn"]
-    #     logger.debug(f"Created KMS key '{kms_key_arn}'")
-
-    #     kms_client.create_alias(AliasName=kms_key_alias, TargetKeyId=kms_key_arn)
-    #     logger.debug(f"Created KMS key alias '{kms_key_alias}'")
-
-    # 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}"
-    #                     }}
-    #                 }}
-    #             }}
-    #         ]
-    #     }}
-    #     """,
-    # )
+    if config["KMS_ENCRYPTION_ENABLED"]:
+        setup_kms_encryption_on_bucket(user_bucket_name)
+    else:
+        logger.debug(f"Skipping KMS encryption setup on bucket '{user_bucket_name}'")
 
     expiration_days = config["S3_OBJECTS_EXPIRATION_DAYS"]
     logger.debug(f"Setting bucket objects expiration to {expiration_days} days")

diff --git a/gen3workflow/config-default.yaml b/gen3workflow/config-default.yaml
@@ -28,6 +28,8 @@ S3_OBJECTS_EXPIRATION_DAYS: 30
 S3_ENDPOINTS_AWS_ACCESS_KEY_ID:
 S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY:
 
+KMS_ENCRYPTION_ENABLED: true
+
 #############
 # DATABASE #
 #############

diff --git a/gen3workflow/config.py b/gen3workflow/config.py
@@ -51,6 +51,7 @@ def validate_top_level_configs(self) -> None:
                 "S3_OBJECTS_EXPIRATION_DAYS": {"type": "integer", "minimum": 1},
                 "S3_ENDPOINTS_AWS_ACCESS_KEY_ID": {"type": ["string", "null"]},
                 "S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY": {"type": ["string", "null"]},
+                "KMS_ENCRYPTION_ENABLED": {"type": "boolean"},
                 "DB_DRIVER": {"type": "string"},
                 "DB_HOST": {"type": "string"},
                 "DB_PORT": {"type": "integer"},

diff --git a/gen3workflow/routes/s3.py b/gen3workflow/routes/s3.py
@@ -135,18 +135,17 @@ async def s3_endpoint(path: str, request: Request):
         assert credentials, "No AWS credentials found"
         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)
-    #     if not kms_key_arn:
-    #         err_msg = "Bucket misconfigured. Hit the `GET /storage/info` endpoint and try again."
-    #         logger.error(
-    #             f"No existing KMS key found for bucket '{user_bucket}'. {err_msg}"
-    #         )
-    #         raise HTTPException(HTTP_400_BAD_REQUEST, err_msg)
-    #     headers["x-amz-server-side-encryption"] = "aws:kms"
-    #     headers["x-amz-server-side-encryption-aws-kms-key-id"] = kms_key_arn
+    if config["KMS_ENCRYPTION_ENABLED"] and request.method == "PUT":
+        _, kms_key_arn = aws_utils.get_existing_kms_key_for_bucket(user_bucket)
+        if not kms_key_arn:
+            err_msg = "Bucket misconfigured. Hit the `GET /storage/info` endpoint and try again."
+            logger.error(
+                f"No existing KMS key found for bucket '{user_bucket}'. {err_msg}"
+            )
+            raise HTTPException(HTTP_400_BAD_REQUEST, err_msg)
+        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
@@ -62,56 +62,55 @@ async def test_storage_info(client, access_token_patcher, mock_aws_services):
         "region": config["USER_BUCKETS_REGION"],
     }
 
-    # 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
-    #                 }
-    #             },
-    #         },
-    #     ],
-    # }
+    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(
@@ -127,7 +126,6 @@ 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.asyncio
 async def test_bucket_enforces_encryption(
     client, access_token_patcher, mock_aws_services