Merge pull request #175 from DreadPirateShawn/SnapshotAttributes

Implementation for ModifySnapshotAttribute and DescribeSnapshotAttribute.

moto/ec2/exceptions.py

@@ -110,6 +110,14 @@ class InvalidAMIIdError(EC2ClientError):
             .format(ami_id))
 
 
+class InvalidAMIAttributeItemValueError(EC2ClientError):
+    def __init__(self, attribute, value):
+        super(InvalidAMIAttributeItemValueError, self).__init__(
+            "InvalidAMIAttributeItemValue",
+            "Invalid attribute item value \"{0}\" for {1} item type."
+            .format(value, attribute))
+
+
 class InvalidSnapshotIdError(EC2ClientError):
     def __init__(self, snapshot_id):
         super(InvalidSnapshotIdError, self).__init__(

moto/ec2/models.py

@@ -29,6 +29,7 @@ from .exceptions import (
     InvalidPermissionNotFoundError,
     InvalidInstanceIdError,
     InvalidAMIIdError,
+    InvalidAMIAttributeItemValueError,
     InvalidSnapshotIdError,
     InvalidVolumeIdError,
     InvalidVolumeAttachmentError,
@@ -656,6 +657,7 @@ class Snapshot(object):
         self.id = snapshot_id
         self.volume = volume
         self.description = description
+        self.create_volume_permission_groups = set()
 
 
 class EBSBackend(object):
@@ -717,11 +719,41 @@ class EBSBackend(object):
     def describe_snapshots(self):
         return self.snapshots.values()
 
+    def get_snapshot(self, snapshot_id):
+        snapshot = self.snapshots.get(snapshot_id, None)
+        if not snapshot:
+            raise InvalidSnapshotIdError(snapshot_id)
+        return snapshot
+
     def delete_snapshot(self, snapshot_id):
         if snapshot_id in self.snapshots:
             return self.snapshots.pop(snapshot_id)
         raise InvalidSnapshotIdError(snapshot_id)
 
+    def get_create_volume_permission_groups(self, snapshot_id):
+        snapshot = self.get_snapshot(snapshot_id)
+        return snapshot.create_volume_permission_groups
+
+    def add_create_volume_permission(self, snapshot_id, user_id=None, group=None):
+        if user_id:
+            ec2_backend.raise_not_implemented_error("The UserId parameter for ModifySnapshotAttribute")
+
+        if group != 'all':
+            raise InvalidAMIAttributeItemValueError("UserGroup", group)
+        snapshot = self.get_snapshot(snapshot_id)
+        snapshot.create_volume_permission_groups.add(group)
+        return True
+
+    def remove_create_volume_permission(self, snapshot_id, user_id=None, group=None):
+        if user_id:
+            ec2_backend.raise_not_implemented_error("The UserId parameter for ModifySnapshotAttribute")
+
+        if group != 'all':
+            raise InvalidAMIAttributeItemValueError("UserGroup", group)
+        snapshot = self.get_snapshot(snapshot_id)
+        snapshot.create_volume_permission_groups.discard(group)
+        return True
+
 
 class VPC(TaggedEC2Instance):
     def __init__(self, vpc_id, cidr_block):
@@ -1380,6 +1412,12 @@ class EC2Backend(BaseBackend, InstanceBackend, TagBackend, AmiBackend,
     def raise_error(self, code, message):
         raise EC2ClientError(code, message)
 
+    def raise_not_implemented_error(self, blurb):
+        msg = "{0} has not been implemented in Moto yet." \
+              " Feel free to open an issue at" \
+              " https://github.com/spulec/moto/issues".format(blurb)
+        raise NotImplementedError(msg)
+
 
 ec2_backends = {}
 for region in boto.ec2.regions():

moto/ec2/responses/elastic_block_store.py

@@ -43,9 +43,6 @@ class ElasticBlockStore(BaseResponse):
         success = ec2_backend.delete_volume(volume_id)
         return DELETE_VOLUME_RESPONSE
 
-    def describe_snapshot_attribute(self):
-        raise NotImplementedError('ElasticBlockStore.describe_snapshot_attribute is not yet implemented')
-
     def describe_snapshots(self):
         snapshots = ec2_backend.describe_snapshots()
         template = Template(DESCRIBE_SNAPSHOTS_RESPONSE)
@@ -77,8 +74,22 @@ class ElasticBlockStore(BaseResponse):
     def import_volume(self):
         raise NotImplementedError('ElasticBlockStore.import_volume is not yet implemented')
 
+    def describe_snapshot_attribute(self):
+        snapshot_id = self.querystring.get('SnapshotId')[0]
+        groups = ec2_backend.get_create_volume_permission_groups(snapshot_id)
+        template = Template(DESCRIBE_SNAPSHOT_ATTRIBUTES_RESPONSE)
+        return template.render(snapshot_id=snapshot_id, groups=groups)
+
     def modify_snapshot_attribute(self):
-        raise NotImplementedError('ElasticBlockStore.modify_snapshot_attribute is not yet implemented')
+        snapshot_id = self.querystring.get('SnapshotId')[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]
+        if (operation_type == 'add'):
+            ec2_backend.add_create_volume_permission(snapshot_id, user_id=user_id, group=group)
+        elif (operation_type == 'remove'):
+            ec2_backend.remove_create_volume_permission(snapshot_id, user_id=user_id, group=group)
+        return MODIFY_SNAPSHOT_ATTRIBUTE_RESPONSE
 
     def modify_volume_attribute(self):
         raise NotImplementedError('ElasticBlockStore.modify_volume_attribute is not yet implemented')
@@ -186,3 +197,29 @@ DELETE_SNAPSHOT_RESPONSE = """<DeleteSnapshotResponse xmlns="http://ec2.amazonaw
   <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
   <return>true</return>
 </DeleteSnapshotResponse>"""
+
+DESCRIBE_SNAPSHOT_ATTRIBUTES_RESPONSE = """
+<DescribeSnapshotAttributeResponse xmlns="http://ec2.amazonaws.com/doc/2013-07-15/">
+    <requestId>a9540c9f-161a-45d8-9cc1-1182b89ad69f</requestId>
+    <snapshotId>snap-a0332ee0</snapshotId>
+   {% if not groups %}
+      <createVolumePermission/>
+   {% endif %}
+   {% if groups %}
+      <createVolumePermission>
+         {% for group in groups %}
+            <item>
+               <group>{{ group }}</group>
+            </item>
+         {% endfor %}
+      </createVolumePermission>
+   {% endif %}
+</DescribeSnapshotAttributeResponse>
+"""
+
+MODIFY_SNAPSHOT_ATTRIBUTE_RESPONSE = """
+<ModifySnapshotAttributeResponse xmlns="http://ec2.amazonaws.com/doc/2013-07-15/">
+    <requestId>666d2944-9276-4d6a-be12-1f4ada972fd8</requestId>
+    <return>true</return>
+</ModifySnapshotAttributeResponse>
+"""

tests/test_ec2/test_elastic_block_store.py

@@ -100,6 +100,87 @@ def test_create_snapshot():
     cm.exception.request_id.should_not.be.none
 
 
+@mock_ec2
+def test_snapshot_attribute():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    volume = conn.create_volume(80, "us-east-1a")
+    snapshot = volume.create_snapshot()
+
+    # Baseline
+    attributes = conn.get_snapshot_attribute(snapshot.id, attribute='createVolumePermission')
+    attributes.name.should.equal('create_volume_permission')
+    attributes.attrs.should.have.length_of(0)
+
+    ADD_GROUP_ARGS = {'snapshot_id': snapshot.id,
+                      'attribute': 'createVolumePermission',
+                      'operation': 'add',
+                      'groups': 'all'}
+
+    REMOVE_GROUP_ARGS = {'snapshot_id': snapshot.id,
+                         'attribute': 'createVolumePermission',
+                         'operation': 'remove',
+                         'groups': 'all'}
+
+    # Add 'all' group and confirm
+    conn.modify_snapshot_attribute(**ADD_GROUP_ARGS)
+
+    attributes = conn.get_snapshot_attribute(snapshot.id, attribute='createVolumePermission')
+    attributes.attrs['groups'].should.have.length_of(1)
+    attributes.attrs['groups'].should.equal(['all'])
+
+    # Add is idempotent
+    conn.modify_snapshot_attribute.when.called_with(**ADD_GROUP_ARGS).should_not.throw(EC2ResponseError)
+
+    # Remove 'all' group and confirm
+    conn.modify_snapshot_attribute(**REMOVE_GROUP_ARGS)
+
+    attributes = conn.get_snapshot_attribute(snapshot.id, attribute='createVolumePermission')
+    attributes.attrs.should.have.length_of(0)
+
+    # Remove is idempotent
+    conn.modify_snapshot_attribute.when.called_with(**REMOVE_GROUP_ARGS).should_not.throw(EC2ResponseError)
+
+    # Error: Add with group != 'all'
+    with assert_raises(EC2ResponseError) as cm:
+        conn.modify_snapshot_attribute(snapshot.id,
+                                       attribute='createVolumePermission',
+                                       operation='add',
+                                       groups='everyone')
+    cm.exception.code.should.equal('InvalidAMIAttributeItemValue')
+    cm.exception.status.should.equal(400)
+    cm.exception.request_id.should_not.be.none
+
+    # Error: Add with invalid snapshot ID
+    with assert_raises(EC2ResponseError) as cm:
+        conn.modify_snapshot_attribute("snapshot-abcd1234",
+                                       attribute='createVolumePermission',
+                                       operation='add',
+                                       groups='all')
+    cm.exception.code.should.equal('InvalidSnapshot.NotFound')
+    cm.exception.status.should.equal(400)
+    cm.exception.request_id.should_not.be.none
+
+    # Error: Remove with invalid snapshot ID
+    with assert_raises(EC2ResponseError) as cm:
+        conn.modify_snapshot_attribute("snapshot-abcd1234",
+                                       attribute='createVolumePermission',
+                                       operation='remove',
+                                       groups='all')
+    cm.exception.code.should.equal('InvalidSnapshot.NotFound')
+    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_snapshot_attribute.when.called_with(snapshot.id,
+                                                    attribute='createVolumePermission',
+                                                    operation='add',
+                                                    user_ids=['user']).should.throw(NotImplementedError)
+    conn.modify_snapshot_attribute.when.called_with(snapshot.id,
+                                                    attribute='createVolumePermission',
+                                                    operation='remove',
+                                                    user_ids=['user']).should.throw(NotImplementedError)
+
+
 @mock_ec2
 def test_modify_attribute_blockDeviceMapping():
     """