Repository: Azure/WALinuxAgent
Pull request #2480 (2022-01-20 22:28:09)

Add check for all Firewall rules periodically
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2480](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (35cf310) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/ec0609d673fa31ee4ad803885a02dbb55411c77c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (ec0609d) will **decrease** coverage by `0.06%`.
> The diff coverage is `68.62%`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480/graphs/tree.svg?width=650&height=150&src=pr&token=ae90E2KYCw&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)

```diff
@@             Coverage Diff             @@
##           develop    #2480      +/-   ##
===========================================
- Coverage    71.96%   71.90%   -0.07%     
===========================================
  Files          101      101              
  Lines        15063    15086      +23     
  Branches      2393     2401       +8     
===========================================
+ Hits         10840    10847       +7     
- Misses        3737     3750      +13     
- Partials       486      489       +3     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/ga/env.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2Vudi5weQ==) | `60.54% <0.00%> (-3.03%)` | :arrow_down: |
| [azurelinuxagent/common/osutil/default.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvZGVmYXVsdC5weQ==) | `59.55% <35.71%> (-1.00%)` | :arrow_down: |
| [azurelinuxagent/common/utils/networkutil.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi91dGlscy9uZXR3b3JrdXRpbC5weQ==) | `97.05% <100.00%> (+0.32%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Last update [ec0609d...35cf310](https://codecov.io/gh/Azure/WALinuxAgent/pull/2480?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).


Patch:
diff --git a/azurelinuxagent/common/conf.py b/azurelinuxagent/common/conf.py
index 65d07030..a666e53e 100644
--- a/azurelinuxagent/common/conf.py
+++ b/azurelinuxagent/common/conf.py
@@ -188,7 +188,8 @@ __INTEGER_OPTIONS__ = {
     "Debug.AgentCpuQuota": 75,
     "Debug.EtpCollectionPeriod": 300,
     "Debug.AutoUpdateHotfixFrequency": 14400,
-    "Debug.AutoUpdateNormalFrequency": 86400
+    "Debug.AutoUpdateNormalFrequency": 86400,
+    "Debug.FirewallRulesLogPeriod": 86400
 }
 
 
@@ -605,3 +606,12 @@ def get_enable_ga_versioning(conf=__conf__):
     NOTE: This option is experimental and may be removed in later versions of the Agent.
     """
     return conf.get_switch("Debug.EnableGAVersioning", False)
+
+
+def get_firewall_rules_log_period(conf=__conf__):
+    """
+    Determine the frequency to perform the periodic operation of logging firewall rules.
+
+    NOTE: This option is experimental and may be removed in later versions of the Agent.
+    """
+    return conf.get_int("Debug.FirewallRulesLogPeriod", 86400)
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index 1c604216..f18ff904 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -75,18 +75,6 @@ def get_iptables_version_command():
     return ["iptables", "--version"]
 
 
-def get_accept_tcp_rule(wait, command, destination):
-    return AddFirewallRules.get_accept_tcp_rule(command, destination, wait=wait)
-
-
-def get_firewall_accept_command(wait, command, destination, owner_uid):
-    return AddFirewallRules.get_iptables_accept_command(wait, command, destination, owner_uid)
-
-
-def get_firewall_drop_command(wait, command, destination):
-    return AddFirewallRules.get_iptables_drop_command(wait, command, destination)
-
-
 def get_firewall_list_command(wait):
     return _add_wait(wait, ["iptables", "-t", "security", "-L", "-nxv"])
 
@@ -228,15 +216,13 @@ class DefaultOSUtil(object):
                 if e.returncode == 2:
                     raise Exception("invalid firewall deletion rule '{0}'".format(rule))
 
-    def remove_firewall(self, dst_ip, uid):
+    def remove_firewall(self, dst_ip, uid, wait):
         # If a previous attempt failed, do not retry
         global _enable_firewall  # pylint: disable=W0603
         if not _enable_firewall:
             return False
 
         try:
-            wait = self.get_firewall_will_wait()
-
             # 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(get_firewall_delete_conntrack_accept_command(wait, dst_ip))
@@ -269,50 +255,56 @@ class DefaultOSUtil(object):
                 "Unable to remove legacy firewall rule, won't try removing it again. Error: {0}".format(ustr(error)))
 
     def enable_firewall(self, dst_ip, uid):
+        """
+        It checks if every iptable rule exists and add them if not present. It returns a tuple(enable firewall success status, update rules flag)
+        enable firewall success status: Returns True if every firewall rule exists otherwise False
+        update rules flag: Returns True if rules are updated otherwise False
+        """
+        # This is to send telemetry when iptable rules updated
+        is_firewall_rules_updated = False
         # If a previous attempt failed, do not retry
         global _enable_firewall  # pylint: disable=W0603
         if not _enable_firewall:
-            return False
+            return False, is_firewall_rules_updated
 
         try:
             wait = self.get_firewall_will_wait()
 
-            # If the DROP rule exists, make no changes
+            # check every iptable rule and delete others if any rule is missing
+            #   and append every iptable rule to the end of the chain.
             try:
-                drop_rule = get_firewall_drop_command(wait, AddFirewallRules.CHECK_COMMAND, dst_ip)
-                shellutil.run_command(drop_rule)
-                logger.verbose("Firewall appears established")
-                return True
+                if not AddFirewallRules.verify_iptables_rules_exist(wait, dst_ip, uid):
+                    self.remove_firewall(dst_ip, uid, wait)
+                    AddFirewallRules.add_iptables_rules(wait, dst_ip, uid)
+                    is_firewall_rules_updated = True
             except CommandError as e:
                 if e.returncode == 2:
-                    self.remove_firewall(dst_ip, uid)
+                    self.remove_firewall(dst_ip, uid, wait)
                     msg = "please upgrade iptables to a version that supports the -C option"
                     logger.warn(msg)
-                    raise Exception(msg)
-
-            # Otherwise, append all rules
-            try:
-                AddFirewallRules.add_iptables_rules(wait, dst_ip, uid)
+                    raise
             except Exception as error:
                 logger.warn(ustr(error))
                 raise
 
-            logger.info("Successfully added Azure fabric firewall rules")
-
-            try:
-                output = shellutil.run_command(get_firewall_list_command(wait))
-                logger.info("Firewall rules:\n{0}".format(output))
-            except Exception as e:
-                logger.warn("Listing firewall rules failed: {0}".format(ustr(e)))
-
-            return True
+            return True, is_firewall_rules_updated
 
         except Exception as e:
             _enable_firewall = False
             logger.info("Unable to establish firewall -- "
                         "no further attempts will be made: "
                         "{0}".format(ustr(e)))
-            return False
+            return False, is_firewall_rules_updated
+
+    def get_firewall_list(self, wait=None):
+        try:
+            if wait is None:
+                wait = self.get_firewall_will_wait()
+            output = shellutil.run_command(get_firewall_list_command(wait))
+            return output
+        except Exception as e:
+            logger.warn("Listing firewall rules failed: {0}".format(ustr(e)))
+            return ""
 
     @staticmethod
     def _correct_instance_id(instance_id):
diff --git a/azurelinuxagent/common/osutil/nsbsd.py b/azurelinuxagent/common/osutil/nsbsd.py
index 31b8a343..016f506f 100644
--- a/azurelinuxagent/common/osutil/nsbsd.py
+++ b/azurelinuxagent/common/osutil/nsbsd.py
@@ -24,7 +24,6 @@ from azurelinuxagent.common.osutil.freebsd import FreeBSDOSUtil
 
 
 class NSBSDOSUtil(FreeBSDOSUtil):
-
     resolver = None
 
     def __init__(self):
@@ -45,7 +44,7 @@ class NSBSDOSUtil(FreeBSDOSUtil):
             for server in output.split("\n"):
                 if server == '':
                     break
-                server = server[:-1] # remove last '='
+                server = server[:-1]  # remove last '='
                 cmd = "grep '{}' /etc/hosts".format(server) + " | awk '{print $1}'"
                 ret, ip = shellutil.run_get_output(cmd)
                 servers.append(ip)
@@ -53,7 +52,8 @@ class NSBSDOSUtil(FreeBSDOSUtil):
             dns.resolver.override_system_resolver(self.resolver)
 
     def set_hostname(self, hostname):
-        self._run_command_without_raising(['/usr/Firewall/sbin/setconf', '/usr/Firewall/System/global', 'SystemName', hostname])
+        self._run_command_without_raising(
+            ['/usr/Firewall/sbin/setconf', '/usr/Firewall/System/global', 'SystemName', hostname])
         self._run_command_without_raising(["/usr/Firewall/sbin/enlog"])
         self._run_command_without_raising(["/usr/Firewall/sbin/enproxy", "-u"])
         self._run_command_without_raising(["/usr/Firewall/sbin/ensl", "-u"])
@@ -94,16 +94,15 @@ class NSBSDOSUtil(FreeBSDOSUtil):
         commands = [['setconf', '/usr/Firewall/ConfigFiles/webadmin', 'ACL', 'any'], ['ensl']]
         self._run_multiple_commands_without_raising(commands, log_error=False, continue_on_error=False)
 
-
     def deploy_ssh_pubkey(self, username, pubkey):
         """
         Deploy authorized_key
         """
         path, thumbprint, value = pubkey  # pylint: disable=W0612
 
-        #overide parameters
+        # overide parameters
         super(NSBSDOSUtil, self).deploy_ssh_pubkey('admin',
-            ["/usr/Firewall/.ssh/authorized_keys", thumbprint, value]) 
+                                                   ["/usr/Firewall/.ssh/authorized_keys", thumbprint, value])
 
     def del_root_password(self):
         logger.warn("Root password deletion disabled")
@@ -130,7 +129,7 @@ class NSBSDOSUtil(FreeBSDOSUtil):
         shellutil.run("ennetwork", chk_err=False)
 
     def set_dhcp_hostname(self, hostname):
-        #already done by the dhcp client
+        # already done by the dhcp client
         pass
 
     def get_firewall_dropped_packets(self, dst_ip=None):
@@ -145,10 +144,14 @@ class NSBSDOSUtil(FreeBSDOSUtil):
         # disable iptables methods
         return
 
-    def remove_firewall(self, dst_ip=None, uid=None):
+    def remove_firewall(self, dst_ip=None, uid=None, wait=""):
         # disable iptables methods
         return True
 
     def enable_firewall(self, dst_ip=None, uid=None):
         # disable iptables methods
-        return True
+        return True, True
+
+    def get_firewall_list(self, wait=""):
+        # disable iptables methods
+        return ""
diff --git a/azurelinuxagent/common/persist_firewall_rules.py b/azurelinuxagent/common/persist_firewall_rules.py
index bd143454..55b19c41 100644
--- a/azurelinuxagent/common/persist_firewall_rules.py
+++ b/azurelinuxagent/common/persist_firewall_rules.py
@@ -143,12 +143,21 @@ if __name__ == '__main__':
 
         return True
 
+    def __remove_firewalld_rules(self):
+        try:
+            AddFirewallRules.remove_firewalld_rules(self._dst_ip, self._uid)
+        except Exception as error:
+            logger.warn(
+                "failed to remove rule using firewalld.service: {0}".format(ustr(error)))
+
     def _setup_permanent_firewalld_rules(self):
         if self.__verify_firewall_rules_enabled():
             logger.info("Firewall rules already set. No change needed.")
             return
 
         logger.info("Firewall rules not added yet, adding them now using firewalld.service")
+        # Remove first if partial list present
+        self.__remove_firewalld_rules()
         # Add rules if not already set
         AddFirewallRules.add_firewalld_rules(self._dst_ip, self._uid)
         logger.info("Successfully added the firewall commands using firewalld.service")
@@ -327,4 +336,3 @@ if __name__ == '__main__':
         logger.info(
             "Unit file version matches with expected version: {0}, not overwriting unit file".format(unit_file_version))
         return False
-
diff --git a/azurelinuxagent/common/protocol/metadata_server_migration_util.py b/azurelinuxagent/common/protocol/metadata_server_migration_util.py
index ab3a6bb5..8a3b9b1c 100644
--- a/azurelinuxagent/common/protocol/metadata_server_migration_util.py
+++ b/azurelinuxagent/common/protocol/metadata_server_migration_util.py
@@ -60,9 +60,9 @@ def _reset_firewall_rules(osutil):
     Removes MetadataServer firewall rule so IMDS can be used. Enables
     WireServer firewall rule based on if firewall is configured to be on.
     """
-    osutil.remove_firewall(dst_ip=_KNOWN_METADATASERVER_IP, uid=os.getuid())
+    osutil.remove_firewall(dst_ip=_KNOWN_METADATASERVER_IP, uid=os.getuid(), wait=osutil.get_firewall_will_wait())
     if conf.enable_firewall():
-        success = osutil.enable_firewall(dst_ip=KNOWN_WIRESERVER_IP, uid=os.getuid())
+        success, _ = osutil.enable_firewall(dst_ip=KNOWN_WIRESERVER_IP, uid=os.getuid())
         add_event(
             AGENT_NAME,
             version=CURRENT_VERSION,
diff --git a/azurelinuxagent/common/utils/networkutil.py b/azurelinuxagent/common/utils/networkutil.py
index 8908d841..160f1751 100644
--- a/azurelinuxagent/common/utils/networkutil.py
+++ b/azurelinuxagent/common/utils/networkutil.py
@@ -26,6 +26,7 @@ class RouteEntry(object):
     Represents a single route. The destination, gateway, and mask members are hex representations of the IPv4 address in
     network byte order.
     """
+
     def __init__(self, interface, destination, gateway, mask, flags, metric):
         self.interface = interface
         self.destination = destination
@@ -40,7 +41,7 @@ class RouteEntry(object):
             raise Exception("String to dotted quad conversion must be 8 characters")
         octets = []
         for idx in range(6, -2, -2):
-            octets.append(str(int(value[idx:idx+2], 16)))
+            octets.append(str(int(value[idx:idx + 2], 16)))
         return ".".join(octets)
 
     def destination_quad(self):
@@ -55,7 +56,7 @@ class RouteEntry(object):
     def to_json(self):
         f = '{{"Iface": "{0}", "Destination": "{1}", "Gateway": "{2}", "Mask": "{3}", "Flags": "{4:#06x}", "Metric": "{5}"}}'
         return f.format(self.interface, self.destination_quad(), self.gateway_quad(), self.mask_quad(),
-                       self.flags, self.metric)
+                        self.flags, self.metric)
 
     def __str__(self):
         f = "Iface: {0}\tDestination: {1}\tGateway: {2}\tMask: {3}\tFlags: {4:#06x}\tMetric: {5}"
@@ -63,7 +64,7 @@ class RouteEntry(object):
                         self.flags, self.metric)
 
     def __repr__(self):
-        return 'RouteEntry("{0}", "{1}", "{2}", "{3}", "{4:#04x}", "{5}")'\
+        return 'RouteEntry("{0}", "{1}", "{2}", "{3}", "{4:#04x}", "{5}")' \
             .format(self.interface, self.destination, self.gateway, self.mask, self.flags, self.metric)
 
 
@@ -112,6 +113,11 @@ class FirewallCmdDirectCommands(object):
     # checks if the firewalld rule is present or not
     QueryPassThrough = "--query-passthrough"
 
+    # firewall-cmd --permanent --direct --remove-passthrough ipv4 -t security -A OUTPUT -d 168.63.129.16 -p tcp -m owner --uid-owner 0 -j ACCEPT
+    # success
+    # remove the firewalld rule
+    RemovePassThrough = "--remove-passthrough"
+
 
 class AddFirewallRules(object):
     """
@@ -136,16 +142,10 @@ class AddFirewallRules(object):
     CHECK_COMMAND = "-C"
 
     @staticmethod
-    def _add_wait(wait, command):
+    def __get_iptables_base_command(wait=""):
         """
         If 'wait' is True, adds the wait option (-w) to the given iptables command line
         """
-        if wait:
-            command.append("-w")
-        return command
-
-    @staticmethod
-    def __get_iptables_base_command(wait=""):
         if wait != "":
             return ["iptables", "-w"]
         return ["iptables"]
@@ -157,56 +157,35 @@ class AddFirewallRules(object):
 
     @staticmethod
     def __get_common_command_params(command, destination):
-        return ["-t", "security", command, "OUTPUT", "-d", destination, "-p", "tcp", "-m"]
-
-    @staticmethod
-    def __get_common_accept_command_params(wait, command, destination, uid):
-        cmd = AddFirewallRules.__get_common_command_params(command, destination)
-        cmd.extend(["owner", "--uid-owner", str(uid), "-j", "ACCEPT"])
-        return AddFirewallRules._add_wait(wait, cmd)
+        return ["-t", "security", command, "OUTPUT", "-d", destination, "-p", "tcp"]
 
     @staticmethod
-    def __get_common_drop_command_params(wait, command, destination):
-        cmd = AddFirewallRules.__get_common_command_params(command, destination)
-        cmd.extend(["conntrack", "--ctstate", "INVALID,NEW", "-j", "DROP"])
-        return AddFirewallRules._add_wait(wait, cmd)
-
-    @staticmethod
-    def get_accept_tcp_rule(command, destination, firewalld_command="", wait=""):
-        # This rule allows DNS TCP request to wireserver ip for non root users
-
+    def __get_firewall_base_command(command, destination, firewalld_command="", wait=""):
+        # Firewalld.service fails if we set `-w` in the iptables command, so not adding it at all for firewalld commands
         if firewalld_command != "":
             cmd = AddFirewallRules.__get_firewalld_base_command(firewalld_command)
         else:
             cmd = AddFirewallRules.__get_iptables_base_command(wait)
-        cmd = cmd + ['-t', 'security', command, 'OUTPUT', '-d', destination, '-p', 'tcp', '--destination-port', '53', '-j', 'ACCEPT']
+        cmd.extend(AddFirewallRules.__get_common_command_params(command, destination))
         return cmd
 
     @staticmethod
-    def get_iptables_accept_command(wait, command, destination, owner_uid):
-        cmd = AddFirewallRules.__get_iptables_base_command()
-        cmd.extend(AddFirewallRules.__get_common_accept_command_params(wait, command, destination, owner_uid))
-        return cmd
-
-    @staticmethod
-    def get_iptables_drop_command(wait, command, destination):
-        cmd = AddFirewallRules.__get_iptables_base_command()
-        cmd.extend(AddFirewallRules.__get_common_drop_command_params(wait, command, destination))
+    def get_accept_tcp_rule(command, destination, firewalld_command="", wait=""):
+        # This rule allows DNS TCP request to wireserver ip for non root users
+        cmd = AddFirewallRules.__get_firewall_base_command(command, destination, firewalld_command, wait)
+        cmd.extend(['--destination-port', '53', '-j', 'ACCEPT'])
         return cmd
 
     @staticmethod
-    def get_firewalld_accept_command(command, destination, uid, wait=""):
-        cmd = AddFirewallRules.__get_firewalld_base_command(command)
-        cmd.extend(
-            AddFirewallRules.__get_common_accept_command_params(wait, AddFirewallRules.APPEND_COMMAND, destination,
-                                                                uid))
+    def get_wire_root_accept_rule(command, destination, owner_uid, firewalld_command="", wait=""):
+        cmd = AddFirewallRules.__get_firewall_base_command(command, destination, firewalld_command, wait)
+        cmd.extend(["-m", "owner", "--uid-owner", str(owner_uid), "-j", "ACCEPT"])
         return cmd
 
     @staticmethod
-    def get_firewalld_drop_command(command, destination, wait=""):
-        cmd = AddFirewallRules.__get_firewalld_base_command(command)
-        cmd.extend(
-            AddFirewallRules.__get_common_drop_command_params(wait, AddFirewallRules.APPEND_COMMAND, destination))
+    def get_wire_non_root_drop_rule(command, destination, firewalld_command="", wait=""):
+        cmd = AddFirewallRules.__get_firewall_base_command(command, destination, firewalld_command, wait)
+        cmd.extend(["-m", "conntrack", "--ctstate", "INVALID,NEW", "-j", "DROP"])
         return cmd
 
     @staticmethod
@@ -226,43 +205,60 @@ class AddFirewallRules(object):
             raise Exception(msg)
 
     @staticmethod
-    def add_iptables_rules(wait, dst_ip, uid):
-        # The order in which the below rules are added matters for the ip table rules to work as expected
-        AddFirewallRules.__raise_if_empty(dst_ip, "Destination IP")
-        AddFirewallRules.__raise_if_empty(uid, "User ID")
-
-        accept_tcp_rule = AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.APPEND_COMMAND, dst_ip, wait=wait)
-        AddFirewallRules.__execute_cmd(accept_tcp_rule)
+    def __execute_check_command(cmd):
+        # Here we primarily check if an  iptable rule exist. True if it exits , false if not
+        try:
+            shellutil.run_command(cmd)
+            return True
+        except CommandError as err:
+            # return code 1 is expected while using the check command. Raise if encounter any other return code
+            if err.returncode != 1:
+                raise
+        return False
 
-        accept_rule = AddFirewallRules.get_iptables_accept_command(wait, AddFirewallRules.APPEND_COMMAND, dst_ip, uid)
-        AddFirewallRules.__execute_cmd(accept_rule)
+    @staticmethod
+    def verify_iptables_rules_exist(wait, dst_ip, uid):
+        check_cmd_tcp_rule = AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, dst_ip, wait=wait)
+        check_cmd_accept_rule = AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, dst_ip, uid,
+                                                                           wait=wait)
+        check_cmd_drop_rule = AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, dst_ip, wait=wait)
 
