From 8689b40d237f05d8233fe22dfb69177dc93f606c Mon Sep 17 00:00:00 2001 From: Costy Petrisor Date: Thu, 20 Oct 2016 17:25:54 +0000 Subject: [PATCH] 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 | 73 +++++++++++++++++++------- tests/test_ec2/test_security_groups.py | 28 +++++++--- 2 files changed, 74 insertions(+), 27 deletions(-) diff --git a/moto/ec2/responses/security_groups.py b/moto/ec2/responses/security_groups.py index 275513082..3451dc1ef 100644 --- a/moto/ec2/responses/security_groups.py +++ b/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] - from_port = querystring.get('IpPermissions.1.FromPort', [None])[0] - to_port = querystring.get('IpPermissions.1.ToPort', [None])[0] - ip_ranges = [] + querytree = {} for key, value in querystring.items(): - if 'IpPermissions.1.IpRanges' in key: - ip_ranges.append(value[0]) + key_splitted = key.split('.') + key_splitted = [try_parse_int(e, e) for e in key_splitted] - source_groups = [] - source_group_ids = [] + d = querytree + 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(): - if 'IpPermissions.1.Groups.1.GroupId' in key: - source_group_ids.append(value[0]) - elif 'IpPermissions.1.Groups' in key: - source_groups.append(value[0]) + ip_permissions = querytree.get('IpPermissions') or {} + for ip_permission_idx in sorted(ip_permissions.keys()): + ip_permission = ip_permissions[ip_permission_idx] - 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)) - if not success: - return "Could not find a matching egress rule", dict(status=404) + for args in process_rules_from_querystring(self.querystring): + success = self.ec2_backend.revoke_security_group_egress(*args) + 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 diff --git a/tests/test_ec2/test_security_groups.py b/tests/test_ec2/test_security_groups.py index 358883902..3ccd1e4ab 100644 --- a/tests/test_ec2/test_security_groups.py +++ b/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) - for ip_permission in ip_permissions: + sg01.ip_permissions.should.have.length_of(3) + 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) - for ip_permission in ip_permissions: + sg01.ip_permissions_egress.should.have.length_of(4) + 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)