Repository: Azure/WALinuxAgent
Pull request #1944 (2020-07-13 22:20:04)

Use run_command to execute iptables
iptables was using the deprecated run/run_get_output functions

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=h1) Report
> Merging [#1944](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=desc) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/5aaf8804790ba408119c4e591641e802ebba6d22&el=desc) will **increase** coverage by `0.05%`.
> The diff coverage is `84.21%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944/graphs/tree.svg?width=650&height=150&src=pr&token=ae90E2KYCw)](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop    #1944      +/-   ##
===========================================
+ Coverage    69.56%   69.61%   +0.05%     
===========================================
  Files           85       85              
  Lines        11907    11928      +21     
  Branches      1673     1667       -6     
===========================================
+ Hits          8283     8304      +21     
- Misses        3252     3255       +3     
+ Partials       372      369       -3     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/osutil/default.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvZGVmYXVsdC5weQ==) | `58.90% <84.21%> (+1.00%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=footer). Last update [5aaf880...5c602c5](https://codecov.io/gh/Azure/WALinuxAgent/pull/1944?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index f84f503d..3ef6be17 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -56,25 +56,50 @@ for all distros. Each concrete distro classes could overwrite default behavior
 if needed.
 """
 
-IPTABLES_VERSION_PATTERN = re.compile("^[^\d\.]*([\d\.]+).*$")
-IPTABLES_VERSION = "iptables --version"
-IPTABLES_LOCKING_VERSION = FlexibleVersion('1.4.21')
-
-FIREWALL_ACCEPT = "iptables {0} -t security -{1} OUTPUT -d {2} -p tcp -m owner --uid-owner {3} -j ACCEPT"
-# Note:
-# -- Initially "flight" the change to ACCEPT packets and develop a metric baseline
-#    A subsequent release will convert the ACCEPT to DROP
-# FIREWALL_DROP = "iptables {0} -t security -{1} OUTPUT -d {2} -p tcp -m conntrack --ctstate INVALID,NEW -j ACCEPT"
-FIREWALL_DROP = "iptables {0} -t security -{1} OUTPUT -d {2} -p tcp -m conntrack --ctstate INVALID,NEW -j DROP"
-FIREWALL_LIST = "iptables {0} -t security -L -nxv"
-FIREWALL_PACKETS = "iptables {0} -t security -L OUTPUT --zero OUTPUT -nxv"
-FIREWALL_FLUSH = "iptables {0} -t security --flush"
+_IPTABLES_VERSION_PATTERN = re.compile("^[^\d\.]*([\d\.]+).*$")
+_IPTABLES_LOCKING_VERSION = FlexibleVersion('1.4.21')
+
+
+def _add_wait(wait, command):
+    """
+    If 'wait' is True, adds the wait option (-w) to the given iptables command line
+    """
+    if wait:
+        command.insert(1, "-w")
+    return command
+
+
+def _get_iptables_version_command():
+    return ["iptables", "--version"]
+
+
+def _get_firewall_accept_command(wait, command, destination, owner_uid):
+    return _add_wait(wait, ["iptables", "-t", "security", command, "OUTPUT", "-d", destination, "-p", "tcp", "-m", "owner", "--uid-owner", owner_uid, "-j" "ACCEPT"])
+
+
+def _get_firewall_drop_command(wait, command, destination):
+    return _add_wait(wait, ["iptables", "-t", "security", command, "OUTPUT", "-d", destination, "-p",  "tcp",  "-m", "conntrack", "--ctstate", "INVALID,NEW", "-j", "DROP"])
+
+
+def _get_firewall_list_command(wait):
+    return _add_wait(wait, ["iptables", "-t", "security", "-L", "-nxv"])
+
+
+def _get_firewall_packets_command(wait):
+    return _add_wait(wait, ["iptables", "-t", "security", "-L", "OUTPUT", "--zero", "OUTPUT", "-nxv"])
+
 
 # Precisely delete the rules created by the agent.
 # this rule was used <= 2.2.25.  This rule helped to validate our change, and determine impact.
-FIREWALL_DELETE_CONNTRACK_ACCEPT = "iptables {0} -t security -D OUTPUT -d {1} -p tcp -m conntrack --ctstate INVALID,NEW -j ACCEPT"
-FIREWALL_DELETE_OWNER_ACCEPT = "iptables {0} -t security -D OUTPUT -d {1} -p tcp -m owner --uid-owner {2} -j ACCEPT"
-FIREWALL_DELETE_CONNTRACK_DROP = "iptables {0} -t security -D OUTPUT -d {1} -p tcp -m conntrack --ctstate INVALID,NEW -j DROP"
+def _get_firewall_delete_conntrack_accept_command(wait, destination):
+    return _add_wait(wait, ["iptables", "-t", "security", "-D", "OUTPUT", "-d",  destination, "-p", "tcp", "-m", "conntrack", "--ctstate", "INVALID,NEW", "-j", "ACCEPT"])
+
+
+def _get_firewall_delete_owner_accept_command(wait, destination, owner_uid):
+    return _add_wait(wait, ["iptables", "-t", "security", "-D", "OUTPUT", "-d", destination, "-p", "tcp", "-m", "owner", "--uid-owner", owner_uid, "-j", "ACCEPT"])
+
+def _get_firewall_delete_conntrack_drop_command(wait, destination):
+    return _add_wait(wait, ["iptables", "-t", "security", "-D", "OUTPUT", "-d", destination, "-p", "tcp", "-m", "conntrack", "--ctstate", "INVALID,NEW", "-j", "DROP"])
 
 PACKET_PATTERN = "^\s*(\d+)\s+(\d+)\s+DROP\s+.*{0}[^\d]*$"
 ALL_CPUS_REGEX = re.compile('^cpu .*')
@@ -119,21 +144,23 @@ class DefaultOSUtil(object):
         try:
             wait = self.get_firewall_will_wait()
 
-            rc, output = shellutil.run_get_output(FIREWALL_PACKETS.format(wait), log_cmd=False, expected_errors=[3])
-            if rc == 3:
-                # Transient error  that we ignore.  This code fires every loop
-                # of the daemon (60m), so we will get the value eventually.
-                return 0
-
-            if rc != 0:
+            try:
+                output = shellutil.run_command(_get_firewall_packets_command(wait))
+
+                pattern = re.compile(PACKET_PATTERN.format(dst_ip))
+                for line in output.split('\n'):
+                    m = pattern.match(line)
+                    if m is not None:
+                        return int(m.group(1))
+
+            except Exception as e:
+                if isinstance(e, CommandError) and e.returncode == 3:
+                    # Transient error  that we ignore.  This code fires every loop
+                    # of the daemon (60m), so we will get the value eventually.
+                    return 0
+                logger.warn("Failed to get firewall packets: {0}", ustr(e))
                 return -1
 
-            pattern = re.compile(PACKET_PATTERN.format(dst_ip))
-            for line in output.split('\n'):
-                m = pattern.match(line)
-                if m is not None:
-                    return int(m.group(1))
-            
             return 0
 
         except Exception as e:
@@ -144,20 +171,21 @@ class DefaultOSUtil(object):
 
     def get_firewall_will_wait(self):
         # Determine if iptables will serialize access
-        rc, output = shellutil.run_get_output(IPTABLES_VERSION)
-        if rc != 0:
-            msg = "Unable to determine version of iptables"
+        try:
+            output = shellutil.run_command(_get_iptables_version_command())
+        except Exception as e:
+            msg = "Unable to determine version of iptables: {0}".format(ustr(e))
             logger.warn(msg)
             raise Exception(msg)
 
-        m = IPTABLES_VERSION_PATTERN.match(output)
+        m = _IPTABLES_VERSION_PATTERN.match(output)
         if m is None:
-            msg = "iptables did not return version information"
+            msg = "iptables did not return version information: {0}".format(output)
             logger.warn(msg)
             raise Exception(msg)
 
         wait = "-w" \
-                if FlexibleVersion(m.group(1)) >= IPTABLES_LOCKING_VERSION \
+                if FlexibleVersion(m.group(1)) >= _IPTABLES_LOCKING_VERSION \
                 else ""
         return wait
 
@@ -167,11 +195,13 @@ class DefaultOSUtil(object):
         code is non-zero or the limit has been reached.
         """
         for i in range(1, 100):
-            rc = shellutil.run(rule, chk_err=False)
-            if rc == 1:
-                return
-            elif rc == 2:
-                raise Exception("invalid firewall deletion rule '{0}'".format(rule))
+            try:
+                rc = shellutil.run_command(rule)
+            except CommandError as e:
+                if e.returncode == 1:
+                    return
+                if e.returncode == 2:
+                    raise Exception("invalid firewall deletion rule '{0}'".format(rule))
 
     def remove_firewall(self, dst_ip=None, uid=None):
         # If a previous attempt failed, do not retry
@@ -189,10 +219,9 @@ class DefaultOSUtil(object):
 
             # This rule was <= 2.2.25 only, and may still exist on some VMs.  Until 2.2.25
             # has aged out, keep this cleanup in place.
-            self._delete_rule(FIREWALL_DELETE_CONNTRACK_ACCEPT.format(wait, dst_ip))
-
-            self._delete_rule(FIREWALL_DELETE_OWNER_ACCEPT.format(wait, dst_ip, uid))
-            self._delete_rule(FIREWALL_DELETE_CONNTRACK_DROP.format(wait, dst_ip))
+            self._delete_rule(_get_firewall_delete_conntrack_accept_command(wait, dst_ip))
+            self._delete_rule(_get_firewall_delete_owner_accept_command(wait, dst_ip, uid))
+            self._delete_rule(_get_firewall_delete_conntrack_drop_command(wait, dst_ip))
 
             return True
 
@@ -203,55 +232,53 @@ class DefaultOSUtil(object):
                         "{0}".format(ustr(e)))
             return False
 
-    def enable_firewall(self, dst_ip=None, uid=None):
+    def enable_firewall(self, dst_ip, uid):
         # If a previous attempt failed, do not retry
         global _enable_firewall
         if not _enable_firewall:
             return False
 
         try:
-            if dst_ip is None or uid is None:
-                msg = "Missing arguments to enable_firewall"
-                logger.warn(msg)
-                raise Exception(msg)
-
             wait = self.get_firewall_will_wait()
 
             # If the DROP rule exists, make no changes
-            drop_rule = FIREWALL_DROP.format(wait, "C", dst_ip)
-            rc = shellutil.run(drop_rule, chk_err=False)
-            if rc == 0:
+            firewall_established = False
+            try:
+                drop_rule = _get_firewall_drop_command(wait, "-C", dst_ip)
+                shellutil.run_command(drop_rule)
                 logger.verbose("Firewall appears established")
                 return True
-            elif rc == 2:
-                self.remove_firewall(dst_ip, uid)
-                msg = "please upgrade iptables to a version that supports the -C option"
-                logger.warn(msg)
-                raise Exception(msg)
+            except CommandError as e:
+                if e.returncode == 2:
+                    self.remove_firewall(dst_ip, uid)
+                    msg = "please upgrade iptables to a version that supports the -C option"
+                    logger.warn(msg)
+                    raise Exception(msg)
 
             # Otherwise, append both rules
-            accept_rule = FIREWALL_ACCEPT.format(wait, "A", dst_ip, uid)
-            drop_rule = FIREWALL_DROP.format(wait, "A", dst_ip)
-
-            if shellutil.run(accept_rule) != 0:
-                msg = "Unable to add ACCEPT firewall rule '{0}'".format(
-                    accept_rule)
+            try:
+                accept_rule = _get_firewall_accept_command(wait, "-A", dst_ip, uid)
+                shellutil.run_command(accept_rule)
+            except Exception as e:
+                msg = "Unable to add ACCEPT firewall rule '{0}' - {1}".format(accept_rule, ustr(e))
                 logger.warn(msg)
                 raise Exception(msg)
 
-            if shellutil.run(drop_rule) != 0:
-                msg = "Unable to add DROP firewall rule '{0}'".format(
-                    drop_rule)
+            try:
+                drop_rule = _get_firewall_drop_command(wait, "-A", dst_ip)
+                shellutil.run_command(drop_rule)
+            except Exception as e:
+                msg = "Unable to add DROP firewall rule '{0}' - {1}".format(drop_rule, ustr(e))
                 logger.warn(msg)
                 raise Exception(msg)
 
             logger.info("Successfully added Azure fabric firewall rules")
 
-            rc, output = shellutil.run_get_output(FIREWALL_LIST.format(wait))
-            if rc == 0:
+            try:
+                output = shellutil.run_command(_get_firewall_list_command(wait))
                 logger.info("Firewall rules:\n{0}".format(output))
-            else:
-                logger.warn("Listing firewall rules failed: {0}".format(output))
+            except Exception as e:
+                logger.warn("Listing firewall rules failed: {0}".format(ustr(e)))
 
             return True
 


Test patch:
diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py
index f2dd87d1..100ff128 100644
--- a/tests/common/osutil/test_default.py
+++ b/tests/common/osutil/test_default.py
@@ -14,13 +14,14 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
-
-import socket
+import contextlib
 import glob
 import mock
-import traceback
 import os
+import socket
+import subprocess
 import tempfile
+import traceback
 import unittest
 
 import azurelinuxagent.common.osutil.default as osutil
@@ -31,6 +32,7 @@ from azurelinuxagent.common.exception import OSUtilError
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.utils import fileutil
+from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from tests.tools import AgentTestCase, call, patch, open_patch, load_data, \
     running_under_travis, skip_if_predicate_true
 
@@ -42,7 +44,6 @@ def fake_is_loopback(_, iface):
 
 
 class TestOSUtil(AgentTestCase):
-
     def setUp(self):
         AgentTestCase.setUp(self)
 
@@ -633,277 +634,254 @@ Match host 192.168.1.2\n\
         print("WRITING TO {0}".format(waagent_sudoers))
         self.assertEqual(1, count)
 
+    @staticmethod
+    def _command_to_string(command):
+        return " ".join([str(c) for c in command]) if isinstance(command, list) else command
+
+    @staticmethod
+    @contextlib.contextmanager
+    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_get_firewall_dropped_packets_returns_zero_if_firewall_disabled(self):
         osutil._enable_firewall = False
         util = osutil.DefaultOSUtil()
 
         self.assertEqual(0, util.get_firewall_dropped_packets("not used"))
 
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    def test_get_firewall_dropped_packets_returns_negative_if_error(self, mock_output):
+    def test_get_firewall_dropped_packets_returns_negative_if_error(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        mock_output.side_effect = [
-            (0, "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)),
-            (1, "not used")]
-        self.assertEqual(-1, util.get_firewall_dropped_packets("not used"))
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            mock_iptables.set_command(osutil._get_firewall_packets_command(mock_iptables.wait), exit_code=1)
 
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    def test_get_firewall_dropped_packets_returns_negative_if_exception(self, mock_output):
-        osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
+            self.assertEqual(-1, osutil.DefaultOSUtil().get_firewall_dropped_packets())
 
-        mock_output.side_effect = [
-            (0, "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)),
-            (1, Exception)]
-        self.assertEqual(-1, util.get_firewall_dropped_packets("not used"))
-
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    def test_get_firewall_dropped_packets_transient_error_ignored(self, mock_output):
+    def test_get_firewall_dropped_packets_should_ignore_transient_errors(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        mock_output.side_effect = [
-            (0, "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)),
-            (3, "can't initialize iptables table `security': iptables who? (do you need to insmod?)")]
-        self.assertEqual(0, util.get_firewall_dropped_packets("not used"))
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            mock_iptables.set_command(osutil._get_firewall_packets_command(mock_iptables.wait), exit_code=3, output="can't initialize iptables table `security': iptables who? (do you need to insmod?)")
 
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    def test_get_firewall_dropped_packets(self, mock_output):
+            self.assertEqual(0, osutil.DefaultOSUtil().get_firewall_dropped_packets())
+
+    def test_get_firewall_dropped_packets(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        mock_output.side_effect = [
-            (0, "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)),
-            (0,
-'''
+        destination = '168.63.129.16'
+
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            mock_iptables.set_command(osutil._get_firewall_packets_command(mock_iptables.wait), output='''
 
 Chain OUTPUT (policy ACCEPT 104 packets, 43628 bytes)
     pkts      bytes target     prot opt in     out     source               destination
        0        0 ACCEPT     tcp  --  any    any     anywhere             168.63.129.16        owner UID match daemon
       32     1920 DROP       tcp  --  any    any     anywhere             168.63.129.16
 
-''')]
-        dst = '168.63.129.16'
+''')
+            self.assertEqual(32, osutil.DefaultOSUtil().get_firewall_dropped_packets(destination))
 
-        self.assertEqual(32, util.get_firewall_dropped_packets(dst))
 
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall(self, mock_run, mock_output, mock_uid):
+    def test_enable_firewall_should_set_up_the_firewall(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        mock_run.side_effect = [1, 0, 0]
-        mock_output.side_effect = [(0, version), (0, "Output")]
-        self.assertTrue(util.enable_firewall(dst_ip=dst, uid=uid))
-
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DROP.format(wait, "C", dst), chk_err=False),
-            call(osutil.FIREWALL_ACCEPT.format(wait, "A", dst, uid)),
-            call(osutil.FIREWALL_DROP.format(wait, "A", dst))
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION),
-            call(osutil.FIREWALL_LIST.format(wait))
-        ])
-        self.assertTrue(osutil._enable_firewall)
-
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall_no_wait(self, mock_run, mock_output, mock_uid):
-        osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
+        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)
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION-1)
-        wait = ""
-
-        mock_run.side_effect = [1, 0, 0]
-        mock_output.side_effect = [(0, version), (0, "Output")]
-        self.assertTrue(util.enable_firewall(dst_ip=dst, uid=uid))
-
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DROP.format(wait, "C", dst), chk_err=False),
-            call(osutil.FIREWALL_ACCEPT.format(wait, "A", dst, uid)),
-            call(osutil.FIREWALL_DROP.format(wait, "A", dst))
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION),
-            call(osutil.FIREWALL_LIST.format(wait))
-        ])
-        self.assertTrue(osutil._enable_firewall)
-
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall_skips_if_drop_exists(self, mock_run, mock_output, mock_uid):
+            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_use_wait_when_iptables_does_not_support_it(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        mock_run.side_effect = [0, 0, 0]
-        mock_output.return_value = (0, version)
-        self.assertTrue(util.enable_firewall(dst_ip=dst, uid=uid))
-
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DROP.format(wait, "C", dst), chk_err=False),
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION)
-        ])
-        self.assertTrue(osutil._enable_firewall)
-
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall_ignores_exceptions(self, mock_run, mock_output, mock_uid):
+        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")
+
+    def test_enable_firewall_should_not_set_firewall_if_the_drop_rule_exists(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        mock_run.side_effect = [1, Exception]
-        mock_output.return_value = (0, version)
-        self.assertFalse(util.enable_firewall(dst_ip=dst, uid=uid))
-
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DROP.format(wait, "C", dst), chk_err=False),
-            call(osutil.FIREWALL_ACCEPT.format(wait, "A", dst, uid))
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION)
-        ])
-        self.assertFalse(osutil._enable_firewall)
+        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)
 
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall_checks_for_invalid_iptables_options(self, mock_run, mock_output):
+            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_check_for_invalid_iptables_options(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        # 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
-        mock_run.side_effect = [2]
-        mock_output.return_value = (0, version)
-
-        self.assertFalse(util.enable_firewall(dst_ip='1.2.3.4', uid=42))
-        self.assertFalse(osutil._enable_firewall)
-
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DROP.format(wait, "C", dst), chk_err=False),
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION)
-        ])
-
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_enable_firewall_skips_if_disabled(self, mock_run, mock_output, mock_uid):
+        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_skips_if_disabled(self):
         osutil._enable_firewall = False
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
 