-        drop_rule = AddFirewallRules.get_iptables_drop_command(wait, AddFirewallRules.APPEND_COMMAND, dst_ip)
-        AddFirewallRules.__execute_cmd(drop_rule)
+        return AddFirewallRules.__execute_check_command(check_cmd_tcp_rule) and AddFirewallRules.__execute_check_command(check_cmd_accept_rule) \
+               and AddFirewallRules.__execute_check_command(check_cmd_drop_rule)
 
     @staticmethod
-    def __execute_firewalld_commands(command, dst_ip, uid):
+    def __execute_firewall_commands(dst_ip, uid, command=APPEND_COMMAND, firewalld_command="", wait=""):
+        # The order in which the below rules are added matters for the ip table rules to work as expected
+
         AddFirewallRules.__raise_if_empty(dst_ip, "Destination IP")
         AddFirewallRules.__raise_if_empty(uid, "User ID")
 
-        # Firewalld.service fails if we set `-w` in the iptables command, so not adding it at all for firewalld commands
-
-        accept_tcp_rule = AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, dst_ip, firewalld_command=command)
+        accept_tcp_rule = AddFirewallRules.get_accept_tcp_rule(command, dst_ip,
+                                                               firewalld_command=firewalld_command, wait=wait)
         AddFirewallRules.__execute_cmd(accept_tcp_rule)
 
