Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions cli/src/pcluster/config/imagebuilder_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -104,13 +104,15 @@ 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)
self.cleanup_lambda_role = Resource.init_param(cleanup_lambda_role)
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]:
Expand All @@ -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."""
Expand Down
1 change: 1 addition & 0 deletions cli/src/pcluster/schemas/imagebuilder_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
27 changes: 21 additions & 6 deletions cli/src/pcluster/templates/imagebuilder_stack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down
134 changes: 134 additions & 0 deletions cli/tests/pcluster/schemas/test_imagebuilder_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
105 changes: 105 additions & 0 deletions cli/tests/pcluster/templates/test_imagebuilder_stack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")