Merge pull request #289 from DreadPirateShawn/ModifyImageAttributeUserPermissions

ModifyImageAttribute: Added support for user permissions.

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):
+    def get_launch_permission_users(self, ami_id):
-        if user_id:
-            self.raise_not_implemented_error("The UserId parameter for ModifyImageAttribute")
-
-        if group != 'all':
-            raise InvalidAMIAttributeItemValueError("UserGroup", group)
         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):
+    def remove_launch_permission(self, ami_id, user_ids=None, group=None):
-        if user_id:
-            self.raise_not_implemented_error("The UserId parameter for ModifyImageAttribute")
-
-        if group != 'all':
-            raise InvalidAMIAttributeItemValueError("UserGroup", group)
         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
 
 

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 = """
 <DescribeImageAttributeResponse xmlns="http://ec2.amazonaws.com/doc/2013-08-15/">
    <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
    <imageId>{{ ami_id }}</imageId>
-   {% if not groups %}
+   {% if not groups and not users %}
       <launchPermission/>
-   {% endif %}
+   {% else %}
-   {% if groups %}
       <launchPermission>
-         {% for group in groups %}
+         {% if groups %}
-            <item>
+            {% for group in groups %}
-               <group>{{ group }}</group>
+               <item>
-            </item>
+                  <group>{{ group }}</group>
-         {% endfor %}
+               </item>
+            {% endfor %}
+         {% endif %}
+         {% if users %}
+            {% for user in users %}
+               <item>
+                  <userId>{{ user }}</userId>
+               </item>
+            {% endfor %}
+         {% endif %}
       </launchPermission>
    {% endif %}
 </DescribeImageAttributeResponse>"""

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)