-        accept_cmd = AddFirewallRules.get_firewalld_accept_command(command, dst_ip, uid)
+        accept_cmd = AddFirewallRules.get_wire_root_accept_rule(command, dst_ip, uid,
+                                                                firewalld_command=firewalld_command, wait=wait)
         AddFirewallRules.__execute_cmd(accept_cmd)
 
-        drop_cmd = AddFirewallRules.get_firewalld_drop_command(command, dst_ip)
+        drop_cmd = AddFirewallRules.get_wire_non_root_drop_rule(command, dst_ip, firewalld_command=firewalld_command, wait=wait)
         AddFirewallRules.__execute_cmd(drop_cmd)
 
+    @staticmethod
+    def add_iptables_rules(wait, dst_ip, uid):
+        AddFirewallRules.__execute_firewall_commands(dst_ip, uid, command=AddFirewallRules.APPEND_COMMAND, wait=wait)
+
     @staticmethod
     def add_firewalld_rules(dst_ip, uid):
         # Firewalld.service fails if we set `-w` in the iptables command, so not adding it at all for firewalld commands
         # Firewalld.service with the "--permanent --passthrough" parameter ensures that a firewall rule is set only once even if command is executed multiple times
 
-        AddFirewallRules.__execute_firewalld_commands(FirewallCmdDirectCommands.PassThrough, dst_ip, uid)
+        AddFirewallRules.__execute_firewall_commands(dst_ip, uid, firewalld_command=FirewallCmdDirectCommands.PassThrough)
 
     @staticmethod
     def check_firewalld_rule_applied(dst_ip, uid):
