chore: consolidate custom resources, additional cleanup

source/patterns/@aws-solutions-constructs/aws-cloudfront-s3/lib/index.ts

@@ -17,11 +17,9 @@ import * as s3 from 'aws-cdk-lib/aws-s3';
 // Note: To ensure CDKv2 compatibility, keep the import statement for Construct separate
 import { Construct } from 'constructs';
 import * as defaults from '@aws-solutions-constructs/core';
-import { CustomResource, aws_iam } from 'aws-cdk-lib';
-import { Provider } from 'aws-cdk-lib/custom-resources';
+import { aws_iam } from 'aws-cdk-lib';
 import { IKey } from 'aws-cdk-lib/aws-kms';
-import { Code, Runtime } from 'aws-cdk-lib/aws-lambda';
-import { Effect, PolicyDocument, PolicyStatement, Role, ServicePrincipal } from 'aws-cdk-lib/aws-iam';
+import * as resources from '@aws-solutions-constructs/resources';
 
 /**
  * @summary The properties for the CloudFrontToS3 Construct
@@ -189,52 +187,10 @@ export class CloudFrontToS3 extends Construct {
     }
 
     if (encryptionKey) {
-      const lambdaHandler = defaults.buildLambdaFunction(this, {
-        lambdaFunctionProps: {
-          runtime: Runtime.NODEJS_18_X,
-          handler: 'index.handler',
-          description: 'kms-key-policy-updater',
-          code: Code.fromAsset(`${__dirname}/../../resources/kms-key-policy-updater`),
-          role: new Role(this, 'KmsKeyPolicyUpdateLambdaRole', {
-            assumedBy: new ServicePrincipal('lambda.amazonaws.com'),
-            description: 'Role to update kms key policy to allow cloudfront access',
-            inlinePolicies: {
-              KmsPolicy: new PolicyDocument({
-                statements: [
-                  new PolicyStatement({
-                    actions: ['kms:PutKeyPolicy', 'kms:GetKeyPolicy', 'kms:DescribeKey'],
-                    effect: Effect.ALLOW,
-                    resources: [ encryptionKey.keyArn ]
-                  })
-                ]
-              }),
-              CWLogsPolicy: new PolicyDocument({
-                statements: [
-                  new PolicyStatement({
-                    actions: ['logs:CreateLogGroup'],
-                    effect: Effect.ALLOW,
-                    resources: [ `arn:${Aws.PARTITION}:logs:${Aws.REGION}:${Aws.ACCOUNT_ID}:*` ]
-                  })
-                ]
-              })
-            }
-          })
-        }
-      });
-
-      const kmsKeyPolicyUpdateProvider = new Provider(this, 'KmsKeyPolicyUpdateProvider', {
-        onEventHandler: lambdaHandler
-      });
-
-      new CustomResource(this, 'KmsKeyPolicyUpdater', {
-        resourceType: 'Custom::KmsKeyPolicyUpdater',
-        serviceToken: kmsKeyPolicyUpdateProvider.serviceToken,
-        properties: {
-          KmsKeyId: encryptionKey.keyId,
-          CloudFrontDistributionId: this.cloudFrontWebDistribution.distributionId,
-          AccountId: Aws.ACCOUNT_ID
-        },
+      resources.createKeyPolicyUpdaterCustomResource(this, {
+        distribution: this.cloudFrontWebDistribution,
+        encryptionKey
       });
     }
   }
-}
+}

source/patterns/@aws-solutions-constructs/resources/.eslintignore

@@ -1,5 +1,3 @@
-index.js
-lib/*.js
-test/*.js
-*.d.ts
-coverage
+**/*.js
+**/*.d.ts
+coverage

source/patterns/@aws-solutions-constructs/resources/.gitignore

@@ -1,9 +1,8 @@
 index.js
-lib/*.js
-test/*.js
 !test/lambda*/*
 *.js.map
