From ac5b229b2ffbb39e69d7fbd4eee87523d20ceb73 Mon Sep 17 00:00:00 2001 From: Leon Luttenberger Date: Mon, 11 Mar 2024 12:50:16 -0500 Subject: [PATCH] chore: Set up CDK Nag (#25) * add CDK nag for fmops * Reduce the scope of CDK nag suppressions for mlflow modules * Reduce the scope of CDK nag suppressions for endpoint module * fix formatting * add CDK nag to sagemaker studio module * fix static checkers * fix mypy error * add cdk nag to the templates module * Add CDK nag to sagemaker-notebook * add CDK nag unit test checks * fix formatting for fmops module * fix sagemaker templates tests * add CDK nag unit test for mlflow-farhate * add CDK nag unit test for mlflow-image * add CDK nag unit test for sagemaker-endpoint * update tests for sagemaker-studio * fix Lambda version suppression for sagemaker-studio --- .../sagemaker-jumpstart-fm-endpoint/index.ts | 3 + .../sagemaker-jumpstart-fm-endpoint-stack.ts | 8 +++ .../package-lock.json | 13 ++-- .../package.json | 1 + .../test/sagemaker-fm-endpoint.test.ts | 19 ++++-- modules/mlflow/mlflow-fargate/app.py | 3 + modules/mlflow/mlflow-fargate/stack.py | 32 ++++++---- .../mlflow/mlflow-fargate/tests/test_stack.py | 23 +++++-- modules/mlflow/mlflow-image/app.py | 3 + modules/mlflow/mlflow-image/stack.py | 8 +-- .../mlflow/mlflow-image/tests/test_stack.py | 20 ++++++- modules/sagemaker/sagemaker-endpoint/app.py | 4 +- modules/sagemaker/sagemaker-endpoint/stack.py | 60 +++++++++---------- .../sagemaker-endpoint/tests/test_stack.py | 35 ++++++++--- modules/sagemaker/sagemaker-notebook/app.py | 3 + .../sagemaker-notebook/requirements.txt | 2 +- .../sagemaker_notebook/stack.py | 31 ++++++++++ .../sagemaker-notebook/tests/test_stack.py | 21 ++++++- modules/sagemaker/sagemaker-studio/app.py | 3 + modules/sagemaker/sagemaker-studio/stack.py | 42 ++++++++++++- .../sagemaker-studio/tests/test_stack.py | 28 ++++++--- .../app.py | 3 + .../stack.py | 27 +++++++++ .../tests/test_stack.py | 20 ++++++- scripts/fix.sh | 2 +- 25 files changed, 321 insertions(+), 93 deletions(-) diff --git a/modules/fmops/sagemaker-jumpstart-fm-endpoint/index.ts b/modules/fmops/sagemaker-jumpstart-fm-endpoint/index.ts index 9aeed0fa..99265b94 100644 --- a/modules/fmops/sagemaker-jumpstart-fm-endpoint/index.ts +++ b/modules/fmops/sagemaker-jumpstart-fm-endpoint/index.ts @@ -1,5 +1,6 @@ import "source-map-support/register"; import * as cdk from "aws-cdk-lib"; +import * as cdk_nag from "cdk-nag"; import { SagemakerJumpStartFmEndpointStack } from "./lib/sagemaker-jumpstart-fm-endpoint-stack"; const account = process.env.CDK_DEFAULT_ACCOUNT; @@ -35,4 +36,6 @@ new cdk.CfnOutput(stack, "metadata", { }), }); +cdk.Aspects.of(app).add(new cdk_nag.AwsSolutionsChecks({ logIgnores: true })); + app.synth(); diff --git a/modules/fmops/sagemaker-jumpstart-fm-endpoint/lib/sagemaker-jumpstart-fm-endpoint-stack.ts b/modules/fmops/sagemaker-jumpstart-fm-endpoint/lib/sagemaker-jumpstart-fm-endpoint-stack.ts index 7428b13c..baf5f2c6 100644 --- a/modules/fmops/sagemaker-jumpstart-fm-endpoint/lib/sagemaker-jumpstart-fm-endpoint-stack.ts +++ b/modules/fmops/sagemaker-jumpstart-fm-endpoint/lib/sagemaker-jumpstart-fm-endpoint-stack.ts @@ -8,6 +8,7 @@ import { SageMakerInstanceType, JumpStartSageMakerEndpoint, } from "@cdklabs/generative-ai-cdk-constructs"; +import * as cdk_nag from "cdk-nag"; interface SagemakerJumpStartFmEndpointStackProps extends cdk.StackProps { projectName?: string; @@ -91,5 +92,12 @@ export class SagemakerJumpStartFmEndpointStack extends cdk.Stack { role: this.role, vpcConfig: vpcConfig, }); + + cdk_nag.NagSuppressions.addResourceSuppressions(this.role, [ + { + id: "AwsSolutions-IAM5", + reason: "Resource access restriced to S3 buckets (with a prefix) and ECR images", + }, + ]); } } diff --git a/modules/fmops/sagemaker-jumpstart-fm-endpoint/package-lock.json b/modules/fmops/sagemaker-jumpstart-fm-endpoint/package-lock.json index c4fdae8e..c64a2eb9 100644 --- a/modules/fmops/sagemaker-jumpstart-fm-endpoint/package-lock.json +++ b/modules/fmops/sagemaker-jumpstart-fm-endpoint/package-lock.json @@ -22,6 +22,7 @@ "@typescript-eslint/eslint-plugin": "^6.21.0", "@typescript-eslint/parser": "^6.21.0", "aws-cdk": "2.130.0", + "cdk-nag": "^2.28.55", "cypress": "^13.6.1", "jest": "^29.7.0", "prettier": "^3.1.1", @@ -2682,9 +2683,9 @@ "dev": true }, "node_modules/cdk-nag": { - "version": "2.28.47", - "resolved": "https://registry.npmjs.org/cdk-nag/-/cdk-nag-2.28.47.tgz", - "integrity": "sha512-QWDAehKW3KKh66dKOYlUdLno+HWdR+KrCw8/8gV/uS2srpKCKcPaK427RXnp6QNO/eV9g7HGIr+GL7fEJCv4RQ==", + "version": "2.28.55", + "resolved": "https://registry.npmjs.org/cdk-nag/-/cdk-nag-2.28.55.tgz", + "integrity": "sha512-ETdEB6zFQqxVrWXMZSI3c3EoMNOp919pdYsb11zlZyyfS99mgKf9wdXHpBWYu2gY+efxXktWj7HLoPU6g1sxrQ==", "peerDependencies": { "aws-cdk-lib": "^2.116.0", "constructs": "^10.0.5" @@ -8881,9 +8882,9 @@ "dev": true }, "cdk-nag": { - "version": "2.28.47", - "resolved": "https://registry.npmjs.org/cdk-nag/-/cdk-nag-2.28.47.tgz", - "integrity": "sha512-QWDAehKW3KKh66dKOYlUdLno+HWdR+KrCw8/8gV/uS2srpKCKcPaK427RXnp6QNO/eV9g7HGIr+GL7fEJCv4RQ==", + "version": "2.28.55", + "resolved": "https://registry.npmjs.org/cdk-nag/-/cdk-nag-2.28.55.tgz", + "integrity": "sha512-ETdEB6zFQqxVrWXMZSI3c3EoMNOp919pdYsb11zlZyyfS99mgKf9wdXHpBWYu2gY+efxXktWj7HLoPU6g1sxrQ==", "requires": {} }, "chalk": { diff --git a/modules/fmops/sagemaker-jumpstart-fm-endpoint/package.json b/modules/fmops/sagemaker-jumpstart-fm-endpoint/package.json index 59901bd0..5e73c630 100644 --- a/modules/fmops/sagemaker-jumpstart-fm-endpoint/package.json +++ b/modules/fmops/sagemaker-jumpstart-fm-endpoint/package.json @@ -18,6 +18,7 @@ "@types/jest": "^29.5.5", "@types/node": "20.7.1", "aws-cdk": "2.130.0", + "cdk-nag": "^2.28.55", "cypress": "^13.6.1", "jest": "^29.7.0", "prettier": "^3.1.1", diff --git a/modules/fmops/sagemaker-jumpstart-fm-endpoint/test/sagemaker-fm-endpoint.test.ts b/modules/fmops/sagemaker-jumpstart-fm-endpoint/test/sagemaker-fm-endpoint.test.ts index 178155a7..c6c394c4 100644 --- a/modules/fmops/sagemaker-jumpstart-fm-endpoint/test/sagemaker-fm-endpoint.test.ts +++ b/modules/fmops/sagemaker-jumpstart-fm-endpoint/test/sagemaker-fm-endpoint.test.ts @@ -1,8 +1,8 @@ import * as cdk from "aws-cdk-lib"; -import { Template } from "aws-cdk-lib/assertions"; +import { Annotations, Match, Template } from "aws-cdk-lib/assertions"; import { SagemakerJumpStartFmEndpointStack } from "../lib/sagemaker-jumpstart-fm-endpoint-stack"; -test("Synth stack", () => { +describe("Sagemaker JumpStart Fm Endpoint Stack", () => { const app = new cdk.App(); const projectName = "mlops"; @@ -26,9 +26,16 @@ test("Synth stack", () => { env: { account, region }, }); - const template = Template.fromStack(stack); + test("Synth stack", () => { + const template = Template.fromStack(stack); - template.hasResource("AWS::SageMaker::Endpoint", {}); - template.hasResource("AWS::IAM::Role", {}); - template.hasResource("AWS::EC2::SecurityGroup", {}); + template.hasResource("AWS::SageMaker::Endpoint", {}); + template.hasResource("AWS::IAM::Role", {}); + template.hasResource("AWS::EC2::SecurityGroup", {}); + }); + + test("No CDK Nag Errors", () => { + const errors = Annotations.fromStack(stack).findError("*", Match.stringLikeRegexp("AwsSolutions-.*")); + expect(errors).toHaveLength(0); + }); }); diff --git a/modules/mlflow/mlflow-fargate/app.py b/modules/mlflow/mlflow-fargate/app.py index a9daa993..0b0727c9 100644 --- a/modules/mlflow/mlflow-fargate/app.py +++ b/modules/mlflow/mlflow-fargate/app.py @@ -6,6 +6,7 @@ from typing import Optional, cast import aws_cdk +import cdk_nag from stack import MlflowFargateStack, RDSSettings @@ -116,4 +117,6 @@ def _param(name: str) -> str: ), ) +aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) + app.synth() diff --git a/modules/mlflow/mlflow-fargate/stack.py b/modules/mlflow/mlflow-fargate/stack.py index 576aacdb..d694e752 100644 --- a/modules/mlflow/mlflow-fargate/stack.py +++ b/modules/mlflow/mlflow-fargate/stack.py @@ -47,7 +47,7 @@ def __init__( cdk.Tags.of(scope=cast(IConstruct, self)).add(key="Deployment", value=app_prefix[:64]) - role = iam.Role( + task_role = iam.Role( self, "TaskRole", assumed_by=iam.ServicePrincipal(service="ecs-tasks.amazonaws.com"), @@ -58,7 +58,7 @@ def __init__( # Grant artifacts bucket read-write permissions model_bucket = s3.Bucket.from_bucket_name(self, "ArtifactsBucket", bucket_name=artifacts_bucket_name) - model_bucket.grant_read_write(role) + model_bucket.grant_read_write(task_role) vpc = ec2.Vpc.from_lookup(self, "Vpc", vpc_id=vpc_id) subnets = [ec2.Subnet.from_subnet_id(self, f"Subnet {subnet_id}", subnet_id) for subnet_id in subnet_ids] @@ -76,7 +76,7 @@ def __init__( task_definition = ecs.FargateTaskDefinition( self, "MlflowTask", - task_role=role, + task_role=task_role, cpu=task_cpu_units, memory_limit_mib=task_memory_limit_mb, ) @@ -172,8 +172,8 @@ def __init__( fs.connections.allow_default_port_from(fargate_service.service.connections) # Setup security group - fargate_securigy_group = fargate_service.service.connections.security_groups[0] - fargate_securigy_group.add_ingress_rule( + fargate_security_group = fargate_service.service.connections.security_groups[0] + fargate_security_group.add_ingress_rule( peer=ec2.Peer.ipv4(vpc.vpc_cidr_block), connection=ec2.Port.tcp(5000), description="Allow inbound from VPC for mlflow", @@ -185,7 +185,7 @@ def __init__( self, "RDS Security Group", rds_settings["security_group_id"] ) rds_security_group.add_ingress_rule( - peer=ec2.Peer.security_group_id(fargate_securigy_group.security_group_id), + peer=ec2.Peer.security_group_id(fargate_security_group.security_group_id), connection=ec2.Port.tcp(rds_settings["port"]), description="Allow inbound access to RDS for mlflow", ) @@ -199,12 +199,10 @@ def __init__( scale_out_cooldown=cdk.Duration.seconds(60), ) - # Add CDK nag solutions checks - cdk.Aspects.of(self).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) - - cdk_nag.NagSuppressions.add_stack_suppressions( - self, - apply_to_nested_stacks=True, + # Add CDK nag suppressions + cdk_nag.NagSuppressions.add_resource_suppressions( + [task_role, task_definition.execution_role], # type: ignore[list-item] + apply_to_children=True, suppressions=[ cdk_nag.NagPackSuppression( id="AwsSolutions-IAM4", @@ -214,10 +212,20 @@ def __init__( id="AwsSolutions-IAM5", reason="Resource access restricted to resources", ), + ], + ) + cdk_nag.NagSuppressions.add_resource_suppressions( + task_definition, + suppressions=[ cdk_nag.NagPackSuppression( id="AwsSolutions-ECS2", reason="Not passing secrets via env variables", ), + ], + ) + cdk_nag.NagSuppressions.add_resource_suppressions( + lb_access_logs_bucket, + suppressions=[ cdk_nag.NagPackSuppression( id="AwsSolutions-S1", reason="Access logs not required for access logs bucket", diff --git a/modules/mlflow/mlflow-fargate/tests/test_stack.py b/modules/mlflow/mlflow-fargate/tests/test_stack.py index 032977b2..9f296590 100644 --- a/modules/mlflow/mlflow-fargate/tests/test_stack.py +++ b/modules/mlflow/mlflow-fargate/tests/test_stack.py @@ -3,8 +3,9 @@ from unittest import mock import aws_cdk as cdk +import cdk_nag import pytest -from aws_cdk.assertions import Template +from aws_cdk.assertions import Annotations, Match, Template @pytest.fixture(scope="function") @@ -20,8 +21,8 @@ def stack_defaults(): yield -@pytest.mark.parametrize("use_rds", [False, True]) -def test_synthesize_stack(stack_defaults, use_rds) -> None: +@pytest.fixture(scope="function") +def stack(stack_defaults, use_rds: bool) -> cdk.Stack: import stack app = cdk.App() @@ -49,7 +50,7 @@ def test_synthesize_stack(stack_defaults, use_rds) -> None: else: rds_settings = None - stack = stack.MlflowFargateStack( + return stack.MlflowFargateStack( scope=app, id=app_prefix, app_prefix=app_prefix, @@ -71,6 +72,20 @@ def test_synthesize_stack(stack_defaults, use_rds) -> None: ), ) + +@pytest.mark.parametrize("use_rds", [False, True]) +def test_synthesize_stack(stack: cdk.Stack, use_rds: bool) -> None: template = Template.from_stack(stack) template.resource_count_is("AWS::ECS::Cluster", 1) template.resource_count_is("AWS::ECS::TaskDefinition", 1) + + +@pytest.mark.parametrize("use_rds", [False, True]) +def test_no_cdk_nag_errors(stack: cdk.Stack, use_rds: bool) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/modules/mlflow/mlflow-image/app.py b/modules/mlflow/mlflow-image/app.py index 6d8cddfb..22e5c94a 100644 --- a/modules/mlflow/mlflow-image/app.py +++ b/modules/mlflow/mlflow-image/app.py @@ -4,6 +4,7 @@ import os import aws_cdk +import cdk_nag from stack import MlflowImagePublishingStack @@ -51,4 +52,6 @@ def _param(name: str) -> str: ), ) +aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) + app.synth() diff --git a/modules/mlflow/mlflow-image/stack.py b/modules/mlflow/mlflow-image/stack.py index d9b68001..81828272 100644 --- a/modules/mlflow/mlflow-image/stack.py +++ b/modules/mlflow/mlflow-image/stack.py @@ -41,16 +41,16 @@ def __init__( dest=ecr_deployment.DockerImageName(self.image_uri), ) - # Add CDK nag solutions checks - cdk.Aspects.of(self).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) - + # Add CDK nag suppressions cdk_nag.NagSuppressions.add_stack_suppressions( self, - apply_to_nested_stacks=True, suppressions=[ cdk_nag.NagPackSuppression( id="AwsSolutions-IAM4", reason="Managed Policies are for src account roles only", + applies_to=[ + "Policy::arn::iam::aws:policy/service-role/AWSLambdaBasicExecutionRole", + ], ), cdk_nag.NagPackSuppression( id="AwsSolutions-IAM5", diff --git a/modules/mlflow/mlflow-image/tests/test_stack.py b/modules/mlflow/mlflow-image/tests/test_stack.py index d60ac344..33445738 100644 --- a/modules/mlflow/mlflow-image/tests/test_stack.py +++ b/modules/mlflow/mlflow-image/tests/test_stack.py @@ -2,8 +2,9 @@ import sys import aws_cdk as cdk +import cdk_nag import pytest -from aws_cdk.assertions import Template +from aws_cdk.assertions import Annotations, Match, Template @pytest.fixture(scope="function") @@ -16,7 +17,8 @@ def stack_defaults() -> None: del sys.modules["stack"] -def test_synthesize_stack() -> None: +@pytest.fixture(scope="function") +def stack(stack_defaults) -> cdk.Stack: import stack app = cdk.App() @@ -28,7 +30,7 @@ def test_synthesize_stack() -> None: ecr_repo_name = "repo" - stack = stack.MlflowImagePublishingStack( + return stack.MlflowImagePublishingStack( scope=app, id=app_prefix, app_prefix=app_prefix, @@ -39,5 +41,17 @@ def test_synthesize_stack() -> None: ), ) + +def test_synthesize_stack(stack: cdk.Stack) -> None: template = Template.from_stack(stack) template.resource_count_is("Custom::CDKBucketDeployment", 1) + + +def test_no_cdk_nag_errors(stack: cdk.Stack) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/modules/sagemaker/sagemaker-endpoint/app.py b/modules/sagemaker/sagemaker-endpoint/app.py index 14a28fb1..a23e50c8 100644 --- a/modules/sagemaker/sagemaker-endpoint/app.py +++ b/modules/sagemaker/sagemaker-endpoint/app.py @@ -5,6 +5,7 @@ import os import aws_cdk +import cdk_nag from stack import DeployEndpointStack @@ -84,7 +85,7 @@ def _param(name: str) -> str: id="metadata", value=stack.to_json_string( { - "ModelExecutionRoleArn": stack.model_execution_role_arn, + "ModelExecutionRoleArn": stack.model_execution_role.role_arn, "ModelName": stack.model.model_name, "ModelPackageArn": stack.model_package_arn, "EndpointName": stack.endpoint.attr_endpoint_name, @@ -93,5 +94,6 @@ def _param(name: str) -> str: ), ) +aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) app.synth() diff --git a/modules/sagemaker/sagemaker-endpoint/stack.py b/modules/sagemaker/sagemaker-endpoint/stack.py index a8fe2f62..877b73e8 100644 --- a/modules/sagemaker/sagemaker-endpoint/stack.py +++ b/modules/sagemaker/sagemaker-endpoint/stack.py @@ -5,13 +5,13 @@ from typing import Any, Dict, List, Optional import constructs -from aws_cdk import Aspects, Stack, Tags +from aws_cdk import Stack, Tags from aws_cdk import aws_ec2 as ec2 from aws_cdk import aws_iam as iam from aws_cdk import aws_kms as kms from aws_cdk import aws_s3 as s3 from aws_cdk import aws_sagemaker as sagemaker -from cdk_nag import AwsSolutionsChecks, NagPackSuppression, NagSuppressions +from cdk_nag import NagPackSuppression, NagSuppressions from scripts.get_approved_package import get_approved_package @@ -54,9 +54,10 @@ def __init__( connection=ec2.Port.all_tcp(), ) + self.model_execution_role: iam.IRole if not model_execution_role_arn: # Create model execution role - model_execution_role = iam.Role( + self.model_execution_role = iam.Role( self, f"{app_prefix}-model-exec", assumed_by=iam.ServicePrincipal("sagemaker.amazonaws.com"), @@ -68,21 +69,21 @@ def __init__( if model_artifacts_bucket_arn: # Grant model assets bucket read permissions model_bucket = s3.Bucket.from_bucket_arn(self, f"{app_prefix}-model-bucket", model_artifacts_bucket_arn) - model_bucket.grant_read(model_execution_role) + model_bucket.grant_read(self.model_execution_role) if ecr_repo_arn: # Add ECR permissions - model_execution_role.add_to_policy( + self.model_execution_role.add_to_policy( iam.PolicyStatement( actions=["ecr:Get*"], effect=iam.Effect.ALLOW, resources=[ecr_repo_arn], ) ) - - model_execution_role_arn = model_execution_role.role_arn - - self.model_execution_role_arn = model_execution_role_arn + else: + self.model_execution_role = iam.Role.from_role_arn( + self, f"{app_prefix}-model-exec", model_execution_role_arn + ) if not model_package_arn: # Get latest approved model package from the model registry @@ -90,6 +91,7 @@ def __init__( model_package_arn = get_approved_package(self.region, model_package_group_name) else: raise ValueError("Either model_package_arn or model_package_group_name is required") + self.model_package_arn = model_package_arn # Create model instance @@ -97,7 +99,7 @@ def __init__( model = sagemaker.CfnModel( self, f"{app_prefix}-model", - execution_role_arn=model_execution_role_arn, + execution_role_arn=self.model_execution_role.role_arn, model_name=model_name, containers=[sagemaker.CfnModel.ContainerDefinitionProperty(model_package_name=model_package_arn)], vpc_config=sagemaker.CfnModel.VpcConfigProperty( @@ -144,25 +146,19 @@ def __init__( f"https://runtime.sagemaker.{self.region}.amazonaws.com/endpoints/{endpoint.attr_endpoint_name}/invocations" ) - # Add CDK nag solutions checks - Aspects.of(self).add(AwsSolutionsChecks()) - - NagSuppressions.add_stack_suppressions( - self, - suppressions=[ - NagPackSuppression( - id="AwsSolutions-IAM4", - reason="Managed Policies are for service account roles only.", - ) - ], - ) - - NagSuppressions.add_stack_suppressions( - self, - suppressions=[ - NagPackSuppression( - id="AwsSolutions-IAM5", - reason="Model execution role requires s3 permissions to the bucket.", - ) - ], - ) + # Add CDK nag suppressions + if not model_execution_role_arn: + NagSuppressions.add_resource_suppressions( + self.model_execution_role, + apply_to_children=True, + suppressions=[ + NagPackSuppression( + id="AwsSolutions-IAM4", + reason="Managed Policies are for service account roles only.", + ), + NagPackSuppression( + id="AwsSolutions-IAM5", + reason="Model execution role requires s3 permissions to the bucket.", + ), + ], + ) diff --git a/modules/sagemaker/sagemaker-endpoint/tests/test_stack.py b/modules/sagemaker/sagemaker-endpoint/tests/test_stack.py index 6cfc223d..868b8bb6 100644 --- a/modules/sagemaker/sagemaker-endpoint/tests/test_stack.py +++ b/modules/sagemaker/sagemaker-endpoint/tests/test_stack.py @@ -4,8 +4,9 @@ import aws_cdk as cdk import botocore.session +import cdk_nag import pytest -from aws_cdk.assertions import Template +from aws_cdk.assertions import Annotations, Match, Template from botocore.stub import Stubber @@ -19,7 +20,8 @@ def stack_defaults() -> None: del sys.modules["stack"] -def test_synthesize_stack_model_package_input() -> None: +@pytest.fixture(scope="function") +def stack_model_package_input() -> cdk.Stack: import stack app = cdk.App() @@ -35,7 +37,7 @@ def test_synthesize_stack_model_package_input() -> None: model_package_arn = "example-arn" model_artifacts_bucket_arn = "arn:aws:s3:::test-bucket" - endpoint_stack = stack.DeployEndpointStack( + return stack.DeployEndpointStack( scope=app, id=app_prefix, app_prefix=app_prefix, @@ -58,12 +60,10 @@ def test_synthesize_stack_model_package_input() -> None: }, ) - template = Template.from_stack(endpoint_stack) - template.resource_count_is("AWS::SageMaker::Endpoint", 1) - +@pytest.fixture(scope="function") @mock.patch("scripts.get_approved_package.boto3.client") -def test_synthesize_stack_latest_approved_model_package(mock_s3_client, stack_defaults) -> None: +def stack_latest_approved_model_package(mock_s3_client) -> cdk.Stack: import stack app = cdk.App() @@ -101,7 +101,7 @@ def test_synthesize_stack_latest_approved_model_package(mock_s3_client, stack_de } stubber.add_response("list_model_packages", response, expected_params) - endpoint_stack = stack.DeployEndpointStack( + return stack.DeployEndpointStack( scope=app, id=app_prefix, app_prefix=app_prefix, @@ -124,5 +124,22 @@ def test_synthesize_stack_latest_approved_model_package(mock_s3_client, stack_de }, ) - template = Template.from_stack(endpoint_stack) + +@pytest.fixture(params=["stack_model_package_input", "stack_latest_approved_model_package"], scope="function") +def stack(request, stack_model_package_input, stack_latest_approved_model_package) -> cdk.Stack: + return request.getfixturevalue(request.param) + + +def test_synthesize_stack(stack: cdk.Stack) -> None: + template = Template.from_stack(stack) template.resource_count_is("AWS::SageMaker::Endpoint", 1) + + +def test_no_cdk_nag_errors(stack: cdk.Stack) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/modules/sagemaker/sagemaker-notebook/app.py b/modules/sagemaker/sagemaker-notebook/app.py index 8a64f869..9251320d 100644 --- a/modules/sagemaker/sagemaker-notebook/app.py +++ b/modules/sagemaker/sagemaker-notebook/app.py @@ -2,6 +2,7 @@ """Create a Sagemaker Model Stack.""" import aws_cdk as cdk +import cdk_nag from sagemaker_notebook.settings import ApplicationSettings from sagemaker_notebook.stack import SagemakerNotebookStack @@ -23,4 +24,6 @@ **app_settings.parameters.model_dump(), ) +cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) + app.synth() diff --git a/modules/sagemaker/sagemaker-notebook/requirements.txt b/modules/sagemaker/sagemaker-notebook/requirements.txt index 1a02af55..8bf170b8 100644 --- a/modules/sagemaker/sagemaker-notebook/requirements.txt +++ b/modules/sagemaker/sagemaker-notebook/requirements.txt @@ -3,4 +3,4 @@ cdk-nag~=2.28.27 constructs~=10.3.0 pydantic~=2.5.3 pydantic-settings~=2.0.3 -configobj~=5.0.8 \ No newline at end of file +configobj~=5.0.8 diff --git a/modules/sagemaker/sagemaker-notebook/sagemaker_notebook/stack.py b/modules/sagemaker/sagemaker-notebook/sagemaker_notebook/stack.py index f5f24c7e..93b7c3b8 100644 --- a/modules/sagemaker/sagemaker-notebook/sagemaker_notebook/stack.py +++ b/modules/sagemaker/sagemaker-notebook/sagemaker_notebook/stack.py @@ -6,6 +6,7 @@ import aws_cdk.aws_iam as iam import aws_cdk.aws_kms as kms import aws_cdk.aws_sagemaker as sagemaker +import cdk_nag from aws_cdk import CfnOutput, Stack, Tags from constructs import Construct @@ -95,6 +96,36 @@ def __init__( self.setup_tags() + cdk_nag.NagSuppressions.add_resource_suppressions( + self.role, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM4", + reason="Managed Policies are for src account roles only", + ), + ], + ) + if not kms_key_arn: + cdk_nag.NagSuppressions.add_resource_suppressions( + self.sm_notebook, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-SM2", + reason="Customers have the option to disable encryption", + ), + ], + ) + if direct_internet_access: + cdk_nag.NagSuppressions.add_resource_suppressions( + self.sm_notebook, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-SM3", + reason="Customers have the option to enable direct internet access", + ), + ], + ) + def setup_resources(self) -> None: """Deploy resources.""" self.vpc = self.get_vpc() diff --git a/modules/sagemaker/sagemaker-notebook/tests/test_stack.py b/modules/sagemaker/sagemaker-notebook/tests/test_stack.py index a82a862f..34129675 100644 --- a/modules/sagemaker/sagemaker-notebook/tests/test_stack.py +++ b/modules/sagemaker/sagemaker-notebook/tests/test_stack.py @@ -1,8 +1,11 @@ import aws_cdk as cdk -from aws_cdk.assertions import Template +import cdk_nag +import pytest +from aws_cdk.assertions import Annotations, Match, Template -def test_synthesize_stack() -> None: +@pytest.fixture(scope="function") +def stack() -> cdk.Stack: from sagemaker_notebook import stack app = cdk.App() @@ -12,7 +15,7 @@ def test_synthesize_stack() -> None: mod_name = "test-module" app_prefix = f"{project_name}-{dep_name}-{mod_name}" - stack = stack.SagemakerNotebookStack( + return stack.SagemakerNotebookStack( scope=app, construct_id=app_prefix, env=cdk.Environment(account="111111111111", region="us-east-1"), @@ -31,5 +34,17 @@ def test_synthesize_stack() -> None: tags={"test": "True"}, ) + +def test_synthesize_stack(stack) -> None: template = Template.from_stack(stack) template.resource_count_is("AWS::SageMaker::NotebookInstance", 1) + + +def test_no_cdk_nag_errors(stack) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/modules/sagemaker/sagemaker-studio/app.py b/modules/sagemaker/sagemaker-studio/app.py index c1244008..4ca927ed 100644 --- a/modules/sagemaker/sagemaker-studio/app.py +++ b/modules/sagemaker/sagemaker-studio/app.py @@ -6,6 +6,7 @@ from typing import cast import aws_cdk +import cdk_nag from aws_cdk import CfnOutput from stack import SagemakerStudioStack @@ -80,4 +81,6 @@ def _param(name: str) -> str: ), ) +aws_cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) + app.synth() diff --git a/modules/sagemaker/sagemaker-studio/stack.py b/modules/sagemaker/sagemaker-studio/stack.py index f0be1e9e..10a140dc 100644 --- a/modules/sagemaker/sagemaker-studio/stack.py +++ b/modules/sagemaker/sagemaker-studio/stack.py @@ -4,6 +4,7 @@ from typing import Any, List, cast import aws_cdk as core +import cdk_nag from aws_cdk import Stack, Tags from aws_cdk import aws_ec2 as ec2 from aws_cdk import aws_iam as iam @@ -111,11 +112,26 @@ def __init__( for user in lead_data_science_users ] + cdk_nag.NagSuppressions.add_resource_suppressions( + self.sm_roles, + apply_to_children=True, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM4", + reason="Managed Policies are for src account roles only", + ), + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM5", + reason="Resource access restricted to resources", + ), + ], + ) + def enable_sagemaker_projects(self, roles: List[str]) -> None: event_handler = PythonFunction( self, "sg-project-function", - runtime=lambda_.Runtime.PYTHON_3_8, + runtime=lambda_.Runtime.PYTHON_3_12, entry="functions/sm_studio/enable_sm_projects", timeout=core.Duration.seconds(120), ) @@ -152,6 +168,30 @@ def enable_sagemaker_projects(self, roles: List[str]) -> None: }, ) + cdk_nag.NagSuppressions.add_resource_suppressions( + [event_handler.role, provider], # type: ignore[list-item] + apply_to_children=True, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM4", + reason="Managed Policies are for src account roles only", + ), + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM5", + reason="Resource access restricted to resources", + ), + ], + ) + cdk_nag.NagSuppressions.add_resource_suppressions( + provider, + apply_to_children=True, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-L1", reason="We don't control the version used by the Provider construct." + ), + ], + ) + def sagemaker_studio_domain( self, domain_name: str, diff --git a/modules/sagemaker/sagemaker-studio/tests/test_stack.py b/modules/sagemaker/sagemaker-studio/tests/test_stack.py index fdbb186f..c739f061 100644 --- a/modules/sagemaker/sagemaker-studio/tests/test_stack.py +++ b/modules/sagemaker/sagemaker-studio/tests/test_stack.py @@ -5,22 +5,23 @@ import sys import aws_cdk as cdk +import cdk_nag import pytest -from aws_cdk.assertions import Template +from aws_cdk.assertions import Annotations, Match, Template @pytest.fixture(scope="function") -def stack_defaults(): +def stack_defaults() -> None: os.environ["CDK_DEFAULT_ACCOUNT"] = "111111111111" os.environ["CDK_DEFAULT_REGION"] = "us-east-1" # Unload the app import so that subsequent tests don't reuse - if "stack" in sys.modules: del sys.modules["stack"] -def test_synthesize_stack(stack_defaults): +@pytest.fixture(scope="function") +def stack(stack_defaults, enable_custom_sagemaker_projects: bool) -> cdk.Stack: import stack app = cdk.App() @@ -33,9 +34,8 @@ def test_synthesize_stack(stack_defaults): lead_data_science_users = ["lead-ds-user-1"] app_image_config_name = None image_name = None - enable_custom_sagemaker_projects = False - stack = stack.SagemakerStudioStack( + return stack.SagemakerStudioStack( app, f"{project_name}-{dep_name}-{mod_name}", project_name=project_name, @@ -56,9 +56,23 @@ def test_synthesize_stack(stack_defaults): enable_custom_sagemaker_projects=enable_custom_sagemaker_projects, ) + +@pytest.mark.parametrize("enable_custom_sagemaker_projects", [True, False]) +def test_synthesize_stack(stack: cdk.Stack, enable_custom_sagemaker_projects: bool) -> None: template = Template.from_stack(stack) template.resource_count_is("AWS::SageMaker::Domain", 1) template.resource_count_is("AWS::SageMaker::UserProfile", 2) template.resource_count_is("AWS::EC2::SecurityGroup", 1) - template.resource_count_is("AWS::IAM::Role", 3) + template.resource_count_is("AWS::IAM::Role", 5 if enable_custom_sagemaker_projects else 3) + + +@pytest.mark.parametrize("enable_custom_sagemaker_projects", [True, False]) +def test_no_cdk_nag_errors(stack: cdk.Stack, enable_custom_sagemaker_projects: bool) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/modules/sagemaker/sagemaker-templates-service-catalog/app.py b/modules/sagemaker/sagemaker-templates-service-catalog/app.py index 22261111..3df783cc 100644 --- a/modules/sagemaker/sagemaker-templates-service-catalog/app.py +++ b/modules/sagemaker/sagemaker-templates-service-catalog/app.py @@ -4,6 +4,7 @@ import os import aws_cdk +import cdk_nag from stack import ServiceCatalogStack @@ -53,4 +54,6 @@ def _param(name: str) -> str: ), ) +aws_cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks(log_ignores=True)) + app.synth() diff --git a/modules/sagemaker/sagemaker-templates-service-catalog/stack.py b/modules/sagemaker/sagemaker-templates-service-catalog/stack.py index c00375a4..c932c87d 100644 --- a/modules/sagemaker/sagemaker-templates-service-catalog/stack.py +++ b/modules/sagemaker/sagemaker-templates-service-catalog/stack.py @@ -5,6 +5,7 @@ import os from typing import Any, Optional, Tuple +import cdk_nag from aws_cdk import BundlingOptions, BundlingOutput, DockerImage, Stack, Tags from aws_cdk import aws_iam as iam from aws_cdk import aws_s3_assets as s3_assets @@ -60,6 +61,9 @@ def __init__( templates_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") for template_name in next(os.walk(templates_dir))[1]: + if template_name == "__pycache__": + continue + build_app_asset, deploy_app_asset = self.upload_assets( portfolio_access_role=portfolio_access_role, template_name=template_name, @@ -96,6 +100,29 @@ def __init__( Tags.of(product).add(key="sagemaker:studio-visibility", value="true") + cdk_nag.NagSuppressions.add_resource_suppressions( + portfolio_access_role, + apply_to_children=True, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM5", + reason="The role needs wildcard permissions to be able to access template assets in S3", + ), + ], + ) + cdk_nag.NagSuppressions.add_resource_suppressions( + product_launch_role, + suppressions=[ + cdk_nag.NagPackSuppression( + id="AwsSolutions-IAM4", + reason=( + "Product launch role needs admin permissions in order to be able " + "to create resources in the AWS account." + ), + ), + ], + ) + def upload_assets( self, portfolio_access_role: iam.IRole, diff --git a/modules/sagemaker/sagemaker-templates-service-catalog/tests/test_stack.py b/modules/sagemaker/sagemaker-templates-service-catalog/tests/test_stack.py index aebeb860..04d6d7f4 100644 --- a/modules/sagemaker/sagemaker-templates-service-catalog/tests/test_stack.py +++ b/modules/sagemaker/sagemaker-templates-service-catalog/tests/test_stack.py @@ -5,8 +5,9 @@ import sys import aws_cdk as cdk +import cdk_nag import pytest -from aws_cdk.assertions import Template +from aws_cdk.assertions import Annotations, Match, Template @pytest.fixture(scope="function") @@ -20,7 +21,8 @@ def stack_defaults(): del sys.modules["stack"] -def test_synthesize_stack(stack_defaults): +@pytest.fixture(scope="function") +def stack(stack_defaults) -> cdk.Stack: import stack app = cdk.App() @@ -31,7 +33,7 @@ def test_synthesize_stack(stack_defaults): portfolio_owner = "owner" portfolio_access_role_arn = "arn:aws:iam::xxxxxxxxxxxx:role/role" - stack = stack.ServiceCatalogStack( + return stack.ServiceCatalogStack( app, f"{project_name}-{dep_name}-{mod_name}", portfolio_name=portfolio_name, @@ -43,7 +45,19 @@ def test_synthesize_stack(stack_defaults): ), ) + +def test_synthesize_stack(stack: cdk.Stack) -> None: template = Template.from_stack(stack) template.resource_count_is("AWS::ServiceCatalog::Portfolio", 1) template.resource_count_is("AWS::ServiceCatalog::CloudFormationProduct", 1) + + +def test_no_cdk_nag_errors(stack: cdk.Stack) -> None: + cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks()) + + nag_errors = Annotations.from_stack(stack).find_error( + "*", + Match.string_like_regexp(r"AwsSolutions-.*"), + ) + assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors" diff --git a/scripts/fix.sh b/scripts/fix.sh index 64b666be..6d190a03 100755 --- a/scripts/fix.sh +++ b/scripts/fix.sh @@ -46,7 +46,7 @@ FIX_PATH=`pwd` echo "Fixing: ${FIX_PATH}, Language: ${LANGUAGE}" if [[ $LANGUAGE == "python" ]]; then - echo "Running isort, black" + echo "Running ruff" ruff format . ruff check --fix . elif [[ $LANGUAGE == "typescript" ]]; then