-        AddFirewallRules.__execute_firewalld_commands(FirewallCmdDirectCommands.QueryPassThrough, dst_ip, uid)
+        AddFirewallRules.__execute_firewall_commands(dst_ip, uid, firewalld_command=FirewallCmdDirectCommands.QueryPassThrough)
+
+    @staticmethod
+    def remove_firewalld_rules(dst_ip, uid):
+        AddFirewallRules.__execute_firewall_commands(dst_ip, uid, firewalld_command=FirewallCmdDirectCommands.RemovePassThrough)
diff --git a/azurelinuxagent/ga/env.py b/azurelinuxagent/ga/env.py
index e8780cbe..771f4765 100644
--- a/azurelinuxagent/ga/env.py
+++ b/azurelinuxagent/ga/env.py
@@ -26,7 +26,7 @@ import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.logger as logger
 
 from azurelinuxagent.common.dhcp import get_dhcp_handler
-from azurelinuxagent.common.event import add_periodic, WALAEventOperation
+from azurelinuxagent.common.event import add_periodic, WALAEventOperation, add_event
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.interfaces import ThreadHandlerInterface
 from azurelinuxagent.common.osutil import get_osutil
@@ -131,7 +131,13 @@ class EnableFirewall(PeriodicOperation):
             self._osutil.remove_legacy_firewall_rule(dst_ip=self._protocol.get_endpoint())
             self._try_remove_legacy_firewall_rule = True
 
-        success = self._osutil.enable_firewall(dst_ip=self._protocol.get_endpoint(), uid=os.getuid())
+        success, is_firewall_rules_updated = self._osutil.enable_firewall(dst_ip=self._protocol.get_endpoint(),
+                                                                          uid=os.getuid())
+
+        if is_firewall_rules_updated:
+            msg = "Successfully added Azure fabric firewall rules. Current Firewall rules:\n{0}".format(self._osutil.get_firewall_list())
+            logger.info(msg)
+            add_event(AGENT_NAME, version=CURRENT_VERSION, op=WALAEventOperation.Firewall, message=msg, log_event=False)
 
         add_periodic(
             logger.EVERY_HOUR,
@@ -142,6 +148,21 @@ class EnableFirewall(PeriodicOperation):
             log_event=False)
 
 
+class LogFirewallRules(PeriodicOperation):
+    """
+    Log firewall rules state once a day.
+    Goal is to capture the firewall state when the agent service startup,
+    in addition to add more debug data and would be more useful long term.
+    """
+    def __init__(self, osutil):
+        super(LogFirewallRules, self).__init__(conf.get_firewall_rules_log_period())
+        self._osutil = osutil
+
+    def _operation(self):
+        # Log firewall rules state once a day
+        logger.info("Current Firewall rules:\n{0}".format(self._osutil.get_firewall_list()))
+
+
 class SetRootDeviceScsiTimeout(PeriodicOperation):
     def __init__(self, osutil):
         super(SetRootDeviceScsiTimeout, self).__init__(conf.get_root_device_scsi_timeout_period())
@@ -223,6 +244,7 @@ class EnvHandler(ThreadHandlerInterface):
 
             if conf.enable_firewall():
                 periodic_operations.append(EnableFirewall(osutil, protocol))
