Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

made the security group endpoints that authorize or revoke firewall rules

Browse Source 
to support batch rules (boto doesn't expose this, but botocore/boto3 does)
This commit is contained in:
Costy Petrisor 2016-10-20 17:25:54 +00:00
parent 6f622b9689
commit 8689b40d23
2 changed files with 74 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
moto/ec2/responses/security_groups.py
View File

@ -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

28
tests/test_ec2/test_security_groups.py
View File

@ -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)