add security group ingress rules

moto/ec2/models.py

@@ -193,11 +193,35 @@ class RegionsAndZonesBackend(object):
         return self.zones
 
 
+class SecurityRule(object):
+    def __init__(self, ip_protocol, from_port, to_port, ip_ranges, source_groups):
+        self.ip_protocol = ip_protocol
+        self.from_port = from_port
+        self.to_port = to_port
+        self.ip_ranges = ip_ranges or []
+        self.source_groups = source_groups
+
+    @property
+    def unique_representation(self):
+        return "{}-{}-{}-{}-{}".format(
+                self.ip_protocol,
+                self.from_port,
+                self.to_port,
+                self.ip_ranges,
+                self.source_groups
+            )
+
+    def __eq__(self, other):
+        return self.unique_representation == other.unique_representation
+
+
 class SecurityGroup(object):
     def __init__(self, group_id, name, description):
         self.id = group_id
         self.name = name
         self.description = description
+        self.ingress_rules = []
+        self.egress_rules = []
 
 
 class SecurityGroupBackend(object):
@@ -232,6 +256,28 @@ class SecurityGroupBackend(object):
             if group.name == name:
                 return group
 
+    def authorize_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None):
+        group = self.get_security_group_from_name(group_name)
+        source_groups = []
+        for source_group_name in source_group_names:
+            source_groups.append(self.get_security_group_from_name(source_group_name))
+
+        security_rule = SecurityRule(ip_protocol, from_port, to_port, ip_ranges, source_groups)
+        group.ingress_rules.append(security_rule)
+
+    def revoke_security_group_ingress(self, group_name, ip_protocol, from_port, to_port, ip_ranges=None, source_group_names=None):
+        group = self.get_security_group_from_name(group_name)
+        source_groups = []
+        for source_group_name in source_group_names:
+            source_groups.append(self.get_security_group_from_name(source_group_name))
+
+        security_rule = SecurityRule(ip_protocol, from_port, to_port, ip_ranges, source_groups)
+        if security_rule in group.ingress_rules:
+            group.ingress_rules.remove(security_rule)
+            return security_rule
+        return False
+
+
 class EC2Backend(BaseBackend, InstanceBackend, TagBackend, AmiBackend, RegionsAndZonesBackend, SecurityGroupBackend):
     pass
 

moto/ec2/responses/security_groups.py

@@ -4,6 +4,23 @@ from moto.ec2.models import ec2_backend
 from moto.ec2.utils import resource_ids_from_querystring
 
 
+def process_rules_from_querystring(querystring):
+    name = querystring.get('GroupName')[0]
+    ip_protocol = querystring.get('IpPermissions.1.IpProtocol')[0]
+    from_port = querystring.get('IpPermissions.1.FromPort')[0]
+    to_port = querystring.get('IpPermissions.1.ToPort')[0]
+    ip_ranges = []
+    for key, value in querystring.iteritems():
+        if 'IpPermissions.1.IpRanges' in key:
+            ip_ranges.append(value[0])
+
+    source_groups = []
+    for key, value in querystring.iteritems():
+        if 'IpPermissions.1.Groups' in key:
+            source_groups.append(value[0])
+    return (name, ip_protocol, from_port, to_port, ip_ranges, source_groups)
+
+
 class SecurityGroups(object):
     def __init__(self, querystring):
         self.querystring = querystring
@@ -12,7 +29,8 @@ class SecurityGroups(object):
         raise NotImplementedError('SecurityGroups.authorize_security_group_egress is not yet implemented')
 
     def authorize_security_group_ingress(self):
-        raise NotImplementedError('SecurityGroups.authorize_security_group_ingress is not yet implemented')
+        ec2_backend.authorize_security_group_ingress(*process_rules_from_querystring(self.querystring))
+        return AUTHORIZE_SECURITY_GROUP_INGRESS_REPONSE
 
     def create_security_group(self):
         name = self.querystring.get('GroupName')[0]
@@ -25,6 +43,7 @@ class SecurityGroups(object):
         return template.render(group=group)
 
     def delete_security_group(self):
+        # TODO this should raise an error if there are instances in the group. See http://docs.aws.amazon.com/AWSEC2/latest/APIReference/ApiReference-query-DeleteSecurityGroup.html
         name = self.querystring.get('GroupName')[0]
         group = ec2_backend.delete_security_group(name)
 
@@ -42,7 +61,10 @@ class SecurityGroups(object):
         raise NotImplementedError('SecurityGroups.revoke_security_group_egress is not yet implemented')
 
     def revoke_security_group_ingress(self):