+                periodic_operations.append(LogFirewallRules(osutil))
             if conf.get_root_device_scsi_timeout() is not None:
                 periodic_operations.append(SetRootDeviceScsiTimeout(osutil))
             if conf.get_monitor_hostname():
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 3c7b6f0a..88525752 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -43,8 +43,6 @@ from azurelinuxagent.common.event import add_event, initialize_event_logger_vmin
 from azurelinuxagent.common.exception import ResourceGoneError, UpdateError, ExitException, AgentUpgradeExitException
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil, systemd
-from azurelinuxagent.common.osutil.default import get_firewall_drop_command, \
-    get_accept_tcp_rule
 from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
 from azurelinuxagent.common.protocol.hostplugin import HostPluginProtocol
 from azurelinuxagent.common.protocol.restapi import VMAgentUpdateStatus, VMAgentUpdateStatuses, ExtHandlerPackageList, \
@@ -1254,7 +1252,7 @@ class UpdateHandler(object):
             wait = self.osutil.get_firewall_will_wait()
 
             # "-C" checks if the iptable rule is available in the chain. It throws an exception with return code 1 if the ip table rule doesnt exist
-            drop_rule = get_firewall_drop_command(wait, AddFirewallRules.CHECK_COMMAND, dst_ip)
+            drop_rule = AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, dst_ip, wait=wait)
             if not _execute_run_command(drop_rule):
                 # DROP command doesn't exist indicates then none of the firewall rules are set yet
                 # exiting here as the environment thread will set up all firewall rules
@@ -1262,12 +1260,12 @@ class UpdateHandler(object):
                 return
             else:
                 # DROP rule exists in the ip table chain. Hence checking if the DNS TCP to wireserver rule exists. If not we add it.
-                accept_tcp_rule = get_accept_tcp_rule(wait, AddFirewallRules.CHECK_COMMAND, dst_ip)
+                accept_tcp_rule = AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, dst_ip, wait=wait)
                 if not _execute_run_command(accept_tcp_rule):
                     try:
                         logger.info(
                             "Firewall rule to allow DNS TCP request to wireserver for a non root user unavailable. Setting it now.")
-                        accept_tcp_rule = get_accept_tcp_rule(wait, AddFirewallRules.INSERT_COMMAND, dst_ip)
+                        accept_tcp_rule = AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, dst_ip, wait=wait)
                         shellutil.run_command(accept_tcp_rule)
                         logger.info(
                             "Succesfully added firewall rule to allow non root users to do a DNS TCP request to wireserver")


Test patch:
diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py
index e34e3d4f..0f7d445f 100644
--- a/tests/common/osutil/test_default.py
+++ b/tests/common/osutil/test_default.py
@@ -661,12 +661,13 @@ Match host 192.168.1.2\n\
         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, AddFirewallRules.CHECK_COMMAND, destination))
-        set_command(osutil.get_firewall_drop_command(wait, AddFirewallRules.APPEND_COMMAND, destination))
-        set_command(osutil.get_firewall_accept_command(wait, AddFirewallRules.APPEND_COMMAND, destination, uid))
-        set_command(osutil.get_accept_tcp_rule(wait, AddFirewallRules.APPEND_COMMAND, destination))
-        set_command(osutil.get_accept_tcp_rule(wait, AddFirewallRules.INSERT_COMMAND, destination))
-        set_command(osutil.get_accept_tcp_rule(wait, AddFirewallRules.CHECK_COMMAND, destination))
+        set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, destination, wait=wait))
+        set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.APPEND_COMMAND, destination, wait=wait))
+        set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, destination, uid, wait=wait))
+        set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.APPEND_COMMAND, destination, uid, wait=wait))
+        set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.APPEND_COMMAND, destination, wait=wait))
+        set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, destination, wait=wait))
+        set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, destination, wait=wait))
         # 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_delete_accept_tcp_rule(wait, destination), exit_code=1)
@@ -737,24 +738,52 @@ Match host 192.168.1.2\n\
         with TestOSUtil._mock_iptables() as mock_iptables:
             with patch.object(osutil, '_enable_firewall', True):
                 # fail the rule check to force enable of the firewall
-                mock_iptables.set_command(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, 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, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                accept_tcp_append_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.APPEND_COMMAND, mock_iptables.destination))
-                accept_command = TestOSUtil._command_to_string(osutil.get_firewall_accept_command(mock_iptables.wait, AddFirewallRules.APPEND_COMMAND, mock_iptables.destination, mock_iptables.uid))
-                drop_add_command = TestOSUtil._command_to_string(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.APPEND_COMMAND, mock_iptables.destination))
+                mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                               wait=mock_iptables.wait), exit_code=0)
+                mock_iptables.set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination, mock_iptables.uid,
+                                                                                     wait=mock_iptables.wait), exit_code=0)
+                mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                       wait=mock_iptables.wait), exit_code=1)
+
+                success, _ = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
+
+                tcp_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                               wait=mock_iptables.wait))
+                accept_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination, mock_iptables.uid,
+                                                                                                                wait=mock_iptables.wait))
+                drop_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                                wait=mock_iptables.wait))
+                delete_conntrack_accept_command = TestOSUtil._command_to_string(osutil.get_firewall_delete_conntrack_accept_command(mock_iptables.wait, mock_iptables.destination))
+                delete_accept_tcp_rule = TestOSUtil._command_to_string(osutil.get_delete_accept_tcp_rule(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))
+                accept_tcp_command = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.APPEND_COMMAND, mock_iptables.destination, wait=mock_iptables.wait))
+                accept_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.APPEND_COMMAND, mock_iptables.destination, mock_iptables.uid, wait=mock_iptables.wait))
+                drop_add_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.APPEND_COMMAND, mock_iptables.destination, wait=mock_iptables.wait))
 
                 self.assertTrue(success, "Enabling the firewall was not successful")
