def test_test_ip_set(self): """Test testing of IP in a given set""" # Disable W0212: Test access protected members of admin module. # pylint: disable=W0212 iptables._ipset.return_value = 42 res = iptables.test_ip_set('foo', '1.2.3.4') treadmill.iptables._ipset.assert_called_with( 'test', 'foo', '1.2.3.4', use_except=False, ) self.assertFalse(res) # Try with success now iptables._ipset.reset_mock() iptables._ipset.return_value = 0 res = iptables.test_ip_set('foo', '1.2.3.4') self.assertTrue(res)
def test_test_ip_set(self): """Test testing of IP in a given set""" # Disable protected-access: Test access protected members . # pylint: disable=protected-access iptables._ipset.return_value = (42, 'foo') res = iptables.test_ip_set('foo', '1.2.3.4') treadmill.iptables._ipset.assert_called_with( 'test', 'foo', '1.2.3.4', use_except=False, ) self.assertFalse(res) # Try with success now iptables._ipset.reset_mock() iptables._ipset.return_value = (0, 'bar') res = iptables.test_ip_set('foo', '1.2.3.4') self.assertTrue(res)
def _add_mark_rule(src_ip, environment): """Add an environment mark for all traffic coming from an IP. :param ``str`` src_ip: Source IP to be marked :param ``str`` environment: Environment to use for the mark """ assert environment in _SET_BY_ENVIRONMENT, \ 'Unknown environment: %r' % environment target_set = _SET_BY_ENVIRONMENT[environment] iptables.add_ip_set(target_set, src_ip) # Check that the IP is not marked in any other environment other_env_sets = { env_set for env_set in six.viewvalues(_SET_BY_ENVIRONMENT) if env_set != target_set } for other_set in other_env_sets: if iptables.test_ip_set(other_set, src_ip) is True: raise Exception('%r is already in %r' % (src_ip, other_set))