-*.d.ts
+**/*.d.ts
+**/*.js
 node_modules
 *.generated.ts
 dist

source/patterns/@aws-solutions-constructs/resources/index.ts

@@ -11,5 +11,5 @@
  *  and limitations under the License.
  */
 
-export * from './kms-key-policy-updater/index';
-export * from './template-writer/lib/template-writer';
+export * from './kms-key-policy-updater';
+export * from './template-writer';

source/patterns/@aws-solutions-constructs/resources/kms-key-policy-updater/index.ts

@@ -11,99 +11,85 @@
  *  and limitations under the License.
  */
 
-import { KMSClient, GetKeyPolicyCommand, DescribeKeyCommand, PutKeyPolicyCommand, KeyManagerType } from "@aws-sdk/client-kms";
+import * as cdk from "aws-cdk-lib";
+import * as iam from 'aws-cdk-lib/aws-iam';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
+import { Aws, CustomResource } from 'aws-cdk-lib';
+import { Provider } from "aws-cdk-lib/custom-resources";
+import { buildLambdaFunction } from "@aws-solutions-constructs/core";
+import { IKey } from "aws-cdk-lib/aws-kms";
+import { Distribution } from "aws-cdk-lib/aws-cloudfront";
+// Note: To ensure CDKv2 compatibility, keep the import statement for Construct separate
+import { Construct } from 'constructs';
 
-const kmsClient = new KMSClient();
+export interface CreateKeyPolicyUpdaterResponse {
+  readonly lambdaFunction: lambda.Function;
+  readonly customResource: CustomResource;
+}
 
