From 2c3f735e8539a401bc474043a37df247140ed295 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Thu, 8 Feb 2024 22:21:42 +0000 Subject: [PATCH] EC2: modify_image_attributes() now supports all LaunchPermissions (#7324) --- IMPLEMENTATION_COVERAGE.md | 2 +- docs/docs/services/ec2.rst | 2 +- moto/ec2/models/amis.py | 81 +++++++++++++++---------------------- moto/ec2/responses/amis.py | 77 ++++++++++++++++++++--------------- tests/test_ec2/test_amis.py | 51 ++++++++++++++++++++++- 5 files changed, 129 insertions(+), 84 deletions(-) diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md index d7cb91469..cb973e06d 100644 --- a/IMPLEMENTATION_COVERAGE.md +++ b/IMPLEMENTATION_COVERAGE.md @@ -2523,7 +2523,7 @@ - [X] modify_hosts - [ ] modify_id_format - [ ] modify_identity_id_format -- [ ] modify_image_attribute +- [X] modify_image_attribute - [X] modify_instance_attribute - [ ] modify_instance_capacity_reservation_attributes - [ ] modify_instance_credit_specification diff --git a/docs/docs/services/ec2.rst b/docs/docs/services/ec2.rst index 4398c475c..2750ebd76 100644 --- a/docs/docs/services/ec2.rst +++ b/docs/docs/services/ec2.rst @@ -545,7 +545,7 @@ ec2 - [X] modify_hosts - [ ] modify_id_format - [ ] modify_identity_id_format -- [ ] modify_image_attribute +- [X] modify_image_attribute - [X] modify_instance_attribute - [ ] modify_instance_capacity_reservation_attributes - [ ] modify_instance_credit_specification diff --git a/moto/ec2/models/amis.py b/moto/ec2/models/amis.py index 12028f2cd..862b97536 100644 --- a/moto/ec2/models/amis.py +++ b/moto/ec2/models/amis.py @@ -101,11 +101,10 @@ class Ami(TaggedEC2Resource): if not description: self.description = source_ami.description - self.launch_permission_groups: Set[str] = set() - self.launch_permission_users: Set[str] = set() + self.launch_permissions: List[Dict[str, str]] = [] if public: - self.launch_permission_groups.add("all") + self.launch_permissions.append({"Group": "all"}) # AWS auto-creates these, we should reflect the same. volume = self.ec2_backend.create_volume(size=15, zone_name=region_name) @@ -119,7 +118,7 @@ class Ami(TaggedEC2Resource): @property def is_public(self) -> bool: - return "all" in self.launch_permission_groups + return {"Group": "all"} in self.launch_permissions @property def is_public_string(self) -> str: @@ -273,8 +272,9 @@ class AmiBackend: tmp_images = [] for ami in images: for user_id in exec_users: - if user_id in ami.launch_permission_users: - tmp_images.append(ami) + for lp in ami.launch_permissions: + if lp.get("UserId") == user_id: + tmp_images.append(ami) images = tmp_images # Limit by owner ids @@ -305,38 +305,39 @@ class AmiBackend: else: raise InvalidAMIIdError(ami_id) - def validate_permission_targets( - self, user_ids: Optional[List[str]] = None, group: Optional[str] = None - ) -> None: - # If anything is invalid, nothing is added. (No partial success.) - if user_ids: - """ - AWS docs: - "The AWS account ID is a 12-digit number, such as 123456789012, that you use to construct Amazon Resource Names (ARNs)." - http://docs.aws.amazon.com/general/latest/gr/acct-identifiers.html - """ - for user_id in user_ids: - if len(user_id) != 12 or not user_id.isdigit(): - raise InvalidAMIAttributeItemValueError("userId", user_id) + def validate_permission_targets(self, permissions: List[Dict[str, str]]) -> None: + for perm in permissions: + # If anything is invalid, nothing is added. (No partial success.) + # AWS docs: + # The AWS account ID is a 12-digit number, such as 123456789012, that you use to construct Amazon Resource Names (ARNs)." + # http://docs.aws.amazon.com/general/latest/gr/acct-identifiers.html - if group and group != "all": - raise InvalidAMIAttributeItemValueError("UserGroup", group) + if "UserId" in perm and ( + len(perm["UserId"]) != 12 or not perm["UserId"].isdigit() + ): + raise InvalidAMIAttributeItemValueError("userId", perm["UserId"]) - def add_launch_permission( + if "Group" in perm and perm["Group"] != "all": + raise InvalidAMIAttributeItemValueError("UserGroup", perm["Group"]) + + def modify_image_attribute( self, ami_id: str, - user_ids: Optional[List[str]] = None, - group: Optional[str] = None, + launch_permissions_to_add: List[Dict[str, str]], + launch_permissions_to_remove: List[Dict[str, str]], ) -> None: ami = self.describe_images(ami_ids=[ami_id])[0] - self.validate_permission_targets(user_ids=user_ids, group=group) - - if user_ids: - for user_id in user_ids: - ami.launch_permission_users.add(user_id) - - if group: - ami.launch_permission_groups.add(group) + self.validate_permission_targets(launch_permissions_to_add) + self.validate_permission_targets(launch_permissions_to_remove) + for lp in launch_permissions_to_add: + if lp not in ami.launch_permissions: + ami.launch_permissions.append(lp) + for lp in launch_permissions_to_remove: + try: + ami.launch_permissions.remove(lp) + except ValueError: + # The LaunchPermission may not exist + pass def register_image( self, name: Optional[str] = None, description: Optional[str] = None @@ -353,21 +354,5 @@ class AmiBackend: self.amis[ami_id] = ami return ami - def remove_launch_permission( - self, - ami_id: str, - user_ids: Optional[List[str]] = None, - group: Optional[str] = None, - ) -> None: - ami = self.describe_images(ami_ids=[ami_id])[0] - self.validate_permission_targets(user_ids=user_ids, group=group) - - if user_ids: - for user_id in user_ids: - ami.launch_permission_users.discard(user_id) - - if group: - ami.launch_permission_groups.discard(group) - def describe_image_attribute(self, ami_id: str, attribute_name: str) -> Any: return self.amis[ami_id].__getattribute__(attribute_name) diff --git a/moto/ec2/responses/amis.py b/moto/ec2/responses/amis.py index c9e36607d..7238c32e5 100644 --- a/moto/ec2/responses/amis.py +++ b/moto/ec2/responses/amis.py @@ -65,10 +65,7 @@ class AmisResponse(EC2BaseResponse): "description": "description", "kernel": "kernel_id", "ramdisk": "ramdisk", - "launchPermission": { - "groups": "launch_permission_groups", - "users": "launch_permission_users", - }, + "launchPermission": "launch_permissions", "productCodes": "product_codes", "blockDeviceMapping": "bdm", "sriovNetSupport": "sriov", @@ -85,45 +82,63 @@ class AmisResponse(EC2BaseResponse): # https://github.com/aws/aws-cli/issues/1083 raise AuthFailureRestricted - groups = None - users = None attribute_value = None + launch_permissions = None if attribute_name == "launchPermission": - groups = self.ec2_backend.describe_image_attribute( - ami_id, valid_atrributes_list[attribute_name]["groups"] # type: ignore[index] - ) - users = self.ec2_backend.describe_image_attribute( - ami_id, valid_atrributes_list[attribute_name]["users"] # type: ignore[index] + launch_permissions = self.ec2_backend.describe_image_attribute( + ami_id, valid_atrributes_list[attribute_name] ) else: attribute_value = self.ec2_backend.describe_image_attribute( ami_id, valid_atrributes_list[attribute_name] ) + template = self.response_template(DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE) return template.render( ami_id=ami_id, - users=users, - groups=groups, + launch_permissions=launch_permissions, attribute_name=attribute_name, attribute_value=attribute_value, ) def modify_image_attribute(self) -> str: ami_id = self._get_param("ImageId") + launch_permissions_to_add = list( + self._get_params().get("LaunchPermission", {}).get("Add", {}).values() + ) + launch_permissions_to_remove = list( + self._get_params().get("LaunchPermission", {}).get("Remove", {}).values() + ) + # If only one OperationType is added, the other attributes are submitted as different variables operation_type = self._get_param("OperationType") - group = self._get_param("UserGroup.1") - user_ids = self._get_multi_param("UserId") + if operation_type in ["add", "remove"]: + group = self._get_param("UserGroup.1") + lp = ( + launch_permissions_to_add + if operation_type == "add" + else launch_permissions_to_remove + ) + if group: + lp.append({"Group": group}) + + for user_id in self._get_multi_param("UserId"): + lp.append({"UserId": user_id}) + + org_arn = self._get_param("OrganizationArn.1") + if org_arn: + lp.append({"OrganizationArn": org_arn}) + + ou_arn = self._get_param("OrganizationalUnitArn.1") + if ou_arn: + lp.append({"OrganizationalUnitArn": ou_arn}) self.error_on_dryrun() - if operation_type == "add": - self.ec2_backend.add_launch_permission( - ami_id, user_ids=user_ids, group=group - ) - elif operation_type == "remove": - self.ec2_backend.remove_launch_permission( - ami_id, user_ids=user_ids, group=group - ) + self.ec2_backend.modify_image_attribute( + ami_id=ami_id, + launch_permissions_to_add=launch_permissions_to_add, + launch_permissions_to_remove=launch_permissions_to_remove, + ) return MODIFY_IMAGE_ATTRIBUTE_RESPONSE def register_image(self) -> str: @@ -232,17 +247,13 @@ DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE = """ {% endfor %} {% endif %} {% if attribute_name == 'launchPermission' %} - {% if groups %} - {% for group in groups %} + {% if launch_permissions %} + {% for lp in launch_permissions %} - {{ group }} - - {% endfor %} - {% endif %} - {% if users %} - {% for user in users %} - - {{ user }} + {% if lp['UserId'] %}{{ lp['UserId'] }}{% endif %} + {% if lp['Group'] %}{{ lp['Group'] }}{% endif %} + {% if lp['OrganizationArn'] %}{{ lp['OrganizationArn'] }}{% endif %} + {% if lp['OrganizationalUnitArn'] %}{{ lp['OrganizationalUnitArn'] }}{% endif %} {% endfor %} {% endif %} diff --git a/tests/test_ec2/test_amis.py b/tests/test_ec2/test_amis.py index e8bff43a0..87905a3ce 100644 --- a/tests/test_ec2/test_amis.py +++ b/tests/test_ec2/test_amis.py @@ -426,6 +426,9 @@ def test_ami_filters(): ] imageB = boto3.resource("ec2", region_name="us-east-1").Image(imageB_id) imageB.modify_attribute(LaunchPermission={"Add": [{"Group": "all"}]}) + assert imageB.describe_attribute(Attribute="launchPermission")[ + "LaunchPermissions" + ] == [{"Group": "all"}] amis_by_architecture = ec2.describe_images( Filters=[{"Name": "architecture", "Values": ["x86_64"]}] @@ -479,7 +482,6 @@ def test_ami_filters(): Filters=[{"Name": "is-public", "Values": ["false"]}] )["Images"] assert imageA.id in [ami["ImageId"] for ami in amis_by_nonpublic] - assert len(amis_by_nonpublic) >= 2, "Should have at least 2 non-public images" @mock_aws @@ -684,6 +686,53 @@ def test_ami_attribute_user_permissions(): image.modify_attribute(**REMOVE_USERS_ARGS) +@mock_aws +def test_ami_attribute_organizations(): + ec2 = boto3.client("ec2", region_name="us-east-1") + reservation = ec2.run_instances(ImageId=EXAMPLE_AMI_ID, MinCount=1, MaxCount=1) + instance_id = reservation["Instances"][0]["InstanceId"] + image_id = ec2.create_image(InstanceId=instance_id, Name="test-ami-A")["ImageId"] + image = boto3.resource("ec2", "us-east-1").Image(image_id) + arn = "someOrganizationArn" + image.modify_attribute( + Attribute="launchPermission", + OperationType="add", + OrganizationArns=[arn], + ) + image.modify_attribute( + Attribute="launchPermission", + OperationType="add", + OrganizationalUnitArns=["ou1"], + ) + + ec2.modify_image_attribute( + Attribute="launchPermission", + ImageId=image_id, + LaunchPermission={ + "Add": [ + {"UserId": "111122223333"}, + {"UserId": "555566667777"}, + {"Group": "all"}, + {"OrganizationArn": "orgarn"}, + {"OrganizationalUnitArn": "ou2"}, + ] + }, + ) + + launch_permissions = image.describe_attribute(Attribute="launchPermission")[ + "LaunchPermissions" + ] + assert launch_permissions == [ + {"OrganizationArn": "someOrganizationArn"}, + {"OrganizationalUnitArn": "ou1"}, + {"UserId": "111122223333"}, + {"UserId": "555566667777"}, + {"Group": "all"}, + {"OrganizationArn": "orgarn"}, + {"OrganizationalUnitArn": "ou2"}, + ] + + @mock_aws def test_ami_describe_executable_users(): conn = boto3.client("ec2", region_name="us-east-1")