diff --git a/cli/src/pcluster/config/imagebuilder_config.py b/cli/src/pcluster/config/imagebuilder_config.py index eaa8841264..c4918f1883 100644 --- a/cli/src/pcluster/config/imagebuilder_config.py +++ b/cli/src/pcluster/config/imagebuilder_config.py @@ -29,7 +29,7 @@ from pcluster.validators.common import ValidatorContext from pcluster.validators.ebs_validators import EbsVolumeTypeSizeValidator from pcluster.validators.ec2_validators import InstanceTypeBaseAMICompatibleValidator -from pcluster.validators.iam_validators import IamPolicyValidator, InstanceProfileValidator, RoleValidator +from pcluster.validators.iam_validators import IamPolicyValidator, IamResourcePrefixValidator, InstanceProfileValidator, RoleValidator from pcluster.validators.imagebuilder_validators import ( AMIVolumeSizeValidator, ComponentsValidator, @@ -104,6 +104,7 @@ def __init__( cleanup_lambda_role: str = None, additional_iam_policies: List[AdditionalIamPolicy] = (), permissions_boundary: str = None, + resource_prefix: str = None, ): super().__init__() self.instance_role = Resource.init_param(instance_role) @@ -111,6 +112,7 @@ def __init__( self.additional_iam_policies = additional_iam_policies self.instance_profile = Resource.init_param(instance_profile) self.permissions_boundary = Resource.init_param(permissions_boundary) + self.resource_prefix = Resource.init_param(resource_prefix) @property def additional_iam_policy_arns(self) -> List[str]: @@ -130,8 +132,10 @@ def _register_validators(self, context: ValidatorContext = None): # noqa: D102 self._register_validator(RoleValidator, role_arn=self.cleanup_lambda_role) if self.permissions_boundary: - self._register_validator(IamPolicyValidator, policy=self.permissions_boundary) - + self._register_validator(IamPolicyValidator, policy=self.permissions_boundary) + + if self.resource_prefix: + self._register_validator(IamResourcePrefixValidator, resource_prefix=self.resource_prefix) class UpdateOsPackages(Resource): """Represent the UpdateOsPackages configuration for the ImageBuilder.""" diff --git a/cli/src/pcluster/schemas/imagebuilder_schema.py b/cli/src/pcluster/schemas/imagebuilder_schema.py index 1d893a12c1..bf7c0d46c0 100644 --- a/cli/src/pcluster/schemas/imagebuilder_schema.py +++ b/cli/src/pcluster/schemas/imagebuilder_schema.py @@ -144,6 +144,7 @@ class IamSchema(BaseSchema): cleanup_lambda_role = fields.Str(validate=validate.Regexp(IAM_ROLE_REGEX)) additional_iam_policies = fields.Nested(AdditionalIamPolicySchema, many=True) permissions_boundary = fields.Str(validate=validate.Regexp(IAM_POLICY_REGEX)) + resource_prefix = fields.Str() @validates_schema def no_coexist_role_policies(self, data, **kwargs): diff --git a/cli/src/pcluster/templates/imagebuilder_stack.py b/cli/src/pcluster/templates/imagebuilder_stack.py index 49bfce6aef..27df1991b5 100644 --- a/cli/src/pcluster/templates/imagebuilder_stack.py +++ b/cli/src/pcluster/templates/imagebuilder_stack.py @@ -50,7 +50,7 @@ wrap_script_to_component, ) from pcluster.models.s3_bucket import S3Bucket, S3FileType, create_s3_presigned_url, parse_bucket_url -from pcluster.templates.cdk_builder_utils import apply_permissions_boundary, get_assume_role_policy_document +from pcluster.templates.cdk_builder_utils import add_cluster_iam_resource_prefix, apply_permissions_boundary, get_assume_role_policy_document from pcluster.utils import get_http_tokens_setting @@ -64,12 +64,14 @@ def __init__( image_config: ImageBuilderConfig, image_id: str, bucket: S3Bucket, + resource_prefix: str = None, **kwargs, ) -> None: super().__init__(scope, construct_id, **kwargs) self.config = image_config self.image_id = image_id self.bucket = bucket + self.resource_prefix = resource_prefix self.custom_instance_role = ( self.config.build.iam.instance_role @@ -591,33 +593,46 @@ def _add_default_instance_role(self, build_tags): ), ) + _, policy_name = add_cluster_iam_resource_prefix( + self.image_id, self.config, "InstanceRoleInlinePolicy", "Policy" + ) + instancerole_policy = iam.CfnRole.PolicyProperty( - policy_name="InstanceRoleInlinePolicy", + policy_name=policy_name, policy_document=instancerole_policy_document, ) + role_path, role_name = add_cluster_iam_resource_prefix( + self.image_id, self.config, self._build_resource_name(IMAGEBUILDER_RESOURCE_NAME_PREFIX), "Role" + ) + instance_role_resource = iam.CfnRole( self, "InstanceRole", - path=IAM_ROLE_PATH, + path=role_path or IAM_ROLE_PATH, managed_policy_arns=managed_policy_arns, assume_role_policy_document=get_assume_role_policy_document("ec2.{0}".format(self.url_suffix)), policies=[ instancerole_policy, ], tags=build_tags, - role_name=self._build_resource_name(IMAGEBUILDER_RESOURCE_NAME_PREFIX), + role_name=role_name, ) return instance_role_resource def _add_instance_profile(self, instance_role=None): """Set default instance profile in imagebuilder cfn template.""" + + instance_profile_path, instance_profile_name = add_cluster_iam_resource_prefix( + self.image_id, self.config, self._build_resource_name(IMAGEBUILDER_RESOURCE_NAME_PREFIX), "InstanceProfile" + ) + instance_profile_resource = iam.CfnInstanceProfile( self, "InstanceProfile", - path=IAM_ROLE_PATH, + path=instance_profile_path or IAM_ROLE_PATH, roles=[instance_role.split("/")[-1] if instance_role else Fn.ref("InstanceRole")], - instance_profile_name=self._build_resource_name(IMAGEBUILDER_RESOURCE_NAME_PREFIX), + instance_profile_name=instance_profile_name, ) return instance_profile_resource diff --git a/cli/tests/pcluster/schemas/test_imagebuilder_schema.py b/cli/tests/pcluster/schemas/test_imagebuilder_schema.py index 03d8595c6b..d0e5f0e372 100644 --- a/cli/tests/pcluster/schemas/test_imagebuilder_schema.py +++ b/cli/tests/pcluster/schemas/test_imagebuilder_schema.py @@ -153,3 +153,137 @@ def test_imagebuilder_deployment_settings_schema(mocker, config_dict, failure_me else: conf = ImagebuilderDeploymentSettingsSchema().load(config_dict) ImagebuilderDeploymentSettingsSchema().dump(conf) + + +@pytest.mark.parametrize( + "config_dict, describe_image_response, failure_message", + [ + pytest.param( + { + "Region": "us-east-1", + "Image": { + "Name": "test-resource-prefix-ami", + "Tags": [ + {"Key": "Name", "Value": "ParallelCluster-ResourcePrefix-AMI"}, + {"Key": "Version", "Value": "3.15.0"} + ] + }, + "Build": { + "Iam": { + "ResourcePrefix": "/test-path/test-prefix" + }, + "InstanceType": "c5.large", + "ParentImage": "ami-12345678", + "SubnetId": "subnet-0d03dc52", + "Tags": [ + {"Key": "Name", "Value": "ParallelCluster-ResourcePrefix-Build"}, + {"Key": "Test", "Value": "ResourcePrefix"} + ], + "UpdateOsPackages": { + "Enabled": True + } + }, + "CustomS3Bucket": "bucket-name" + }, + { + "Architecture": "x86_64", + "BlockDeviceMappings": [ + { + "DeviceName": "/dev/xvda", + "Ebs": { + "DeleteOnTermination": True, + "SnapshotId": "snap-0a20b6671bc5e3ead", + "VolumeSize": 25, + "VolumeType": "gp2", + "Encrypted": False, + }, + } + ], + }, + None, + id="Test with ResourcePrefix path and name prefix", + ), + pytest.param( + { + "Region": "us-east-1", + "Image": { + "Name": "test-resource-prefix-policies-ami", + "Tags": [ + {"Key": "Name", "Value": "ParallelCluster-ResourcePrefix-Policies-AMI"} + ] + }, + "Build": { + "Iam": { + "ResourcePrefix": "test-prefix-only", + "AdditionalIamPolicies": [ + {"Policy": "arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess"} + ], + "PermissionsBoundary": "arn:aws:iam::111122223333:policy/test-boundary" + }, + "InstanceType": "c5.xlarge", + "ParentImage": "ami-87654321", + "SubnetId": "subnet-0d03dc52", + "SecurityGroupIds": ["sg-0058826a55bae6679"], + "Tags": [ + {"Key": "Name", "Value": "ParallelCluster-ResourcePrefix-Policies-Build"} + ], + "UpdateOsPackages": { + "Enabled": False + } + }, + "CustomS3Bucket": "test-bucket-name" + }, + { + "Architecture": "x86_64", + "BlockDeviceMappings": [ + { + "DeviceName": "/dev/xvda", + "Ebs": { + "DeleteOnTermination": True, + "SnapshotId": "snap-0a20b6671bc5e3ead", + "VolumeSize": 25, + "VolumeType": "gp2", + "Encrypted": False, + }, + } + ], + }, + None, + id="Test with ResourcePrefix name prefix only and additional IAM policies", + ), + ], +) +def test_imagebuilder_schema_with_resource_prefix(mocker, config_dict, describe_image_response, failure_message): + """Test ImageBuilder schema with ResourcePrefix support.""" + mock_aws_api(mocker) + mocker.patch("pcluster.imagebuilder_utils.get_ami_id", return_value="ami-0185634c5a8a37250") + mocker.patch( + "pcluster.aws.ec2.Ec2Client.describe_image", + return_value=describe_image_response, + ) + + print(config_dict) + + if failure_message: + with pytest.raises(ValidationError, match=failure_message): + ImageBuilderSchema().load(config_dict) + else: + imagebuilder_config = ImageBuilderSchema().load(config_dict) + print(imagebuilder_config) + + # Verify ResourcePrefix is properly loaded + if imagebuilder_config.build.iam and imagebuilder_config.build.iam.resource_prefix: + assert_that(imagebuilder_config.build.iam.resource_prefix).is_not_none() + print(f"ResourcePrefix: {imagebuilder_config.build.iam.resource_prefix}") + + # Re-create Yaml file from model and compare content + image_builder_schema = ImageBuilderSchema() + image_builder_schema.context = {"delete_defaults_when_dump": True} + output_json = image_builder_schema.dump(imagebuilder_config) + + # Assert imagebuilder config file can be convert to imagebuilder config + assert_that(json.dumps(config_dict, sort_keys=True)).is_equal_to(json.dumps(output_json, sort_keys=True)) + + # Print output yaml + output_yaml = yaml.dump(output_json) + print(output_yaml) diff --git a/cli/tests/pcluster/templates/test_imagebuilder_stack.py b/cli/tests/pcluster/templates/test_imagebuilder_stack.py index 721d03d23c..d393de93cb 100644 --- a/cli/tests/pcluster/templates/test_imagebuilder_stack.py +++ b/cli/tests/pcluster/templates/test_imagebuilder_stack.py @@ -2705,3 +2705,108 @@ def _lambda_functions_vpc_config_with_camel_cased_keys(resource): ) def test_parse_bucket_url(url, expect_output): assert_that(parse_bucket_url(url)).is_equal_to(expect_output) + + +@pytest.mark.parametrize( + "resource, response, expected_role_path, expected_role_name, expected_profile_path, expected_profile_name", + [ + ( + { + "imagebuilder": { + "build": { + "parent_image": "ami-12345678", + "instance_type": "c5.large", + "iam": {"resource_prefix": "/test-path/test-prefix"}, + }, + } + }, + { + "Architecture": "x86_64", + "BlockDeviceMappings": [ + { + "DeviceName": "/dev/xvda", + "Ebs": { + "DeleteOnTermination": True, + "SnapshotId": "snap-0a20b6671bc5e3ead", + "VolumeSize": 25, + "VolumeType": "gp2", + "Encrypted": False, + }, + } + ], + }, + "/test-path/", # Expected role path + "test-prefix-ParallelClusterImageBuilderRole-fakestackid", # Expected role name with prefix + "/test-path/", # Expected instance profile path + "test-prefix-ParallelClusterImageBuilderInstanceProfile-fakestackid", # Expected instance profile name with prefix + ), + ( + { + "imagebuilder": { + "build": { + "parent_image": "ami-87654321", + "instance_type": "c5.xlarge", + "iam": {"resource_prefix": "name-prefix-only"}, + }, + } + }, + { + "Architecture": "x86_64", + "BlockDeviceMappings": [ + { + "DeviceName": "/dev/xvda", + "Ebs": { + "DeleteOnTermination": True, + "SnapshotId": "snap-0a20b6671bc5e3ead", + "VolumeSize": 25, + "VolumeType": "gp2", + "Encrypted": False, + }, + } + ], + }, + "/", # Default path when only name prefix provided + "name-prefix-only-ParallelClusterImageBuilderRole-fakestackid", # Expected role name with prefix + "/", # Default path when only name prefix provided + "name-prefix-only-ParallelClusterImageBuilderInstanceProfile-fakestackid", # Expected instance profile name with prefix + ), + ], +) +def test_imagebuilder_resource_prefix( + mocker, resource, response, expected_role_path, expected_role_name, expected_profile_path, expected_profile_name +): + """Test ImageBuilder ResourcePrefix support in IAM resources.""" + # Mock image describe response + image_mock = ImageInfo({**response, "ImageId": "ami-12345678"}) + mock_aws_api(mocker) + mocker.patch("pcluster.aws.ec2.Ec2Client.describe_image", return_value=image_mock) + mocker.patch("pcluster.aws.sts.StsClient.get_account_id", return_value="123456789012") + + # Mock stack unique id to make test deterministic + mocker.patch("pcluster.templates.imagebuilder_stack.ImageBuilderCdkStack._stack_unique_id", return_value="fakestackid") + + # Create imagebuilder config and build template + imagebuilder_config = imagebuilder_factory(resource).get("imagebuilder") + bucket = dummy_imagebuilder_bucket() + + generated_template = CDKTemplateBuilder().build_imagebuilder_template( + image_config=imagebuilder_config, + image_id="test-image", + bucket=bucket, + ) + + # Assert Role has correct path and name with ResourcePrefix + assert_that(generated_template["Resources"]["InstanceRole"]["Properties"]["Path"]).is_equal_to(expected_role_path) + assert_that(generated_template["Resources"]["InstanceRole"]["Properties"]["RoleName"]).is_equal_to(expected_role_name) + + # Assert InstanceProfile has correct path and name with ResourcePrefix + assert_that(generated_template["Resources"]["InstanceProfile"]["Properties"]["Path"]).is_equal_to(expected_profile_path) + assert_that(generated_template["Resources"]["InstanceProfile"]["Properties"]["InstanceProfileName"]).is_equal_to(expected_profile_name) + + # Assert inline policy name includes ResourcePrefix + policies = generated_template["Resources"]["InstanceRole"]["Properties"]["Policies"] + assert_that(len(policies)).is_equal_to(1) + if "test-prefix" in expected_role_name: + assert_that(policies[0]["PolicyName"]).starts_with("test-prefix") + elif "name-prefix-only" in expected_role_name: + assert_that(policies[0]["PolicyName"]).starts_with("name-prefix-only")