-                # Exactly 4 calls have to be made.
-                # First is the DROP rule check which was mocked to fail, Second, third and forth to Append the three IPTable rules
-                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")
-                self.assertEqual(mock_iptables.command_calls[1], accept_tcp_append_rule,
-                                "The second command should add the dns tcp accept rule")
-                self.assertEqual(mock_iptables.command_calls[2], accept_command, "The third command should add the accept rule")
-                self.assertEqual(mock_iptables.command_calls[3], drop_add_command, "The fourth command should add the drop rule")
+                # Exactly 10 calls have to be made.
+                # First is the check rule check which was mocked to fail, and delete call and then append calls
+                self.assertEqual(len(mock_iptables.command_calls), 10, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
+                self.assertEqual(mock_iptables.command_calls[0], tcp_check_command, "The first command should check the tcp rule")
+                self.assertEqual(mock_iptables.command_calls[1], accept_check_command, "The second command should check the accept rule")
+                self.assertEqual(mock_iptables.command_calls[2], drop_check_command, "The third command should check the drop rule")
+                self.assertEqual(mock_iptables.command_calls[3], delete_conntrack_accept_command,
+                                 "The fourth command should delete the conntrack accept rule: {0}".format(
+                                     mock_iptables.command_calls[3]))
+                self.assertEqual(mock_iptables.command_calls[4], delete_accept_tcp_rule,
+                                 "The fifth command should delete the dns tcp accept rule: {0}".format(
+                                     mock_iptables.command_calls[4]))
+                self.assertEqual(mock_iptables.command_calls[5], delete_owner_accept_command,
+                                 "The sixth command should delete the owner accept rule: {0}".format(
+                                     mock_iptables.command_calls[5]))
+                self.assertEqual(mock_iptables.command_calls[6], delete_conntrack_drop_command,
+                                 "The seventh command should delete the conntrack accept rule : {0}".format(
+                                     mock_iptables.command_calls[6]))
+                self.assertEqual(mock_iptables.command_calls[7], accept_tcp_command,
+                                "The eighth command should add the dns tcp accept rule")
+                self.assertEqual(mock_iptables.command_calls[8], accept_command, "The ninth command should add the accept rule")
+                self.assertEqual(mock_iptables.command_calls[9], drop_add_command, "The tenth command should add the drop rule")
 
                 self.assertTrue(osutil._enable_firewall, "The firewall should not have been disabled")
 
@@ -762,34 +791,44 @@ Match host 192.168.1.2\n\
         with TestOSUtil._mock_iptables(version=osutil._IPTABLES_LOCKING_VERSION - 1) as mock_iptables:
             with patch.object(osutil, '_enable_firewall', True):
                 # fail the rule check to force enable of the firewall
-                mock_iptables.set_command(
-                    osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND,
-                                                     mock_iptables.destination), exit_code=1)
+                mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                               wait=mock_iptables.wait), exit_code=1)
+                mock_iptables.set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination, mock_iptables.uid,
+                                                                                     wait=mock_iptables.wait), exit_code=1)
+                mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                       wait=mock_iptables.wait), exit_code=1)
 
-                success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
+                success, _ = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
 
                 self.assertTrue(success, "Enabling the firewall was not successful")
-                # Exactly 4 calls have to be made.
-                # First is the DROP rule check which was mocked to fail, Second, third and forth
-                # to Append the three IPTable rules
-                self.assertEqual(len(mock_iptables.command_calls), 4,
+                # Exactly 8 calls have to be made.
+                # First check rule, delete 4 rules,
+                # and Append the IPTable 3 rules.
+                self.assertEqual(len(mock_iptables.command_calls), 8,
                                  "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):
+    def test_enable_firewall_should_not_set_firewall_if_the_all_the_rules_exists(self):
 
         with TestOSUtil._mock_iptables() as mock_iptables:
             with patch.object(osutil, '_enable_firewall', True):
-                drop_check_command = mock_iptables.set_command(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=0)
+                tcp_check_command = mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                               wait=mock_iptables.wait), exit_code=0)
+                accept_check_command = mock_iptables.set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination, mock_iptables.uid,
+                                                                                                            wait=mock_iptables.wait), exit_code=0)
+                drop_check_command = mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                            wait=mock_iptables.wait), exit_code=0)
 
-                success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
+                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.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], tcp_check_command, "Unexpected command: {0}".format(mock_iptables.command_calls[0]))
+                self.assertEqual(mock_iptables.command_calls[1], accept_check_command, "Unexpected command: {0}".format(mock_iptables.command_calls[1]))
+                self.assertEqual(mock_iptables.command_calls[2], drop_check_command, "Unexpected command: {0}".format(mock_iptables.command_calls[2]))
 
                 self.assertTrue(osutil._enable_firewall)
 
@@ -802,9 +841,14 @@ Match host 192.168.1.2\n\
                 #  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, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=2)
+                tcp_check_command = mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                               wait=mock_iptables.wait), exit_code=0)
+                accept_check_command = mock_iptables.set_command(AddFirewallRules.get_wire_root_accept_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination, mock_iptables.uid,
+                                                                                                            wait=mock_iptables.wait), exit_code=0)
+                drop_check_command = mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                            wait=mock_iptables.wait), exit_code=2)
 
-                success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
+                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_accept_tcp_rule = TestOSUtil._command_to_string(osutil.get_delete_accept_tcp_rule(mock_iptables.wait, mock_iptables.destination))
@@ -812,14 +856,16 @@ Match host 192.168.1.2\n\
                 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), 5, "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_accept_tcp_rule,
-                                 "The third command should delete the dns tcp accept rule: {0}".format(
-                                     mock_iptables.command_calls[2]))
-                self.assertEqual(mock_iptables.command_calls[3], delete_owner_accept_command, "The forth command should delete the owner accept rule: {0}".format(mock_iptables.command_calls[3]))
-                self.assertEqual(mock_iptables.command_calls[4], delete_conntrack_drop_command, "The fifth command should delete the conntrack accept rule : {0}".format(mock_iptables.command_calls[4]))
+                self.assertEqual(len(mock_iptables.command_calls), 7, "Incorrect number of calls to iptables: [{0}]". format(mock_iptables.command_calls))
+                self.assertEqual(mock_iptables.command_calls[0], tcp_check_command, "The first command should check the tcp rule: {0}".format(mock_iptables.command_calls[0]))
+                self.assertEqual(mock_iptables.command_calls[1], accept_check_command, "The second command should check the accept rule: {0}".format(mock_iptables.command_calls[1]))
+                self.assertEqual(mock_iptables.command_calls[2], drop_check_command, "The third command should check the drop rule: {0}".format(mock_iptables.command_calls[2]))
+                self.assertEqual(mock_iptables.command_calls[3], delete_conntrack_accept_command, "The fourth command should delete the conntrack accept rule: {0}".format(mock_iptables.command_calls[3]))
+                self.assertEqual(mock_iptables.command_calls[4], delete_accept_tcp_rule,
+                                 "The fifth command should delete the dns tcp accept rule: {0}".format(
+                                     mock_iptables.command_calls[4]))
+                self.assertEqual(mock_iptables.command_calls[5], delete_owner_accept_command, "The sixth command should delete the owner accept rule: {0}".format(mock_iptables.command_calls[5]))
+                self.assertEqual(mock_iptables.command_calls[6], delete_conntrack_drop_command, "The seventh command should delete the conntrack accept rule : {0}".format(mock_iptables.command_calls[6]))
 
                 self.assertFalse(osutil._enable_firewall)
 
@@ -827,7 +873,7 @@ Match host 192.168.1.2\n\
 
         with TestOSUtil._mock_iptables() as mock_iptables:
             with patch.object(osutil, '_enable_firewall', False):
-                success = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.uid)
+                success, _ = osutil.DefaultOSUtil().enable_firewall(dst_ip=mock_iptables.destination, uid=mock_iptables.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))
@@ -858,7 +904,7 @@ Match host 192.168.1.2\n\
                 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)
