def _mock_iptables(version=osutil._IPTABLES_LOCKING_VERSION, destination='168.63.129.16'):
"""
Mock for the iptable commands used to set up the firewall.
Returns a patch of subprocess.Popen augmented with these properties:
* wait - True if the iptable commands use the -w option
* destination - The target IP address
* uid - The uid used for the -owner option
* command_calls - A list of the iptable commands executed by the mock (the --version and -L commands are omitted)
* set_command - By default all the mocked commands succeed and produce no output; this method can be used to override
the return value and output of these commands (or to add other commands)
"""
mocked_commands = {}
def set_command(command, output='', exit_code=0):
command_string = TestOSUtil._command_to_string(command)
mocked_commands[command_string] = (output.replace("'", "'\"'\"'"), exit_code)
return command_string
wait = FlexibleVersion(version) >= osutil._IPTABLES_LOCKING_VERSION
uid = 42
version_command = set_command(osutil._get_iptables_version_command(), output=str(version))
list_command = set_command(osutil._get_firewall_list_command(wait), output="Mock Output")
set_command(osutil._get_firewall_packets_command(wait))
set_command(osutil._get_firewall_drop_command(wait, "-C", destination))
set_command(osutil._get_firewall_drop_command(wait, "-A", destination))
set_command(osutil._get_firewall_accept_command(wait, "-A", destination, uid))
# the agent assumes the rules have been deleted when these commands return 1
set_command(osutil._get_firewall_delete_conntrack_accept_command(wait, destination), exit_code=1)
set_command(osutil._get_firewall_delete_owner_accept_command(wait, destination, uid), exit_code=1)
set_command(osutil._get_firewall_delete_conntrack_drop_command(wait, destination), exit_code=1)
command_calls = []
def mock_popen(command, *args, **kwargs):
command_string = TestOSUtil._command_to_string(command)
if command_string in mocked_commands:
if command_string != version_command and command_string != list_command:
command_calls.append(command_string)
output, exit_code = mocked_commands[command_string]
command = "echo '{0}' && exit {1}".format(output, exit_code)
kwargs["shell"] = True
return mock_popen.original(command, *args, **kwargs)
mock_popen.original = subprocess.Popen
with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as popen_patcher:
with patch('os.getuid', return_value=uid):
popen_patcher.wait = wait
popen_patcher.destination = destination
popen_patcher.uid = uid
popen_patcher.set_command = set_command
popen_patcher.command_calls = command_calls
yield popen_patcher
def test_enable_firewall_should_check_for_invalid_iptables_options(self):
osutil._enable_firewall = True
with TestOSUtil._mock_iptables() as mock_iptables:
# iptables uses the following exit codes
# 0 - correct function
# 1 - other errors
# 2 - errors which appear to be caused by invalid or abused command
# line parameters
drop_check_command = mock_iptables.set_command(osutil._get_firewall_drop_command(mock_iptables.wait, "-C", mock_iptables.destination), exit_code=2)
success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
delete_conntrack_accept_command = TestOSUtil._command_to_string(osutil._get_firewall_delete_conntrack_accept_command(mock_iptables.wait, mock_iptables.destination))
delete_owner_accept_command = TestOSUtil._command_to_string(osutil._get_firewall_delete_owner_accept_command(mock_iptables.wait, mock_iptables.destination, mock_iptables.uid))
delete_conntrack_drop_command = TestOSUtil._command_to_string(osutil._get_firewall_delete_conntrack_drop_command(mock_iptables.wait, mock_iptables.destination))
self.assertFalse(success, "Enable firewall should have failed")
self.assertEqual(len(mock_iptables.command_calls), 4, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
self.assertEqual(mock_iptables.command_calls[0], drop_check_command, "The first command should check the drop rule: {0}".format(mock_iptables.command_calls[0]))
self.assertEqual(mock_iptables.command_calls[1], delete_conntrack_accept_command, "The second command should delete the conntrack accept rule: {0}".format(mock_iptables.command_calls[1]))
self.assertEqual(mock_iptables.command_calls[2], delete_owner_accept_command, "The third command should delete the owner accept rule: {0}".format(mock_iptables.command_calls[2]))
self.assertEqual(mock_iptables.command_calls[3], delete_conntrack_drop_command, "The fourth command should delete the conntrack accept rule : {0}".format(mock_iptables.command_calls[3]))
self.assertFalse(osutil._enable_firewall)
def test_enable_firewall_should_set_up_the_firewall(self):
osutil._enable_firewall = True
with TestOSUtil._mock_iptables() as mock_iptables:
# fail the rule check to force enable of the firewall
mock_iptables.set_command(osutil._get_firewall_drop_command(mock_iptables.wait, "-C", mock_iptables.destination), exit_code=1)
success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
drop_check_command = TestOSUtil._command_to_string(osutil._get_firewall_drop_command(mock_iptables.wait, "-C", mock_iptables.destination))
accept_command = TestOSUtil._command_to_string(osutil._get_firewall_accept_command(mock_iptables.wait, "-A", mock_iptables.destination, mock_iptables.uid))
drop_add_command = TestOSUtil._command_to_string(osutil._get_firewall_drop_command(mock_iptables.wait, "-A", mock_iptables.destination))
self.assertTrue(success, "Enabling the firewall was not successful")
self.assertEqual(len(mock_iptables.command_calls), 3, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
self.assertEqual(mock_iptables.command_calls[0], drop_check_command, "The first command should check the drop rule")
self.assertEqual(mock_iptables.command_calls[1], accept_command, "The second command should add the accept rule")
self.assertEqual(mock_iptables.command_calls[2], drop_add_command, "The third command should add the drop rule")
self.assertTrue(osutil._enable_firewall, "The firewall should not have been disabled")
def test_enable_firewall_should_not_set_firewall_if_the_drop_rule_exists(self):
osutil._enable_firewall = True
with TestOSUtil._mock_iptables() as mock_iptables:
drop_check_command = mock_iptables.set_command(osutil._get_firewall_drop_command(mock_iptables.wait, "-C", mock_iptables.destination), exit_code=0)
success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
self.assertTrue(success, "Enabling the firewall was not successful")
self.assertEqual(len(mock_iptables.command_calls), 1, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
self.assertEqual(mock_iptables.command_calls[0], drop_check_command, "Unexpected command: {0}".format(mock_iptables.command_calls[0]))
self.assertTrue(osutil._enable_firewall)
def test_enable_firewall_should_not_use_wait_when_iptables_does_not_support_it(self):
osutil._enable_firewall = True
with TestOSUtil._mock_iptables(version=osutil._IPTABLES_LOCKING_VERSION - 1) as mock_iptables:
# fail the rule check to force enable of the firewall
mock_iptables.set_command(osutil._get_firewall_drop_command(mock_iptables.wait, "-C", mock_iptables.destination), exit_code=1)
success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
self.assertTrue(success, "Enabling the firewall was not successful")
self.assertEqual(len(mock_iptables.command_calls), 3, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
for command in mock_iptables.command_calls:
self.assertNotIn("-w", command, "The -w option should have been used in {0}".format(command))
self.assertTrue(osutil._enable_firewall, "The firewall should not have been disabled")