From 9cbe8d57f40eb49f2be88ab22dde67bd33e4f1c1 Mon Sep 17 00:00:00 2001
From: Tyler Sanders <tylerssanders@gmail.com>
Date: Thu, 27 Nov 2014 11:05:39 -0600
Subject: [PATCH] Support instance security group/groupSet attribute
 modification and description

---
 moto/ec2/models.py               |  7 +++++++
 moto/ec2/responses/instances.py  | 33 ++++++++++++++++++++++++++++++--
 tests/test_ec2/test_instances.py | 16 ++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/moto/ec2/models.py b/moto/ec2/models.py
index e3a828954..ceb25d2cc 100644
--- a/moto/ec2/models.py
+++ b/moto/ec2/models.py
@@ -543,7 +543,14 @@ class InstanceBackend(object):
         setattr(instance, key, value)
         return instance
 
+    def modify_instance_security_groups(self, instance_id, new_group_list):
+        instance = self.get_instance(instance_id)
+        setattr(instance, 'security_groups', new_group_list)
+        return instance
+
     def describe_instance_attribute(self, instance_id, key):
+        if key == 'group_set':
+            key = 'security_groups'
         instance = self.get_instance(instance_id)
         value = getattr(instance, key)
         return instance, value
diff --git a/moto/ec2/responses/instances.py b/moto/ec2/responses/instances.py
index 8605969b8..f2fb290d2 100644
--- a/moto/ec2/responses/instances.py
+++ b/moto/ec2/responses/instances.py
@@ -88,12 +88,18 @@ class InstanceResponse(BaseResponse):
         instance_ids = instance_ids_from_querystring(self.querystring)
         instance_id = instance_ids[0]
         instance, value = self.ec2_backend.describe_instance_attribute(instance_id, key)
-        template = Template(EC2_DESCRIBE_INSTANCE_ATTRIBUTE)
+
+        if key == "group_set":
+            template = Template(EC2_DESCRIBE_INSTANCE_GROUPSET_ATTRIBUTE)
+        else:
+            template = Template(EC2_DESCRIBE_INSTANCE_ATTRIBUTE)
+
         return template.render(instance=instance, attribute=attribute, value=value)
 
     def modify_instance_attribute(self):
         handlers = [self._dot_value_instance_attribute_handler,
-                    self._block_device_mapping_handler]
+                    self._block_device_mapping_handler,
+                    self._security_grp_instance_attribute_handler]
 
         for handler in handlers:
             success = handler()
@@ -163,6 +169,17 @@ class InstanceResponse(BaseResponse):
         self.ec2_backend.modify_instance_attribute(instance_id, normalized_attribute, value)
         return EC2_MODIFY_INSTANCE_ATTRIBUTE
 
+    def _security_grp_instance_attribute_handler(self):
+        new_security_grp_list = []
+        for key, value in self.querystring.items():
+            if 'GroupId.' in key:
+                new_security_grp_list.append(self.querystring.get(key)[0])
+
+        instance_ids = instance_ids_from_querystring(self.querystring)
+        instance_id = instance_ids[0]
+        self.ec2_backend.modify_instance_security_groups(instance_id, new_security_grp_list)
+        return EC2_MODIFY_INSTANCE_ATTRIBUTE
+
 
 EC2_RUN_INSTANCES = """<RunInstancesResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
   <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
@@ -500,6 +517,18 @@ EC2_DESCRIBE_INSTANCE_ATTRIBUTE = """<DescribeInstanceAttributeResponse xmlns="h
   </{{ attribute }}>
 </DescribeInstanceAttributeResponse>"""
 
+EC2_DESCRIBE_INSTANCE_GROUPSET_ATTRIBUTE = """<DescribeInstanceAttributeResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
+  <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
+  <instanceId>{{ instance.id }}</instanceId>
+  <{{ attribute }}>
+    {% for sg_id in value %}
+      <item>
+        <groupId>{{ sg_id }}</groupId>
+      </item>
+    {% endfor %}
+  </{{ attribute }}>
+</DescribeInstanceAttributeResponse>"""
+
 EC2_MODIFY_INSTANCE_ATTRIBUTE = """<ModifyInstanceAttributeResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
   <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
   <return>true</return>
diff --git a/tests/test_ec2/test_instances.py b/tests/test_ec2/test_instances.py
index ec404b641..a63f939b1 100644
--- a/tests/test_ec2/test_instances.py
+++ b/tests/test_ec2/test_instances.py
@@ -274,6 +274,22 @@ def test_instance_attribute_instance_type():
     instance_attribute.should.be.a(InstanceAttribute)
     instance_attribute.get('instanceType').should.equal("m1.small")
 
+@mock_ec2
+def test_modify_instance_attribute_security_groups():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    reservation = conn.run_instances('ami-1234abcd')
+    instance = reservation.instances[0]
+
+    sg_id = 'sg-1234abcd'
+    sg_id2 = 'sg-abcd4321'
+    instance.modify_attribute("groupSet", [sg_id, sg_id2])
+
+    instance_attribute = instance.get_attribute("groupSet")
+    instance_attribute.should.be.a(InstanceAttribute)
+    group_list = instance_attribute.get('groupSet')
+    any(g.id == sg_id for g in group_list).should.be.ok
+    any(g.id == sg_id2 for g in group_list).should.be.ok
+
 
 @mock_ec2
 def test_instance_attribute_user_data():