+                    success = osutil.DefaultOSUtil().remove_firewall(mock_iptables.destination, mock_iptables.uid, mock_iptables.wait)
 
                     delete_conntrack_accept_command = TestOSUtil._command_to_string(osutil.get_firewall_delete_conntrack_accept_command(mock_iptables.wait, mock_iptables.destination))
                     delete_accept_tcp_rule = TestOSUtil._command_to_string(
@@ -889,7 +935,7 @@ Match host 192.168.1.2\n\
                 # 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)
 
-                success = osutil.DefaultOSUtil().remove_firewall(mock_iptables.destination, mock_iptables.uid)
+                success = osutil.DefaultOSUtil().remove_firewall(mock_iptables.destination, mock_iptables.uid, mock_iptables.wait)
 
                 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))
diff --git a/tests/common/test_persist_firewall_rules.py b/tests/common/test_persist_firewall_rules.py
index 29d845d2..307c8536 100644
--- a/tests/common/test_persist_firewall_rules.py
+++ b/tests/common/test_persist_firewall_rules.py
@@ -99,21 +99,25 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
 
     def __assert_firewall_called(self, cmd, validate_command_called=True):
         if validate_command_called:
-            self.assertIn(AddFirewallRules.get_firewalld_accept_command(command=cmd,
-                                                                        destination=self.__test_dst_ip,
-                                                                        uid=self.__test_uid),
+            self.assertIn(AddFirewallRules.get_wire_root_accept_rule(command=AddFirewallRules.APPEND_COMMAND,
+                                                                     destination=self.__test_dst_ip,
+                                                                     owner_uid=self.__test_uid,
+                                                                     firewalld_command=cmd),
                           self.__executed_commands, "Firewall {0} command not found".format(cmd))
-            self.assertIn(AddFirewallRules.get_firewalld_drop_command(command=cmd,
-                                                                      destination=self.__test_dst_ip),
+            self.assertIn(AddFirewallRules.get_wire_non_root_drop_rule(command=AddFirewallRules.APPEND_COMMAND,
+                                                                       destination=self.__test_dst_ip,
+                                                                       firewalld_command=cmd),
                           self.__executed_commands, "Firewall {0} command not found".format(cmd))
         else:
-            self.assertNotIn(AddFirewallRules.get_firewalld_accept_command(command=cmd,
-                                                                           destination=self.__test_dst_ip,
-                                                                           uid=self.__test_uid),
+            self.assertNotIn(AddFirewallRules.get_wire_root_accept_rule(command=AddFirewallRules.APPEND_COMMAND,
+                                                                        destination=self.__test_dst_ip,
+                                                                        owner_uid=self.__test_uid,
+                                                                        firewalld_command=cmd),
                              self.__executed_commands,
                              "Firewall {0} command found".format(cmd))
-            self.assertNotIn(AddFirewallRules.get_firewalld_drop_command(command=cmd,
-                                                                         destination=self.__test_dst_ip),
+            self.assertNotIn(AddFirewallRules.get_wire_non_root_drop_rule(command=AddFirewallRules.APPEND_COMMAND,
+                                                                          destination=self.__test_dst_ip,
+                                                                          firewalld_command=cmd),
                              self.__executed_commands, "Firewall {0} command found".format(cmd))
 
     def __assert_systemctl_called(self, cmd="enable", validate_command_called=True):
@@ -173,6 +177,22 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
             return True, ["echo", "enabled"]
         return False, []
 
+    @staticmethod
+    def __mock_firewalld_running_and_remove_not_successful(cmd):
+        if cmd == PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD:
+            return True, ["echo", "running"]
+        # This is to fail the check if firewalld-rules are already applied
+        cmds_to_fail = ["firewall-cmd", FirewallCmdDirectCommands.QueryPassThrough, "conntrack"]
+        if all(cmd_to_fail in cmd for cmd_to_fail in cmds_to_fail):
+            return True, ["exit", "1"]
+        # This is to fail the remove if firewalld-rules fails to remove rule
+        cmds_to_fail = ["firewall-cmd", FirewallCmdDirectCommands.RemovePassThrough, "conntrack"]
+        if all(cmd_to_fail in cmd for cmd_to_fail in cmds_to_fail):
+            return True, ["exit", "2"]
+        if "firewall-cmd" in cmd:
+            return True, ["echo", "enabled"]
+        return False, []
+
     def __setup_and_assert_network_service_setup_scenario(self, handler, mock_popen=None):
         mock_popen = TestPersistFirewallRulesHandler.__mock_network_setup_service_disabled if mock_popen is None else mock_popen
         self.__replace_popen_cmd = mock_popen
@@ -183,6 +203,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         self.__assert_systemctl_reloaded(validate_command_called=True)
         self.__assert_firewall_cmd_running_called(validate_command_called=True)
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=False)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=False)
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=False)
         self.assertTrue(os.path.exists(handler.get_service_file_path()), "Service unit file not found")
 
@@ -194,6 +215,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         # Assert we verified that rules were set using firewall-cmd
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=True)
         # Assert no commands for adding rules using firewall-cmd were called
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=False)
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=False)
         # Assert no commands for systemctl were called
         self.assertFalse(any("systemctl" in cmd for cmd in self.__executed_commands), "Systemctl shouldn't be called")
@@ -215,6 +237,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
             self.__assert_systemctl_reloaded(validate_command_called=False)
             self.__assert_firewall_cmd_running_called(validate_command_called=True)
             self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=False)
+            self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=False)
             self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=False)
             self.assertTrue(os.path.exists(handler.get_service_file_path()), "Service unit file not found")
 
@@ -280,6 +303,19 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
 
         self.__assert_firewall_cmd_running_called(validate_command_called=True)
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=True)
+        self.assertFalse(any("systemctl" in cmd for cmd in self.__executed_commands), "Systemctl shouldn't be called")
+
+    def test_it_should_add_firewalld_rules_if_remove_raises_exception(self):
+
+        self.__replace_popen_cmd = self.__mock_firewalld_running_and_remove_not_successful
+        with self._get_persist_firewall_rules_handler() as handler:
+            handler.setup()
+
+        self.__assert_firewall_cmd_running_called(validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=True)
         self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=True)
         self.assertFalse(any("systemctl" in cmd for cmd in self.__executed_commands), "Systemctl shouldn't be called")
 
@@ -356,6 +392,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
 
             self.__assert_firewall_cmd_running_called(validate_command_called=True)
             self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.QueryPassThrough, validate_command_called=True)
+            self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.RemovePassThrough, validate_command_called=True)
             self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=True)
             self.__assert_systemctl_called(cmd="is-enabled", validate_command_called=False)
             self.__assert_systemctl_called(cmd="enable", validate_command_called=False)
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 9cc5bdb1..b0f15b80 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -1668,14 +1668,19 @@ class TestUpdate(UpdateTestCase):
             with _get_update_handler(test_data=DATA_FILE) as (update_handler, _):
                 with patch.object(osutil, '_enable_firewall', True):
                     # drop rule is present
-                    mock_iptables.set_command(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=0)
+                    mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                           wait=mock_iptables.wait), exit_code=0)
                     # non root tcp iptable rule is absent
-                    mock_iptables.set_command(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=1)
+                    mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait), exit_code=1)
                     update_handler.run(debug=True)
 