-export const handler = async (e: any, context: any) => {
+export interface KeyPolicyUpdaterProps {
+  readonly encryptionKey: IKey;
+  readonly distribution: Distribution;
+  readonly timeout?: cdk.Duration;
+  readonly memorySize?: number;
+}
 
-  let status = 'SUCCESS';
-  let responseData = {};
+export function createKeyPolicyUpdaterCustomResource(
+  scope: Construct,
+  props: KeyPolicyUpdaterProps
+): CreateKeyPolicyUpdaterResponse {
 
-  if (e.RequestType === 'Create' || e.RequestType === 'Update') {
-
-    try {
-      const kmsKeyId = e.ResourceProperties.KmsKeyId;
-      const cloudFrontDistributionId = e.ResourceProperties.CloudFrontDistributionId;
-      const accountId = e.ResourceProperties.AccountId;
-      const region = process.env.AWS_REGION;
-
-      const describeKeyCommandResponse = await kmsClient.send(new DescribeKeyCommand({
-        KeyId: kmsKeyId
-      }));
-
-      if (describeKeyCommandResponse.KeyMetadata?.KeyManager === KeyManagerType.AWS) {
-        return {
-          Status: 'SUCCESS',
-          Reason: 'An AWS managed key was provided, no action needed from the custom resource, exiting now.',
-          PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
-          StackId: e.StackId,
-          RequestId: e.RequestId,
-          LogicalResourceId: e.LogicalResourceId,
-          Data: 'An AWS managed key was provided, no action needed from the custom resource, exiting now.',
-        };
-      }
-
-      const getKeyPolicyCommandResponse = await kmsClient.send(new GetKeyPolicyCommand({
-        KeyId: kmsKeyId,
-        PolicyName: 'default'
-      }));
-
-      if (!getKeyPolicyCommandResponse.Policy) {
-        return {
-          Status: 'FAILED',
-          Reason: 'An error occurred while retrieving the key policy',
-          PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
-          StackId: e.StackId,
-          RequestId: e.RequestId,
-          LogicalResourceId: e.LogicalResourceId,
-          Data: 'An error occurred while retrieving the key policy',
-        };
-      }
-
-      const keyPolicy = JSON.parse(getKeyPolicyCommandResponse?.Policy);
-
-      // Update the existing key policy to allow the cloudfront distribution to use the key
-      keyPolicy.Statement.push({
-        Sid: 'Grant-CloudFront-Distribution-Key-Usage',
-        Effect: 'Allow',
-        Principal: {
-          Service: 'cloudfront.amazonaws.com',
-        },
-        Action: [
-          'kms:Decrypt',
-          'kms:Encrypt',
-          'kms:GenerateDataKey*',
-          'kms:ReEncrypt*'
-        ],
-        Resource: `arn:aws:kms:${region}:${accountId}:key/${kmsKeyId}`,
-        Condition: {
-          StringEquals: {
-            'AWS:SourceArn': `arn:aws:cloudfront::${accountId}:distribution/${cloudFrontDistributionId}`
-          }
+  const lambdaFunction = buildLambdaFunction(scope, {
+    lambdaFunctionProps: {
+      runtime: lambda.Runtime.NODEJS_18_X,
+      handler: 'index.handler',
+      description: 'kms-key-policy-updater',
+      timeout: props.timeout,
+      memorySize: props.memorySize,
+      code: lambda.Code.fromAsset(`${__dirname}/custom-resource/`),
+      role: new iam.Role(scope, 'KmsKeyPolicyUpdateLambdaRole', {
+        assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
+        description: 'Role to update kms key policy to allow cloudfront access',
+        inlinePolicies: {
+          KmsPolicy: new iam.PolicyDocument({
+            statements: [
+              new iam.PolicyStatement({
+                actions: ['kms:PutKeyPolicy', 'kms:GetKeyPolicy', 'kms:DescribeKey'],
+                effect: iam.Effect.ALLOW,
+                resources: [ props.encryptionKey.keyArn ]
+              })
+            ]
+          }),
+          CWLogsPolicy: new iam.PolicyDocument({
+            statements: [
+              new iam.PolicyStatement({
+                actions: ['logs:CreateLogGroup'],
+                effect: iam.Effect.ALLOW,
+                resources: [ `arn:${Aws.PARTITION}:logs:${Aws.REGION}:${Aws.ACCOUNT_ID}:*` ]
+              })
+            ]
+          })
         }
-      });
-
-      await kmsClient.send(new PutKeyPolicyCommand({
-        KeyId: kmsKeyId,
-        Policy: JSON.stringify(keyPolicy),
-        PolicyName: 'default'
-      }));
-    } catch (err) {
-      status = 'FAILED';
-      responseData = {
-        Error: JSON.stringify(err)
-      };
+      })
     }
-  }
+  });
+
+  const kmsKeyPolicyUpdateProvider = new Provider(scope, 'KmsKeyPolicyUpdateProvider', {
+    onEventHandler: lambdaFunction
+  });
+
+  const customResource = new CustomResource(scope, 'KmsKeyPolicyUpdater', {
+    resourceType: 'Custom::KmsKeyPolicyUpdater',
+    serviceToken: kmsKeyPolicyUpdateProvider.serviceToken,
+    properties: {
+      KmsKeyId: props.encryptionKey.keyId,
+      CloudFrontDistributionId: props.distribution.distributionId,
+      AccountId: Aws.ACCOUNT_ID
+    },
+  });
 
   return {
-    Status: status,
-    Reason: JSON.stringify(responseData),
-    PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
-    StackId: e.StackId,
-    RequestId: e.RequestId,
-    LogicalResourceId: e.LogicalResourceId,
-    Data: responseData,
+    lambdaFunction,
+    customResource
   };
-};
+}

source/patterns/@aws-solutions-constructs/resources/kms-key-policy-updater/lib/index.ts

@@ -0,0 +1,109 @@
+/**
+ *  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ *  with the License. A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+ *  OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+ *  and limitations under the License.
+ */
+
+import { KMSClient, GetKeyPolicyCommand, DescribeKeyCommand, PutKeyPolicyCommand, KeyManagerType } from "@aws-sdk/client-kms";
+
+const kmsClient = new KMSClient();
+
+export const handler = async (e: any, context: any) => {
+
+  let status = 'SUCCESS';
+  let responseData = {};
+
+  if (e.RequestType === 'Create' || e.RequestType === 'Update') {
+
+    try {
+      const kmsKeyId = e.ResourceProperties.KmsKeyId;
+      const cloudFrontDistributionId = e.ResourceProperties.CloudFrontDistributionId;
+      const accountId = e.ResourceProperties.AccountId;
+      const region = process.env.AWS_REGION;
+
+      const describeKeyCommandResponse = await kmsClient.send(new DescribeKeyCommand({
+        KeyId: kmsKeyId
+      }));
+
+      if (describeKeyCommandResponse.KeyMetadata?.KeyManager === KeyManagerType.AWS) {
+        return {
+          Status: 'SUCCESS',
+          Reason: 'An AWS managed key was provided, no action needed from the custom resource, exiting now.',
+          PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
+          StackId: e.StackId,
+          RequestId: e.RequestId,
+          LogicalResourceId: e.LogicalResourceId,
+          Data: 'An AWS managed key was provided, no action needed from the custom resource, exiting now.',
+        };
+      }
+
+      const getKeyPolicyCommandResponse = await kmsClient.send(new GetKeyPolicyCommand({
+        KeyId: kmsKeyId,
+        PolicyName: 'default'
+      }));
+
+      if (!getKeyPolicyCommandResponse.Policy) {
+        return {
+          Status: 'FAILED',
+          Reason: 'An error occurred while retrieving the key policy',
+          PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
+          StackId: e.StackId,
+          RequestId: e.RequestId,
+          LogicalResourceId: e.LogicalResourceId,
+          Data: 'An error occurred while retrieving the key policy',
+        };
+      }
+
+      const keyPolicy = JSON.parse(getKeyPolicyCommandResponse?.Policy);
+
+      // Update the existing key policy to allow the cloudfront distribution to use the key
+      keyPolicy.Statement.push({
+        Sid: 'Grant-CloudFront-Distribution-Key-Usage',
+        Effect: 'Allow',
+        Principal: {
+          Service: 'cloudfront.amazonaws.com',
+        },
+        Action: [
+          'kms:Decrypt',
+          'kms:Encrypt',
+          'kms:GenerateDataKey*',
+          'kms:ReEncrypt*'
+        ],
+        Resource: `arn:aws:kms:${region}:${accountId}:key/${kmsKeyId}`,
+        Condition: {
+          StringEquals: {
+            'AWS:SourceArn': `arn:aws:cloudfront::${accountId}:distribution/${cloudFrontDistributionId}`
+          }
+        }
+      });
+
+      await kmsClient.send(new PutKeyPolicyCommand({
+        KeyId: kmsKeyId,
+        Policy: JSON.stringify(keyPolicy),
+        PolicyName: 'default'
+      }));
+    } catch (err) {
+      status = 'FAILED';
+      responseData = {
+        Error: JSON.stringify(err)
+      };
+    }
+  }
+
+  return {
+    Status: status,
+    Reason: JSON.stringify(responseData),
+    PhysicalResourceId: e.PhysicalResourceId ?? context.logStreamName,
+    StackId: e.StackId,
+    RequestId: e.RequestId,
+    LogicalResourceId: e.LogicalResourceId,
+    Data: responseData,
+  };
+};

source/patterns/@aws-solutions-constructs/resources/kms-key-policy-updater/tests/kms-key-policy-updater.test.ts

@@ -11,9 +11,9 @@
  *  and limitations under the License.
  */
 
-import { handler } from "..";
 import { mockClient } from "aws-sdk-client-mock";
 import { KMSClient, DescribeKeyCommand, KeyManagerType, GetKeyPolicyCommand, PutKeyPolicyCommand } from "@aws-sdk/client-kms";
+import { handler } from "../lib";
 
 const kmsMock = mockClient(KMSClient);