diff --git a/moto/ec2/models.py b/moto/ec2/models.py index a83284ef6..b4d491264 100644 --- a/moto/ec2/models.py +++ b/moto/ec2/models.py @@ -820,6 +820,7 @@ class Ami(TaggedEC2Resource): self.description = description if description else source_ami.description self.launch_permission_groups = set() + self.launch_permission_users = set() # AWS auto-creates these, we should reflect the same. volume = self.ec2_backend.create_volume(15, "us-east-1a") @@ -902,24 +903,49 @@ class AmiBackend(object): ami = self.describe_images(ami_ids=[ami_id])[0] return ami.launch_permission_groups - def add_launch_permission(self, ami_id, user_id=None, group=None): - if user_id: - self.raise_not_implemented_error("The UserId parameter for ModifyImageAttribute") - - if group != 'all': - raise InvalidAMIAttributeItemValueError("UserGroup", group) + def get_launch_permission_users(self, ami_id): ami = self.describe_images(ami_ids=[ami_id])[0] - ami.launch_permission_groups.add(group) + return ami.launch_permission_users + + def validate_permission_targets(self, user_ids=None, group=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) + + if group and group != 'all': + raise InvalidAMIAttributeItemValueError("UserGroup", group) + + def add_launch_permission(self, ami_id, user_ids=None, group=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) + return True - def remove_launch_permission(self, ami_id, user_id=None, group=None): - if user_id: - self.raise_not_implemented_error("The UserId parameter for ModifyImageAttribute") - - if group != 'all': - raise InvalidAMIAttributeItemValueError("UserGroup", group) + def remove_launch_permission(self, ami_id, user_ids=None, group=None): ami = self.describe_images(ami_ids=[ami_id])[0] - ami.launch_permission_groups.discard(group) + 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) + return True diff --git a/moto/ec2/responses/amis.py b/moto/ec2/responses/amis.py index 4af2b9be6..7d715fd51 100644 --- a/moto/ec2/responses/amis.py +++ b/moto/ec2/responses/amis.py @@ -1,6 +1,7 @@ from __future__ import unicode_literals from moto.core.responses import BaseResponse -from moto.ec2.utils import instance_ids_from_querystring, image_ids_from_querystring, filters_from_querystring +from moto.ec2.utils import instance_ids_from_querystring, image_ids_from_querystring, \ + filters_from_querystring, sequence_from_querystring class AmisResponse(BaseResponse): @@ -41,18 +42,19 @@ class AmisResponse(BaseResponse): def describe_image_attribute(self): ami_id = self.querystring.get('ImageId')[0] groups = self.ec2_backend.get_launch_permission_groups(ami_id) + users = self.ec2_backend.get_launch_permission_users(ami_id) template = self.response_template(DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE) - return template.render(ami_id=ami_id, groups=groups) + return template.render(ami_id=ami_id, groups=groups, users=users) def modify_image_attribute(self): ami_id = self.querystring.get('ImageId')[0] operation_type = self.querystring.get('OperationType')[0] group = self.querystring.get('UserGroup.1', [None])[0] - user_id = self.querystring.get('UserId.1', [None])[0] + user_ids = sequence_from_querystring('UserId', self.querystring) if (operation_type == 'add'): - self.ec2_backend.add_launch_permission(ami_id, user_id=user_id, group=group) + 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_id=user_id, group=group) + self.ec2_backend.remove_launch_permission(ami_id, user_ids=user_ids, group=group) return MODIFY_IMAGE_ATTRIBUTE_RESPONSE def register_image(self): @@ -140,16 +142,24 @@ DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE = """ 59dbff89-35bd-4eac-99ed-be587EXAMPLE {{ ami_id }} - {% if not groups %} + {% if not groups and not users %} - {% endif %} - {% if groups %} + {% else %} - {% for group in groups %} - - {{ group }} - - {% endfor %} + {% if groups %} + {% for group in groups %} + + {{ group }} + + {% endfor %} + {% endif %} + {% if users %} + {% for user in users %} + + {{ user }} + + {% endfor %} + {% endif %} {% endif %} """ diff --git a/tests/test_ec2/test_amis.py b/tests/test_ec2/test_amis.py index efc2c8f2d..71c47cf9a 100644 --- a/tests/test_ec2/test_amis.py +++ b/tests/test_ec2/test_amis.py @@ -238,7 +238,7 @@ def test_getting_malformed_ami(): @mock_ec2 -def test_ami_attribute(): +def test_ami_attribute_group_permissions(): conn = boto.connect_ec2('the_key', 'the_secret') reservation = conn.run_instances('ami-1234abcd') instance = reservation.instances[0] @@ -283,6 +283,132 @@ def test_ami_attribute(): # Remove is idempotent conn.modify_image_attribute.when.called_with(**REMOVE_GROUP_ARGS).should_not.throw(EC2ResponseError) + +@mock_ec2 +def test_ami_attribute_user_permissions(): + conn = boto.connect_ec2('the_key', 'the_secret') + reservation = conn.run_instances('ami-1234abcd') + instance = reservation.instances[0] + image_id = conn.create_image(instance.id, "test-ami", "this is a test ami") + image = conn.get_image(image_id) + + # Baseline + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.name.should.equal('launch_permission') + attributes.attrs.should.have.length_of(0) + + # Both str and int values should work. + USER1 = '123456789011' + USER2 = 123456789022 + + ADD_USERS_ARGS = {'image_id': image.id, + 'attribute': 'launchPermission', + 'operation': 'add', + 'user_ids': [USER1, USER2]} + + REMOVE_USERS_ARGS = {'image_id': image.id, + 'attribute': 'launchPermission', + 'operation': 'remove', + 'user_ids': [USER1, USER2]} + + REMOVE_SINGLE_USER_ARGS = {'image_id': image.id, + 'attribute': 'launchPermission', + 'operation': 'remove', + 'user_ids': [USER1]} + + # Add multiple users and confirm + conn.modify_image_attribute(**ADD_USERS_ARGS) + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs['user_ids'].should.have.length_of(2) + set(attributes.attrs['user_ids']).should.equal(set([str(USER1), str(USER2)])) + image = conn.get_image(image_id) + image.is_public.should.equal(False) + + # Add is idempotent + conn.modify_image_attribute.when.called_with(**ADD_USERS_ARGS).should_not.throw(EC2ResponseError) + + # Remove single user and confirm + conn.modify_image_attribute(**REMOVE_SINGLE_USER_ARGS) + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs['user_ids'].should.have.length_of(1) + set(attributes.attrs['user_ids']).should.equal(set([str(USER2)])) + image = conn.get_image(image_id) + image.is_public.should.equal(False) + + # Remove multiple users and confirm + conn.modify_image_attribute(**REMOVE_USERS_ARGS) + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs.should.have.length_of(0) + image = conn.get_image(image_id) + image.is_public.should.equal(False) + + # Remove is idempotent + conn.modify_image_attribute.when.called_with(**REMOVE_USERS_ARGS).should_not.throw(EC2ResponseError) + + +@mock_ec2 +def test_ami_attribute_user_and_group_permissions(): + """ + Boto supports adding/removing both users and groups at the same time. + Just spot-check this -- input variations, idempotency, etc are validated + via user-specific and group-specific tests above. + """ + conn = boto.connect_ec2('the_key', 'the_secret') + reservation = conn.run_instances('ami-1234abcd') + instance = reservation.instances[0] + image_id = conn.create_image(instance.id, "test-ami", "this is a test ami") + image = conn.get_image(image_id) + + # Baseline + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.name.should.equal('launch_permission') + attributes.attrs.should.have.length_of(0) + + USER1 = '123456789011' + USER2 = '123456789022' + + ADD_ARGS = {'image_id': image.id, + 'attribute': 'launchPermission', + 'operation': 'add', + 'groups': ['all'], + 'user_ids': [USER1, USER2]} + + REMOVE_ARGS = {'image_id': image.id, + 'attribute': 'launchPermission', + 'operation': 'remove', + 'groups': ['all'], + 'user_ids': [USER1, USER2]} + + # Add and confirm + conn.modify_image_attribute(**ADD_ARGS) + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs['user_ids'].should.have.length_of(2) + set(attributes.attrs['user_ids']).should.equal(set([USER1, USER2])) + set(attributes.attrs['groups']).should.equal(set(['all'])) + image = conn.get_image(image_id) + image.is_public.should.equal(True) + + # Remove and confirm + conn.modify_image_attribute(**REMOVE_ARGS) + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs.should.have.length_of(0) + image = conn.get_image(image_id) + image.is_public.should.equal(False) + + +@mock_ec2 +def test_ami_attribute_error_cases(): + conn = boto.connect_ec2('the_key', 'the_secret') + reservation = conn.run_instances('ami-1234abcd') + instance = reservation.instances[0] + image_id = conn.create_image(instance.id, "test-ami", "this is a test ami") + image = conn.get_image(image_id) + # Error: Add with group != 'all' with assert_raises(EC2ResponseError) as cm: conn.modify_image_attribute(image.id, @@ -293,6 +419,49 @@ def test_ami_attribute(): cm.exception.status.should.equal(400) cm.exception.request_id.should_not.be.none + # Error: Add with user ID that isn't an integer. + with assert_raises(EC2ResponseError) as cm: + conn.modify_image_attribute(image.id, + attribute='launchPermission', + operation='add', + user_ids='12345678901A') + cm.exception.code.should.equal('InvalidAMIAttributeItemValue') + cm.exception.status.should.equal(400) + cm.exception.request_id.should_not.be.none + + # Error: Add with user ID that is > length 12. + with assert_raises(EC2ResponseError) as cm: + conn.modify_image_attribute(image.id, + attribute='launchPermission', + operation='add', + user_ids='1234567890123') + cm.exception.code.should.equal('InvalidAMIAttributeItemValue') + cm.exception.status.should.equal(400) + cm.exception.request_id.should_not.be.none + + # Error: Add with user ID that is < length 12. + with assert_raises(EC2ResponseError) as cm: + conn.modify_image_attribute(image.id, + attribute='launchPermission', + operation='add', + user_ids='12345678901') + cm.exception.code.should.equal('InvalidAMIAttributeItemValue') + cm.exception.status.should.equal(400) + cm.exception.request_id.should_not.be.none + + # Error: Add with one invalid user ID among other valid IDs, ensure no partial changes. + with assert_raises(EC2ResponseError) as cm: + conn.modify_image_attribute(image.id, + attribute='launchPermission', + operation='add', + user_ids=['123456789011', 'foo', '123456789022']) + cm.exception.code.should.equal('InvalidAMIAttributeItemValue') + cm.exception.status.should.equal(400) + cm.exception.request_id.should_not.be.none + + attributes = conn.get_image_attribute(image.id, attribute='launchPermission') + attributes.attrs.should.have.length_of(0) + # Error: Add with invalid image ID with assert_raises(EC2ResponseError) as cm: conn.modify_image_attribute("ami-abcd1234", @@ -313,12 +482,3 @@ def test_ami_attribute(): cm.exception.status.should.equal(400) cm.exception.request_id.should_not.be.none - # Error: Add or remove with user ID instead of group - conn.modify_image_attribute.when.called_with(image.id, - attribute='launchPermission', - operation='add', - user_ids=['user']).should.throw(NotImplementedError) - conn.modify_image_attribute.when.called_with(image.id, - attribute='launchPermission', - operation='remove', - user_ids=['user']).should.throw(NotImplementedError)