-                    drop_check_command = TestOSUtil._command_to_string(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_check_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_insert_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.INSERT_COMMAND, mock_iptables.destination))
+                    drop_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                                    wait=mock_iptables.wait))
+                    accept_tcp_check_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
+                    accept_tcp_insert_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
 
                     # Filtering the mock iptable command calls with only the once related to this test.
                     filtered_mock_iptable_calls = [cmd for cmd in mock_iptables.command_calls if cmd in [drop_check_command, accept_tcp_check_rule, accept_tcp_insert_rule]]
@@ -1694,13 +1699,17 @@ class TestUpdate(UpdateTestCase):
             with _get_update_handler(test_data=DATA_FILE) as (update_handler, _):
                 with patch.object(osutil, '_enable_firewall', True):
                     # drop rule is not available
-                    mock_iptables.set_command(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=1)
+                    mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                           wait=mock_iptables.wait), exit_code=1)
 
                     update_handler.run(debug=True)
 
-                    drop_check_command = TestOSUtil._command_to_string(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_check_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_insert_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.INSERT_COMMAND, mock_iptables.destination))
+                    drop_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                                    wait=mock_iptables.wait))
+                    accept_tcp_check_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
+                    accept_tcp_insert_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
 
                     # Filtering the mock iptable command calls with only the once related to this test.
                     filtered_mock_iptable_calls = [cmd for cmd in mock_iptables.command_calls if cmd in [drop_check_command, accept_tcp_check_rule, accept_tcp_insert_rule]]
@@ -1715,15 +1724,20 @@ class TestUpdate(UpdateTestCase):
             with _get_update_handler(test_data=DATA_FILE) as (update_handler, _):
                 with patch.object(osutil, '_enable_firewall', True):
                     # drop rule is available
-                    mock_iptables.set_command(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=0)
+                    mock_iptables.set_command(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                           wait=mock_iptables.wait), exit_code=0)
                     # non root tcp iptable rule is available
-                    mock_iptables.set_command(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination), exit_code=0)
+                    mock_iptables.set_command(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait), exit_code=0)
 
                     update_handler.run(debug=True)
 
-                    drop_check_command = TestOSUtil._command_to_string(osutil.get_firewall_drop_command(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_check_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.CHECK_COMMAND, mock_iptables.destination))
-                    accept_tcp_insert_rule = TestOSUtil._command_to_string(osutil.get_accept_tcp_rule(mock_iptables.wait, AddFirewallRules.INSERT_COMMAND, mock_iptables.destination))
+                    drop_check_command = TestOSUtil._command_to_string(AddFirewallRules.get_wire_non_root_drop_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                                                                                    wait=mock_iptables.wait))
+                    accept_tcp_check_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.CHECK_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
+                    accept_tcp_insert_rule = TestOSUtil._command_to_string(AddFirewallRules.get_accept_tcp_rule(AddFirewallRules.INSERT_COMMAND, mock_iptables.destination,
+                                                             wait=mock_iptables.wait))
 
                     # Filtering the mock iptable command calls with only the once related to this test.
                     filtered_mock_iptable_calls = [cmd for cmd in mock_iptables.command_calls if cmd in [drop_check_command, accept_tcp_check_rule, accept_tcp_insert_rule]]
diff --git a/tests/protocol/test_metadata_server_migration_util.py b/tests/protocol/test_metadata_server_migration_util.py
index 3ac69030..3f5c2cf9 100644
--- a/tests/protocol/test_metadata_server_migration_util.py
+++ b/tests/protocol/test_metadata_server_migration_util.py
@@ -70,6 +70,7 @@ class TestMetadataServerMigrationUtil(AgentTestCase):
         fixed_uid = 0
         mock_os_getuid.return_value = fixed_uid
         osutil = MagicMock()  # pylint: disable=redefined-outer-name
+        osutil.enable_firewall.return_value = (MagicMock(), MagicMock())
 
         # Run
         migration_util.cleanup_metadata_server_artifacts(osutil)
@@ -80,7 +81,7 @@ class TestMetadataServerMigrationUtil(AgentTestCase):
         self.assertFalse(os.path.exists(metadata_server_p7b_file))
 
         # Assert Firewall rule calls
-        osutil.remove_firewall.assert_called_once_with(dst_ip=_KNOWN_METADATASERVER_IP, uid=fixed_uid)
+        osutil.remove_firewall.assert_called_once_with(dst_ip=_KNOWN_METADATASERVER_IP, uid=fixed_uid, wait=osutil.get_firewall_will_wait())
         osutil.enable_firewall.assert_called_once_with(dst_ip=KNOWN_WIRESERVER_IP, uid=fixed_uid)
 
     @patch('azurelinuxagent.common.conf.enable_firewall')
@@ -112,7 +113,7 @@ class TestMetadataServerMigrationUtil(AgentTestCase):
         self.assertFalse(os.path.exists(metadata_server_p7b_file))
 
         # Assert Firewall rule calls
-        osutil.remove_firewall.assert_called_once_with(dst_ip=_KNOWN_METADATASERVER_IP, uid=fixed_uid)
+        osutil.remove_firewall.assert_called_once_with(dst_ip=_KNOWN_METADATASERVER_IP, uid=fixed_uid, wait=osutil.get_firewall_will_wait())
         osutil.enable_firewall.assert_not_called()
 
     # Cleanup certificate files
diff --git a/tests/protocol/test_protocol_util.py b/tests/protocol/test_protocol_util.py
index 216c6cb4..6af6dbf8 100644
--- a/tests/protocol/test_protocol_util.py
+++ b/tests/protocol/test_protocol_util.py
@@ -174,6 +174,7 @@ class TestProtocolUtil(AgentTestCase):
         mock_enable_firewall.return_value = True
         protocol_util = get_protocol_util()
         protocol_util.osutil = MagicMock()
+        protocol_util.osutil.enable_firewall.return_value = (MagicMock(), MagicMock())
         protocol_util.dhcp_handler = MagicMock()
         protocol_util.dhcp_handler.endpoint = KNOWN_WIRESERVER_IP
 
@@ -213,6 +214,7 @@ class TestProtocolUtil(AgentTestCase):
         mock_enable_firewall.return_value = True
         protocol_util = get_protocol_util()
         protocol_util.osutil = MagicMock()
+        protocol_util.osutil.enable_firewall.return_value = (MagicMock(), MagicMock())
         mock_wire_client.return_value = MagicMock()
         protocol_util.dhcp_handler = MagicMock()
         protocol_util.dhcp_handler.endpoint = KNOWN_WIRESERVER_IP
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 1ce32129..386fab61 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -41,6 +41,7 @@ Debug.CgroupMonitorExtensionName = Microsoft.Azure.Monitor.AzureMonitorLinuxAgen
 Debug.EnableFastTrack = True
 Debug.EnableGAVersioning = False
 Debug.EtpCollectionPeriod = 300
+Debug.FirewallRulesLogPeriod = 86400
 DetectScvmmEnv = False
 EnableOverProvisioning = True
 Extension.LogDir = /var/log/azure