From fc3e9907e75299e3fb55b93aa32503cf159133b3 Mon Sep 17 00:00:00 2001 From: cameron-r Date: Wed, 5 Nov 2014 17:07:29 -0600 Subject: [PATCH] Ed & Cameron | Remove rules from security groups on EC2 in GroupRuleRefresher --- group_rule_refresher.py | 27 ++++++++++++++++++++----- tests/unit/test_group_rule_refresher.py | 26 ++++++++++++++++++------ 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/group_rule_refresher.py b/group_rule_refresher.py index e8dafbb..b473731 100644 --- a/group_rule_refresher.py +++ b/group_rule_refresher.py @@ -6,13 +6,30 @@ class GroupRuleRefresher: self.ec2_rule_service = ec2_rule_service def refresh(self, group_name): - openstack_rules = self.openstack_rule_service.get_rules_for_group(group_name) - ec2_rules = self.ec2_rule_service.get_rules_for_group(group_name) + openstack_rules = self.openstack_rule_service.get_rules_for_group(group_name) + ec2_rules = self.ec2_rule_service.get_rules_for_group(group_name) - for rule in openstack_rules - ec2_rules: - self._create_rule_on_ec2(group_name, rule) + self._add_rules_to_ec2(ec2_rules, group_name, openstack_rules) + self._remove_rules_from_ec2(ec2_rules, group_name, openstack_rules) - def _create_rule_on_ec2(self, group_name, rule): + def _add_rules_to_ec2(self, ec2_rules, group_name, openstack_rules): + for rule in openstack_rules - ec2_rules: + self._add_rule_on_ec2(group_name, rule) + + def _remove_rules_from_ec2(self, ec2_rules, group_name, openstack_rules): + for rule in ec2_rules - openstack_rules: + self._remove_rule_from_ec2(group_name, rule) + + def _remove_rule_from_ec2(self, group_name, rule): + self.ec2_conn.revoke_security_group( + group_name=group_name, + ip_protocol=rule.ip_protocol, + from_port=rule.from_port, + to_port=rule.to_port, + cidr_ip=rule.ip_range + ) + + def _add_rule_on_ec2(self, group_name, rule): self.ec2_conn.authorize_security_group( group_name=group_name, ip_protocol=rule.ip_protocol, diff --git a/tests/unit/test_group_rule_refresher.py b/tests/unit/test_group_rule_refresher.py index 1166a71..9d4aabe 100644 --- a/tests/unit/test_group_rule_refresher.py +++ b/tests/unit/test_group_rule_refresher.py @@ -13,7 +13,7 @@ OTHER_GROUP_NAME = "otherSecGroup" class TestGroupRuleRefresher(unittest.TestCase): def setUp(self): - self.new_rule = Rule('hjkl', 7, 8, '9.9.9.9/99') + self.rule = Rule('hjkl', 7, 8, '9.9.9.9/99') self.openstack_instance = Mock() self.ec2_connection = Mock(EC2Connection) @@ -27,15 +27,29 @@ class TestGroupRuleRefresher(unittest.TestCase): ) def test_should_add_rule_to_ec2_security_group_when_rule_associated_with_group_on_openstack(self): - self.openstack_rule_service.get_rules_for_group.return_value = set([self.new_rule]) + self.openstack_rule_service.get_rules_for_group.return_value = set([self.rule]) self.ec2_rule_service.get_rules_for_group.return_value = set() self.group_rule_refresher.refresh(GROUP_NAME) self.ec2_connection.authorize_security_group.assert_called_once_with( group_name=GROUP_NAME, - ip_protocol=self.new_rule.ip_protocol, - from_port=self.new_rule.from_port, - to_port=self.new_rule.to_port, - cidr_ip=self.new_rule.ip_range + ip_protocol=self.rule.ip_protocol, + from_port=self.rule.from_port, + to_port=self.rule.to_port, + cidr_ip=self.rule.ip_range + ) + + def test_should_remove_rule_from_ec2_security_group_when_rule_not_associated_with_group_on_openstack(self): + self.openstack_rule_service.get_rules_for_group.return_value = set() + self.ec2_rule_service.get_rules_for_group.return_value = set([self.rule]) + + self.group_rule_refresher.refresh(GROUP_NAME) + + self.ec2_connection.revoke_security_group.assert_called_once_with( + group_name=GROUP_NAME, + ip_protocol=self.rule.ip_protocol, + from_port=self.rule.from_port, + to_port=self.rule.to_port, + cidr_ip=self.rule.ip_range ) \ No newline at end of file