-        mock_run.side_effect = [1, 0, 0]
-        mock_output.side_effect = [(0, version), (0, "Output")]
-        self.assertFalse(util.enable_firewall(dst_ip=dst, uid=uid))
+            self.assertFalse(success, "The firewall should not have been disabled")
+            self.assertEqual(len(mock_iptables.command_calls), 0, "iptables should not have been invoked: [{0}]". format(mock_iptables.command_calls))
 
-        mock_run.assert_not_called()
-        mock_output.assert_not_called()
-        mock_uid.assert_not_called()
-        self.assertFalse(osutil._enable_firewall)
+            self.assertFalse(osutil._enable_firewall)
 
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_remove_firewall(self, mock_run, mock_output, mock_uid):
+    def test_remove_firewall(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
 
-        dst = '1.2.3.4'
-        uid = 42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        mock_run.side_effect = [0, 1, 0, 1, 0, 1]
-        mock_output.side_effect = [(0, version), (0, "Output")]
-        self.assertTrue(util.remove_firewall(dst, uid))
-
-        mock_run.assert_has_calls([
-            # delete rules < 2.2.26
-            call(osutil.FIREWALL_DELETE_CONNTRACK_ACCEPT.format(wait, dst), chk_err=False),
-            call(osutil.FIREWALL_DELETE_CONNTRACK_ACCEPT.format(wait, dst), chk_err=False),
-            call(osutil.FIREWALL_DELETE_OWNER_ACCEPT.format(wait, dst, uid), chk_err=False),
-            call(osutil.FIREWALL_DELETE_OWNER_ACCEPT.format(wait, dst, uid), chk_err=False),
-
-            # delete rules >= 2.2.26
-            call(osutil.FIREWALL_DELETE_CONNTRACK_DROP.format(wait, dst), chk_err=False),
-            call(osutil.FIREWALL_DELETE_CONNTRACK_DROP.format(wait, dst), chk_err=False),
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION)
-        ])
-        self.assertTrue(osutil._enable_firewall)
-
-    @patch('os.getuid', return_value=42)
-    @patch('azurelinuxagent.common.utils.shellutil.run_get_output')
-    @patch('azurelinuxagent.common.utils.shellutil.run')
-    def test_remove_firewall_does_not_repeat(self, mock_run, mock_output, _):
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            delete_commands = {}
+
+            def mock_popen(command, *args, **kwargs):
+                command_string = TestOSUtil._command_to_string(command)
+                if " -D " in command_string:
+                    # The agent invokes the delete commands continuously until they return 1 to indicate the rules has been removed
+                    # The mock returns 0 (success) the first time it is invoked and 1 (rule does not exist) thereafter
+                    if command_string not in delete_commands:
+                        exit_code = 0
+                        delete_commands[command_string] = 1
+                    else:
+                        exit_code = 1
+                        delete_commands[command_string] += 1
+
+                    command = "echo '' && exit {0}".format(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):
+                success = osutil.DefaultOSUtil().remove_firewall(mock_iptables.destination, 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.assertTrue(success, "Removing the firewall should have succeeded")
+                self.assertEquals(len(delete_commands), 3, "Expected 3 delete commands: [{0}]".format(delete_commands))
+                # delete rules < 2.2.26
+                self.assertIn(delete_conntrack_accept_command, delete_commands, "The delete conntrack accept command was not executed")
+                self.assertEqual(delete_commands[delete_conntrack_accept_command], 2, "The delete conntrack accept command should have been executed twice")
+                self.assertIn(delete_owner_accept_command, delete_commands, "The delete owner accept command was not executed")
+                self.assertEqual(delete_commands[delete_owner_accept_command], 2, "The delete owner accept command should have been executed twice")
+                # delete rules >= 2.2.26
+                self.assertIn(delete_conntrack_drop_command, delete_commands, "The delete conntrack drop command was not executed")
+                self.assertEqual(delete_commands[delete_conntrack_drop_command], 2, "The delete conntrack drop command should have been executed twice")
+
+                self.assertTrue(osutil._enable_firewall)
+
+    def test_remove_firewall_should_not_retry_invalid_rule(self):
         osutil._enable_firewall = True
-        util = osutil.DefaultOSUtil()
-
-        dst_ip='1.2.3.4'
-        uid=42
-        version = "iptables v{0}".format(osutil.IPTABLES_LOCKING_VERSION)
-        wait = "-w"
-
-        mock_run.side_effect = [2]
-        mock_output.side_effect = [(0, version), (1, "Output")]
-        self.assertFalse(util.remove_firewall(dst_ip, uid))
 
-        mock_run.assert_has_calls([
-            call(osutil.FIREWALL_DELETE_CONNTRACK_ACCEPT.format(wait, dst_ip), chk_err=False),
-        ])
-        mock_output.assert_has_calls([
-            call(osutil.IPTABLES_VERSION)
-        ])
-        self.assertFalse(osutil._enable_firewall)
+        with TestOSUtil._mock_iptables() as mock_iptables:
+            command = osutil._get_firewall_delete_conntrack_accept_command(mock_iptables.wait, mock_iptables.destination)
+            # Note that the command is actually a valid rule, but we use the mock to report it as invalid (exit code 2)
+            delete_conntrack_accept_command = mock_iptables.set_command(command, exit_code=2)
 
-        self.assertTrue(mock_run.call_count == 1)
-        self.assertTrue(mock_output.call_count == 1)
+            success = osutil.DefaultOSUtil().remove_firewall(mock_iptables.destination, mock_iptables.uid)
 
-        self.assertFalse(util.remove_firewall())
-        self.assertFalse(util.remove_firewall())
+            self.assertFalse(success, "Removing the firewall should not have succeeded")
+            self.assertEqual(len(mock_iptables.command_calls), 1, "Expected a single call to iptables: [{0}]". format(mock_iptables.command_calls))
+            self.assertEqual(mock_iptables.command_calls[0], delete_conntrack_accept_command, "Expected call to delete conntrack accept command: {0}".format(mock_iptables.command_calls[0]))
 
-        self.assertTrue(mock_run.call_count == 1)
-        self.assertTrue(mock_output.call_count == 1)
+            self.assertFalse(osutil._enable_firewall)
 
     @skip_if_predicate_true(running_under_travis, "The ip command isn't available in Travis")
     def test_get_nic_state(self):