made the security group endpoints that authorize or revoke firewall rules

to support batch rules (boto doesn't expose this, but botocore/boto3 does)

moto/ec2/responses/security_groups.py

@@ -1,43 +1,74 @@
 from __future__ import unicode_literals
+
+import collections
+
 from moto.core.responses import BaseResponse
 from moto.ec2.utils import filters_from_querystring
 
 
+def try_parse_int(value, default=None):
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return default
+
+
 def process_rules_from_querystring(querystring):
     try:
         group_name_or_id = querystring.get('GroupName')[0]
     except:
         group_name_or_id = querystring.get('GroupId')[0]
 
-    ip_protocol = querystring.get('IpPermissions.1.IpProtocol', [None])[0]
+    querytree = {}
-    from_port = querystring.get('IpPermissions.1.FromPort', [None])[0]
-    to_port = querystring.get('IpPermissions.1.ToPort', [None])[0]
-    ip_ranges = []
     for key, value in querystring.items():
-        if 'IpPermissions.1.IpRanges' in key:
+        key_splitted = key.split('.')
-            ip_ranges.append(value[0])
+        key_splitted = [try_parse_int(e, e) for e in key_splitted]
 
-    source_groups = []
+        d = querytree
-    source_group_ids = []
+        for subkey in key_splitted[:-1]:
+            if subkey not in d:
+                d[subkey] = {}
+            d = d[subkey]
+        d[key_splitted[-1]] = value
 
-    for key, value in querystring.items():
+    ip_permissions = querytree.get('IpPermissions') or {}
-        if 'IpPermissions.1.Groups.1.GroupId' in key:
+    for ip_permission_idx in sorted(ip_permissions.keys()):
-            source_group_ids.append(value[0])
+        ip_permission = ip_permissions[ip_permission_idx]
-        elif 'IpPermissions.1.Groups' in key:
-            source_groups.append(value[0])
 
-    return (group_name_or_id, ip_protocol, from_port, to_port, ip_ranges, source_groups, source_group_ids)
+        ip_protocol = ip_permission.get('IpProtocol', [None])[0]
+        from_port = ip_permission.get('FromPort', [None])[0]
+        to_port = ip_permission.get('ToPort', [None])[0]
+
+        ip_ranges = []
+        ip_ranges_tree = ip_permission.get('IpRanges') or {}
+        for ip_range_idx in sorted(ip_ranges_tree.keys()):
+            ip_ranges.append(ip_ranges_tree[ip_range_idx]['CidrIp'][0])
+
+        source_groups = []
+        source_group_ids = []
+        groups_tree = ip_permission.get('Groups') or {}
+        for group_idx in sorted(groups_tree.keys()):
+            group_dict = groups_tree[group_idx]
+            if 'GroupId' in group_dict:
+                source_group_ids.append(group_dict['GroupId'][0])
+            elif 'GroupName' in group_dict:
+                source_groups.append(group_dict['GroupName'][0])
+
+        yield (group_name_or_id, ip_protocol, from_port, to_port, ip_ranges,
+                source_groups, source_group_ids)
 
 
 class SecurityGroups(BaseResponse):
     def authorize_security_group_egress(self):
         if self.is_not_dryrun('GrantSecurityGroupEgress'):
-            self.ec2_backend.authorize_security_group_egress(*process_rules_from_querystring(self.querystring))
+            for args in process_rules_from_querystring(self.querystring):
+                self.ec2_backend.authorize_security_group_egress(*args)
             return AUTHORIZE_SECURITY_GROUP_EGRESS_RESPONSE
 
     def authorize_security_group_ingress(self):
         if self.is_not_dryrun('GrantSecurityGroupIngress'):
-            self.ec2_backend.authorize_security_group_ingress(*process_rules_from_querystring(self.querystring))
+            for args in process_rules_from_querystring(self.querystring):
+                self.ec2_backend.authorize_security_group_ingress(*args)
             return AUTHORIZE_SECURITY_GROUP_INGRESS_REPONSE
 
     def create_security_group(self):
@@ -80,14 +111,16 @@ class SecurityGroups(BaseResponse):
 
     def revoke_security_group_egress(self):
         if self.is_not_dryrun('RevokeSecurityGroupEgress'):
-            success = self.ec2_backend.revoke_security_group_egress(*process_rules_from_querystring(self.querystring))
+            for args in process_rules_from_querystring(self.querystring):
-            if not success:
+                success = self.ec2_backend.revoke_security_group_egress(*args)
-                return "Could not find a matching egress rule", dict(status=404)
+                if not success:
+                    return "Could not find a matching egress rule", dict(status=404)
             return REVOKE_SECURITY_GROUP_EGRESS_RESPONSE
 
     def revoke_security_group_ingress(self):
         if self.is_not_dryrun('RevokeSecurityGroupIngress'):
-            self.ec2_backend.revoke_security_group_ingress(*process_rules_from_querystring(self.querystring))
+            for args in process_rules_from_querystring(self.querystring):
+                self.ec2_backend.revoke_security_group_ingress(*args)
             return REVOKE_SECURITY_GROUP_INGRESS_REPONSE
 
 

tests/test_ec2/test_security_groups.py

@@ -1,4 +1,7 @@
 from __future__ import unicode_literals
+
+import copy
+
 # Ensure 'assert_raises' context manager support for Python 2.6
 import tests.backport_assert_raises  # noqa
 from nose.tools import assert_raises
@@ -406,6 +409,7 @@ def test_authorize_and_revoke_in_bulk():
 
     sg01 = ec2.create_security_group(GroupName='sg01', Description='Test security group sg01', VpcId=vpc.id)
     sg02 = ec2.create_security_group(GroupName='sg02', Description='Test security group sg02', VpcId=vpc.id)
+    sg03 = ec2.create_security_group(GroupName='sg03', Description='Test security group sg03')
 
     ip_permissions = [
         {
@@ -420,27 +424,37 @@ def test_authorize_and_revoke_in_bulk():
             'IpProtocol': 'tcp',
             'FromPort': 27017,
             'ToPort': 27017,
-            'UserIdGroupPairs': [{'GroupId': sg02.id, 'GroupName': 'sg02', 'UserId': sg02.owner_id}],
+            'UserIdGroupPairs': [{'GroupId': sg02.id, 'UserId': sg02.owner_id}],
+            'IpRanges': []
+        },
+        {
+            'IpProtocol': 'tcp',
+            'FromPort': 27017,
+            'ToPort': 27017,
+            'UserIdGroupPairs': [{'GroupName': 'sg03', 'UserId': sg03.owner_id}],
             'IpRanges': []
         }
     ]
+    expected_ip_permissions = copy.deepcopy(ip_permissions)
+    expected_ip_permissions[1]['UserIdGroupPairs'][0]['GroupName'] = 'sg02'
+    expected_ip_permissions[2]['UserIdGroupPairs'][0]['GroupId'] = sg03.id
 
     sg01.authorize_ingress(IpPermissions=ip_permissions)
-    sg01.ip_permissions.should.have.length_of(2)
+    sg01.ip_permissions.should.have.length_of(3)
-    for ip_permission in ip_permissions:
+    for ip_permission in expected_ip_permissions:
         sg01.ip_permissions.should.contain(ip_permission)
 
     sg01.revoke_ingress(IpPermissions=ip_permissions)
     sg01.ip_permissions.should.be.empty
-    for ip_permission in ip_permissions:
+    for ip_permission in expected_ip_permissions:
         sg01.ip_permissions.shouldnt.contain(ip_permission)
 
     sg01.authorize_egress(IpPermissions=ip_permissions)
-    sg01.ip_permissions_egress.should.have.length_of(3)
+    sg01.ip_permissions_egress.should.have.length_of(4)
-    for ip_permission in ip_permissions:
+    for ip_permission in expected_ip_permissions:
         sg01.ip_permissions_egress.should.contain(ip_permission)
 
     sg01.revoke_egress(IpPermissions=ip_permissions)
     sg01.ip_permissions_egress.should.have.length_of(1)
-    for ip_permission in ip_permissions:
+    for ip_permission in expected_ip_permissions:
         sg01.ip_permissions_egress.shouldnt.contain(ip_permission)