-        raise NotImplementedError('SecurityGroups.revoke_security_group_ingress is not yet implemented')
+        success = ec2_backend.revoke_security_group_ingress(*process_rules_from_querystring(self.querystring))
+        if not success:
+            return "Could not find a matching ingress rule", dict(status=404)
+        return REVOKE_SECURITY_GROUP_INGRESS_REPONSE
 
 
 CREATE_SECURITY_GROUP_RESPONSE = """<CreateSecurityGroupResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
@@ -67,20 +89,42 @@ DESCRIBE_SECURITY_GROUPS_RESPONSE = """<DescribeSecurityGroupsResponse xmlns="ht
              <groupDescription>{{ group.description }}</groupDescription>
              <vpcId/>
              <ipPermissions>
-                <item>
+               {% for rule in group.ingress_rules %}
-                   <ipProtocol>tcp</ipProtocol>
+                    <item>
-                   <fromPort>80</fromPort>
+                       <ipProtocol>{{ rule.ip_protocol }}</ipProtocol>
-                   <toPort>80</toPort>
+                       <fromPort>{{ rule.from_port }}</fromPort>
-                   <groups/>
+                       <toPort>{{ rule.to_port }}</toPort>
-                   <ipRanges>
+                       <groups>
-                      <item>
+                          {% for source_group in rule.source_groups %}
-                         <cidrIp>0.0.0.0/0</cidrIp>
+                              <item>
-                      </item>
+                                 <userId>111122223333</userId>
-                   </ipRanges>
+                                 <groupId>{{ source_group.id }}</groupId>
-                </item>
+                                 <groupName>{{ source_group.name }}</groupName>
+                              </item>
+                          {% endfor %}
+                       </groups>
+                       <ipRanges>
+                          {% for ip_range in rule.ip_ranges %}
+                              <item>
+                                 <cidrIp>{{ ip_range }}</cidrIp>
+                              </item>
+                          {% endfor %}
+                       </ipRanges>
+                    </item>
+                {% endfor %}
              </ipPermissions>
              <ipPermissionsEgress/>
           </item>
       {% endfor %}
    </securityGroupInfo>
 </DescribeSecurityGroupsResponse>"""
+
+AUTHORIZE_SECURITY_GROUP_INGRESS_REPONSE = """<AuthorizeSecurityGroupIngressResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
+  <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
+  <return>true</return>
+</AuthorizeSecurityGroupIngressResponse>"""
+
+REVOKE_SECURITY_GROUP_INGRESS_REPONSE = """<RevokeSecurityGroupIngressResponse xmlns="http://ec2.amazonaws.com/doc/2012-12-01/">
+  <requestId>59dbff89-35bd-4eac-99ed-be587EXAMPLE</requestId>
+  <return>true</return>
+</RevokeSecurityGroupIngressResponse>"""

tests/test_ec2/test_security_groups.py

@@ -39,3 +39,48 @@ def test_deleting_security_groups():
     # Delete by group id
     conn.delete_security_group(security_group1.id)
     conn.get_all_security_groups().should.have.length_of(0)
+
+
+@mock_ec2
+def test_authorize_ip_range_and_revoke():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    security_group = conn.create_security_group('test', 'test')
+
+    success = security_group.authorize(ip_protocol="tcp", from_port="22", to_port="2222", cidr_ip="123.123.123.123/32")
+    assert success.should.be.true
+
+    security_group = conn.get_all_security_groups()[0]
+    int(security_group.rules[0].to_port).should.equal(2222)
+    security_group.rules[0].grants[0].cidr_ip.should.equal("123.123.123.123/32")
+
+    # Wrong Cidr should throw error
+    security_group.revoke.when.called_with(ip_protocol="tcp", from_port="22", to_port="2222", cidr_ip="123.123.123.122/32").should.throw(EC2ResponseError)
+
+    # Actually revoke
+    security_group.revoke(ip_protocol="tcp", from_port="22", to_port="2222", cidr_ip="123.123.123.123/32")
+
+    security_group = conn.get_all_security_groups()[0]
+    security_group.rules.should.have.length_of(0)
+
+@mock_ec2
+def test_authorize_other_group_and_revoke():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    security_group = conn.create_security_group('test', 'test')
+    other_security_group = conn.create_security_group('other', 'other')
+    wrong_group = conn.create_security_group('wrong', 'wrong')
+
+    success = security_group.authorize(ip_protocol="tcp", from_port="22", to_port="2222", src_group=other_security_group)
+    assert success.should.be.true
+
+    security_group = [group for group in conn.get_all_security_groups() if group.name == 'test'][0]
+    int(security_group.rules[0].to_port).should.equal(2222)
+    security_group.rules[0].grants[0].group_id.should.equal(other_security_group.id)
+
+    # Wrong source group should throw error
+    security_group.revoke.when.called_with(ip_protocol="tcp", from_port="22", to_port="2222", src_group=wrong_group).should.throw(EC2ResponseError)
+
+    # Actually revoke
+    security_group.revoke(ip_protocol="tcp", from_port="22", to_port="2222", src_group=other_security_group)
+
+    security_group = [group for group in conn.get_all_security_groups() if group.name == 'test'][0]
+    security_group.rules.should.have.length_of(0)