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