Repository: Azure/WALinuxAgent
Pull request #2154 (2021-02-01 22:17:21)

Setup persistent firewall rules on service restart
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

This PR ensures that the firewall rules for the agent are persisted over system restarts. This PR does the following - 
- Ensures Firewall rules are persisted over restarts by either using `firewalld.service` or creating a new unit file for doing the same
- If firewalld is available on the box, prioritize that because it already has provisions for making the rules persist over reboots.
- If not, create a new custom agent-network-setup.service which would always setup firewall rules on reboots before network of the VM is even established.

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).
- [ ] Except for special cases involving multiple contributors, the PR is started from a fork of the main repository, not a branch.
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made and Travis.CI is passing.

### 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/2154?src=pr&el=h1) Report
> Merging [#2154](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154?src=pr&el=desc) (776ff42) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/04b59cbf1f43612d3df0d858fefae77f33ae6220?el=desc) (04b59cb) will **increase** coverage by `0.09%`.
> The diff coverage is `85.95%`.

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

```diff
@@             Coverage Diff             @@
##           develop    #2154      +/-   ##
===========================================
+ Coverage    70.47%   70.57%   +0.09%     
===========================================
  Files           95       97       +2     
  Lines        13064    13305     +241     
  Branches      1859     1871      +12     
===========================================
+ Hits          9207     9390     +183     
- Misses        3441     3497      +56     
- Partials       416      418       +2     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/waagent\_network\_setup.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi93YWFnZW50X25ldHdvcmtfc2V0dXAucHk=) | `0.00% <0.00%> (ø)` | |
| [azurelinuxagent/common/osutil/arch.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvYXJjaC5weQ==) | `65.62% <66.66%> (-1.05%)` | :arrow_down: |
| [azurelinuxagent/common/osutil/clearlinux.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvY2xlYXJsaW51eC5weQ==) | `58.20% <66.66%> (+0.83%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/coreos.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvY29yZW9zLnB5) | `67.39% <66.66%> (-0.06%)` | :arrow_down: |
| [azurelinuxagent/common/osutil/freebsd.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvZnJlZWJzZC5weQ==) | `44.09% <66.66%> (+0.03%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/iosxe.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvaW9zeGUucHk=) | `43.90% <66.66%> (+1.79%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/mariner.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvbWFyaW5lci5weQ==) | `54.28% <66.66%> (+2.56%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/openbsd.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvb3BlbmJzZC5weQ==) | `23.93% <66.66%> (+0.55%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/redhat.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvcmVkaGF0LnB5) | `57.31% <66.66%> (+0.35%)` | :arrow_up: |
| [azurelinuxagent/common/osutil/suse.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9vc3V0aWwvc3VzZS5weQ==) | `63.79% <66.66%> (+0.15%)` | :arrow_up: |
| ... and [6 more](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154?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/2154?src=pr&el=footer). Last update [04b59cb...809dec3](https://codecov.io/gh/Azure/WALinuxAgent/pull/2154?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

narrieta: There are a couple of INFO messages in the log that look like errors and may be misleading. Can we rephrase them?

```
>>>>>>>>>>>> 2021-02-05T15:06:39.031001Z INFO ExtHandler ExtHandler firewall-cmd --state command returned error/not running: [Errno 2] No such file or directory: 'firewall-cmd': 'firewall-cmd'**
2021-02-05T15:06:39.031472Z INFO ExtHandler ExtHandler Firewalld service not running/unavailable, trying to set up walinuxagent-network-setup.service
>>>>>>>>>>>2021-02-05T15:06:39.130154Z INFO ExtHandler ExtHandler walinuxagent-network-setup.service not enabled. Command: systemctl is-enabled walinuxagent-network-setup.service, ExitCode: 1
Stdout:
Stderr: Failed to get unit file state for walinuxagent-network-setup.service: No such file or directory

2021-02-05T15:06:39.371742Z INFO ExtHandler ExtHandler Successfully added and enabled the walinuxagent-network-setup.service
2021-02-05T15:06:39.372280Z INFO ExtHandler ExtHandler Drop-in file /lib/systemd/system/walinuxagent-network-setup.service.d/10-environment.conf successfully updated
```

Patch:
diff --git a/azurelinuxagent/agent.py b/azurelinuxagent/agent.py
index 26ad7b21..b5ee2f3e 100644
--- a/azurelinuxagent/agent.py
+++ b/azurelinuxagent/agent.py
@@ -38,6 +38,7 @@ from azurelinuxagent.common.logcollector import LogCollector, OUTPUT_RESULTS_FIL
 from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
+from azurelinuxagent.common.utils.networkutil import AddFirewallRules
 from azurelinuxagent.common.version import AGENT_NAME, AGENT_LONG_VERSION, AGENT_VERSION, \
     DISTRO_NAME, DISTRO_VERSION, \
     PY_VERSION_MAJOR, PY_VERSION_MINOR, \
@@ -46,6 +47,24 @@ from azurelinuxagent.common.version import AGENT_NAME, AGENT_LONG_VERSION, AGENT
 from azurelinuxagent.pa.provision.default import ProvisionHandler
 
 
+class AgentCommands(object):
+    """
+    This is the list of all commands that the Linux Guest Agent supports
+    """
+    DeprovisionUser = "deprovision+user"
+    Deprovision = "deprovision"
+    Daemon = "daemon"
+    Start = "start"
+    RegisterService = "register-service"
+    RunExthandlers = "run-exthandlers"
+    Version = "version"
+    ShowConfig = "show-configuration"
+    Help = "help"
+    CollectLogs = "collect-logs"
+    SetupFirewall = "setup-firewall"
+    Provision = "provision"
+
+
 class Agent(object):
     def __init__(self, verbose, conf_file_path=None):
         """
@@ -190,6 +209,18 @@ class Agent(object):
             print("Detailed log output can be found at {0}".format(OUTPUT_RESULTS_FILE_PATH))
             sys.exit(1)
 
+    @staticmethod
+    def setup_firewall(firewall_metadata):
+
+        print("Setting up firewall for the WALinux Agent with args: {0}".format(firewall_metadata))
+        try:
+            AddFirewallRules.add_iptables_rules(firewall_metadata['wait'], firewall_metadata['dst_ip'],
+                                                firewall_metadata['uid'])
+            print("Successfully set the firewall rules")
+        except Exception as error:
+            print("Unable to add firewall rules. Error: {0}".format(ustr(error)))
+            sys.exit(1)
+
 
 def main(args=None):
     """
@@ -200,32 +231,34 @@ def main(args=None):
         args = []
     if len(args) <= 0:
         args = sys.argv[1:]
-    command, force, verbose, debug, conf_file_path, log_collector_full_mode = parse_args(args)
-    if command == "version":
+    command, force, verbose, debug, conf_file_path, log_collector_full_mode, firewall_metadata = parse_args(args)
+    if command == AgentCommands.Version:
         version()
-    elif command == "help":
+    elif command == AgentCommands.Help:
         print(usage())
-    elif command == "start":
+    elif command == AgentCommands.Start:
         start(conf_file_path=conf_file_path)
     else:
         try:
             agent = Agent(verbose, conf_file_path=conf_file_path)
-            if command == "deprovision+user":
+            if command == AgentCommands.DeprovisionUser:
                 agent.deprovision(force, deluser=True)
-            elif command == "deprovision":
+            elif command == AgentCommands.Deprovision:
                 agent.deprovision(force, deluser=False)
-            elif command == "provision":
+            elif command == AgentCommands.Provision:
                 agent.provision()
-            elif command == "register-service":
+            elif command == AgentCommands.RegisterService:
                 agent.register_service()
-            elif command == "daemon":
+            elif command == AgentCommands.Daemon:
                 agent.daemon()
-            elif command == "run-exthandlers":
+            elif command == AgentCommands.RunExthandlers:
                 agent.run_exthandlers(debug)
-            elif command == "show-configuration":
+            elif command == AgentCommands.ShowConfig:
                 agent.show_configuration()
-            elif command == "collect-logs":
+            elif command == AgentCommands.CollectLogs:
                 agent.collect_logs(log_collector_full_mode)
+            elif command == AgentCommands.SetupFirewall:
+                agent.setup_firewall(firewall_metadata)
         except Exception:
             logger.error(u"Failed to run '{0}': {1}",
                          command,
@@ -236,14 +269,24 @@ def parse_args(sys_args):
     """
     Parse command line arguments
     """
-    cmd = "help"
+    cmd = AgentCommands.Help
     force = False
     verbose = False
     debug = False
     conf_file_path = None
     log_collector_full_mode = False
+    firewall_metadata = {
+        "dst_ip": None,
+        "uid": None,
+        "wait": ""
+    }
+
+    regex_cmd_format = "^([-/]*){0}"
 
     for arg in sys_args:
+        if arg == "":
+            # Don't parse an empty parameter
+            continue
         m = re.match("^(?:[-/]*)configuration-path:([\w/\.\-_]+)", arg)  # pylint: disable=W1401
         if not m is None:
             conf_file_path = m.group(1)
@@ -253,38 +296,47 @@ def parse_args(sys_args):
                 print(usage())
                 sys.exit(1)
         elif re.match("^([-/]*)deprovision\\+user", arg):
-            cmd = "deprovision+user"
-        elif re.match("^([-/]*)deprovision", arg):
-            cmd = "deprovision"
-        elif re.match("^([-/]*)daemon", arg):
-            cmd = "daemon"
-        elif re.match("^([-/]*)start", arg):
-            cmd = "start"
-        elif re.match("^([-/]*)register-service", arg):
-            cmd = "register-service"
-        elif re.match("^([-/]*)run-exthandlers", arg):
-            cmd = "run-exthandlers"
-        elif re.match("^([-/]*)version", arg):
-            cmd = "version"
-        elif re.match("^([-/]*)verbose", arg):
+            cmd = AgentCommands.DeprovisionUser
+        elif re.match(regex_cmd_format.format(AgentCommands.Deprovision), arg):
+            cmd = AgentCommands.Deprovision
+        elif re.match(regex_cmd_format.format(AgentCommands.Daemon), arg):
+            cmd = AgentCommands.Daemon
+        elif re.match(regex_cmd_format.format(AgentCommands.Start), arg):
+            cmd = AgentCommands.Start
+        elif re.match(regex_cmd_format.format(AgentCommands.RegisterService), arg):
+            cmd = AgentCommands.RegisterService
+        elif re.match(regex_cmd_format.format(AgentCommands.RunExthandlers), arg):
+            cmd = AgentCommands.RunExthandlers
+        elif re.match(regex_cmd_format.format(AgentCommands.Version), arg):
+            cmd = AgentCommands.Version
+        elif re.match(regex_cmd_format.format("verbose"), arg):
             verbose = True
-        elif re.match("^([-/]*)debug", arg):
+        elif re.match(regex_cmd_format.format("debug"), arg):
             debug = True
-        elif re.match("^([-/]*)force", arg):
+        elif re.match(regex_cmd_format.format("force"), arg):
             force = True
-        elif re.match("^([-/]*)show-configuration", arg):
-            cmd = "show-configuration"
+        elif re.match(regex_cmd_format.format(AgentCommands.ShowConfig), arg):
+            cmd = AgentCommands.ShowConfig
         elif re.match("^([-/]*)(help|usage|\\?)", arg):
-            cmd = "help"
-        elif re.match("^([-/]*)collect-logs", arg):
-            cmd = "collect-logs"
-        elif re.match("^([-/]*)full", arg):
+            cmd = AgentCommands.Help
+        elif re.match(regex_cmd_format.format(AgentCommands.CollectLogs), arg):
+            cmd = AgentCommands.CollectLogs
+        elif re.match(regex_cmd_format.format("full"), arg):
             log_collector_full_mode = True
+        elif re.match(regex_cmd_format.format(AgentCommands.SetupFirewall), arg):
+            cmd = AgentCommands.SetupFirewall
+        elif re.match(regex_cmd_format.format("dst_ip=(?P<dst_ip>[\\d.]{7,})"), arg):
+            firewall_metadata['dst_ip'] = re.match(regex_cmd_format.format("dst_ip=(?P<dst_ip>[\\d.]{7,})"), arg).group(
+                'dst_ip')
+        elif re.match(regex_cmd_format.format("uid=(?P<uid>[\\d]+)"), arg):
+            firewall_metadata['uid'] = re.match(regex_cmd_format.format("uid=(?P<uid>[\\d]+)"), arg).group('uid')
+        elif re.match(regex_cmd_format.format("(w|wait)$"), arg):
+            firewall_metadata['wait'] = "-w"
         else:
-            cmd = "help"
+            cmd = AgentCommands.Help
             break
 
-    return cmd, force, verbose, debug, conf_file_path, log_collector_full_mode
+    return cmd, force, verbose, debug, conf_file_path, log_collector_full_mode, firewall_metadata
 
 
 def version():
@@ -308,7 +360,7 @@ def usage():
     s += ("usage: {0} [-verbose] [-force] [-help] "
            "-configuration-path:<path to configuration file>" 
            "-deprovision[+user]|-register-service|-version|-daemon|-start|"
-           "-run-exthandlers|-show-configuration|-collect-logs [-full]"
+           "-run-exthandlers|-show-configuration|-collect-logs [-full]|-setup-firewall [-dst_ip=<IP> -uid=<UID> [-w/--wait]]"
            "").format(sys.argv[0])
     s += "\n"
     return s
diff --git a/azurelinuxagent/common/event.py b/azurelinuxagent/common/event.py
index 172bfae5..c03ddd27 100644
--- a/azurelinuxagent/common/event.py
+++ b/azurelinuxagent/common/event.py
@@ -104,6 +104,7 @@ class WALAEventOperation:  # pylint: disable=no-init
     LogCollection = "LogCollection"
     OSInfo = "OSInfo"
     Partition = "Partition"
+    PersistFirewallRules = "PersistFirewallRules"
     PluginSettingsVersionMismatch = "PluginSettingsVersionMismatch"
     InvalidExtensionConfig = "InvalidExtensionConfig"
     ProcessGoalState = "ProcessGoalState"
diff --git a/azurelinuxagent/common/osutil/arch.py b/azurelinuxagent/common/osutil/arch.py
index e04e5573..7a863eeb 100644
--- a/azurelinuxagent/common/osutil/arch.py
+++ b/azurelinuxagent/common/osutil/arch.py
@@ -16,7 +16,6 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-import os  # pylint: disable=W0611
 import azurelinuxagent.common.utils.shellutil as shellutil
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
 
@@ -25,6 +24,14 @@ class ArchUtil(DefaultOSUtil):
     def __init__(self):
         super(ArchUtil, self).__init__()
         self.jit_enabled = True
+
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/bin"
     
     def is_dhcp_enabled(self):
         return True
diff --git a/azurelinuxagent/common/osutil/clearlinux.py b/azurelinuxagent/common/osutil/clearlinux.py
index 49626110..fda72c16 100644
--- a/azurelinuxagent/common/osutil/clearlinux.py
+++ b/azurelinuxagent/common/osutil/clearlinux.py
@@ -42,6 +42,14 @@ class ClearLinuxUtil(DefaultOSUtil):
         self.agent_conf_file_path = '/usr/share/defaults/waagent/waagent.conf'
         self.jit_enabled = True
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/bin"
+
     def is_dhcp_enabled(self):
         return True
 
diff --git a/azurelinuxagent/common/osutil/coreos.py b/azurelinuxagent/common/osutil/coreos.py
index 2226f1be..fc0a6604 100644
--- a/azurelinuxagent/common/osutil/coreos.py
+++ b/azurelinuxagent/common/osutil/coreos.py
@@ -42,6 +42,10 @@ class CoreOSUtil(DefaultOSUtil):
             py_path = self.waagent_path
         os.environ['PYTHONPATH'] = py_path
 
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/share/oem/bin"
+
     def is_sys_user(self, username):
         # User 'core' is not a sysuser.
         if username == 'core':
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index dda0b898..505f2e13 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -143,6 +143,14 @@ class DefaultOSUtil(object):
     def get_service_name():
         return "waagent"
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/lib/systemd/system"
+
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/sbin"
+
     def get_firewall_dropped_packets(self, dst_ip=None):
         # If a previous attempt failed, do not retry
         global _enable_firewall  # pylint: disable=W0603
@@ -245,7 +253,6 @@ class DefaultOSUtil(object):
             wait = self.get_firewall_will_wait()
 
             # If the DROP rule exists, make no changes
-            firewall_established = False  # pylint: disable=W0612
             try:
                 drop_rule = _get_firewall_drop_command(wait, "-C", dst_ip)
                 shellutil.run_command(drop_rule)
diff --git a/azurelinuxagent/common/osutil/freebsd.py b/azurelinuxagent/common/osutil/freebsd.py
index f0cfed20..f8ee6db8 100644
--- a/azurelinuxagent/common/osutil/freebsd.py
+++ b/azurelinuxagent/common/osutil/freebsd.py
@@ -22,7 +22,6 @@ import binascii
 import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.utils.shellutil as shellutil
 import azurelinuxagent.common.utils.textutil as textutil
-from azurelinuxagent.common.utils.networkutil import RouteEntry  # pylint: disable=W0611
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.exception import OSUtilError
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
@@ -36,6 +35,10 @@ class FreeBSDOSUtil(DefaultOSUtil):
         self._scsi_disks_timeout_set = False
         self.jit_enabled = True
 
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/local/sbin"
+
     def set_hostname(self, hostname):
         rc_file_path = '/etc/rc.conf'
         conf_file = fileutil.read_file(rc_file_path).split("\n")
diff --git a/azurelinuxagent/common/osutil/iosxe.py b/azurelinuxagent/common/osutil/iosxe.py
index f09996f0..ace28f07 100644
--- a/azurelinuxagent/common/osutil/iosxe.py
+++ b/azurelinuxagent/common/osutil/iosxe.py
@@ -41,6 +41,10 @@ class IosxeOSUtil(DefaultOSUtil):
     def __init__(self):  # pylint: disable=W0235
         super(IosxeOSUtil, self).__init__()
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
     def set_hostname(self, hostname):
         """
         Unlike redhat 6.x, redhat 7.x will set hostname via hostnamectl
diff --git a/azurelinuxagent/common/osutil/mariner.py b/azurelinuxagent/common/osutil/mariner.py
index 44e6800a..b6ea5900 100644
--- a/azurelinuxagent/common/osutil/mariner.py
+++ b/azurelinuxagent/common/osutil/mariner.py
@@ -24,6 +24,14 @@ class MarinerOSUtil(DefaultOSUtil):
         super(MarinerOSUtil, self).__init__()
         self.jit_enabled = True
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/bin"
+
     def is_dhcp_enabled(self):
         return True
 
diff --git a/azurelinuxagent/common/osutil/openbsd.py b/azurelinuxagent/common/osutil/openbsd.py
index a945cac9..0a80fc93 100644
--- a/azurelinuxagent/common/osutil/openbsd.py
+++ b/azurelinuxagent/common/osutil/openbsd.py
@@ -43,6 +43,10 @@ class OpenBSDOSUtil(DefaultOSUtil):
         self.jit_enabled = True
         self._scsi_disks_timeout_set = False
 
+    @staticmethod
+    def get_agent_bin_path():
+        return "/usr/local/sbin"
+
     def get_instance_id(self):
         ret, output = shellutil.run_get_output("sysctl -n hw.uuid")
         if ret != 0 or UUID_PATTERN.match(output) is None:
diff --git a/azurelinuxagent/common/osutil/redhat.py b/azurelinuxagent/common/osutil/redhat.py
index a1470cc1..9759d113 100644
--- a/azurelinuxagent/common/osutil/redhat.py
+++ b/azurelinuxagent/common/osutil/redhat.py
@@ -99,6 +99,10 @@ class RedhatOSUtil(Redhat6xOSUtil):
         super(RedhatOSUtil, self).__init__()
         self.service_name = self.get_service_name()
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
     def set_hostname(self, hostname):
         """
         Unlike redhat 6.x, redhat 7.x will set hostname via hostnamectl
diff --git a/azurelinuxagent/common/osutil/suse.py b/azurelinuxagent/common/osutil/suse.py
index 1aa5231b..9b10b00e 100644
--- a/azurelinuxagent/common/osutil/suse.py
+++ b/azurelinuxagent/common/osutil/suse.py
@@ -69,6 +69,10 @@ class SUSEOSUtil(SUSE11OSUtil):
         super(SUSEOSUtil, self).__init__()
         self.dhclient_name = 'wickedd-dhcp4'
 
+    @staticmethod
+    def get_systemd_unit_file_install_path():
+        return "/usr/lib/systemd/system"
+
     def set_hostname(self, hostname):
         self._run_command_without_raising(["hostnamectl", "set-hostname", hostname], log_error=False)
 
diff --git a/azurelinuxagent/common/persist_firewall_rules.py b/azurelinuxagent/common/persist_firewall_rules.py
new file mode 100644
index 00000000..e75b627e
--- /dev/null
+++ b/azurelinuxagent/common/persist_firewall_rules.py
@@ -0,0 +1,238 @@
+#
+# Copyright 2018 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.6+ and Openssl 1.0+
+#
+import os
+import sys
+
+from azurelinuxagent.common import logger
+from azurelinuxagent.common.cgroupapi import CGroupsApi
+from azurelinuxagent.common.event import add_event, WALAEventOperation
+from azurelinuxagent.common.future import ustr
+from azurelinuxagent.common.osutil import get_osutil
+from azurelinuxagent.common.utils import shellutil, fileutil
+from azurelinuxagent.common.utils.networkutil import AddFirewallRules
+from azurelinuxagent.common.utils.shellutil import CommandError
+
+
+class PersistFirewallRulesHandler(object):
+
+    __SERVICE_FILE_CONTENT = """
+# This unit file was created by the Azure VM Agent.
+# Do not edit.
+[Unit]
+Description=Setup network rules for WALinuxAgent 
+Before=network-pre.target
+Wants=network-pre.target
+DefaultDependencies=no
+
+[Service]
+Type=oneshot
+Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
+ExecStart={py_path} ${{EGG}} --setup-firewall --dst_ip=${{DST_IP}} --uid=${{UID}} ${{WAIT}}
+RemainAfterExit=false
+
+[Install]
+WantedBy=network.target
+"""
+
+    __OVERRIDE_CONTENT = """
+# This drop-in unit file was created by the Azure VM Agent.
+# Do not edit.
+[Service]
+# The first line clears the old data, and the 2nd line overwrites it.
+Environment=
+Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
+"""
+
+    _AGENT_NETWORK_SETUP_NAME_FORMAT = "{0}-network-setup.service"
+    _DROP_IN_ENV_FILE_NAME = "10-environment.conf"
+
+    _FIREWALLD_RUNNING_CMD = ["firewall-cmd", "--state"]
+
+    def __init__(self, dst_ip, uid):
+        """
+        This class deals with ensuring that Firewall rules are persisted over system reboots.
+        It tries to employ using Firewalld.service if present first as it already has provisions for persistent rules.
+        If not, it then creates a new agent-network-setup.service file and copy it over to the osutil.get_systemd_unit_file_install_path() dynamically
+        On top of it, on every service restart it ensures that the WireIP is overwritten and the new IP is blocked as well.
+        """
+        osutil = get_osutil()
+        self._network_setup_service_name = self._AGENT_NETWORK_SETUP_NAME_FORMAT.format(osutil.get_service_name())
+        self._is_systemd = CGroupsApi.is_systemd()
+        self._systemd_file_path = osutil.get_systemd_unit_file_install_path()
+        self._dst_ip = dst_ip
+        self._uid = uid
+        self._wait = osutil.get_firewall_will_wait()
+        # The custom service will try to call the current agent executable to setup the firewall rules
+        self._current_agent_executable_path = os.path.join(os.getcwd(), sys.argv[0])
+
+    @staticmethod
+    def _is_firewall_service_running():
+        # Check if firewall-cmd can connect to the daemon
+        # https://docs.fedoraproject.org/en-US/Fedora/19/html/Security_Guide/sec-Check_if_firewalld_is_running.html
+        # Eg:    firewall-cmd --state
+        #           >   running
+        firewalld_state = PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD
+        try:
+            return shellutil.run_command(firewalld_state).rstrip() == "running"
+        except Exception as error:
+            logger.verbose("{0} command failed: {1}".format(' '.join(firewalld_state), ustr(error)))
+        return False
+
+    def setup(self):
+        if self._is_firewall_service_running():
+            logger.info("Firewalld.service present on the VM, setting up permanent rules on the VM")
+            # In case of a failure, this would throw. In such a case, we don't need to try to setup our custom service
+            # because on system reboot, all iptable rules are reset by firewalld.service so it would be a no-op.
+            self._setup_permanent_firewalld_rules()
+            return
+
+        logger.info(
+            "Firewalld service not running/unavailable, trying to set up {0}".format(self._network_setup_service_name))
+
+        if not CGroupsApi.is_systemd():
+            raise Exception("Did not detect Systemd, unable to set {0}".format(self._network_setup_service_name))
+
+        self._setup_network_setup_service()
+
+    def __verify_firewall_rules_enabled(self):
+        # Check if firewall-rules have already been enabled
+        # This function would also return False if the dest-ip is changed. So no need to check separately for that
+        try:
+            AddFirewallRules.check_firewalld_rule_applied(self._dst_ip, self._uid)
+        except Exception as error:
+            logger.verbose(
+                "Check if Firewall rules already applied using firewalld.service failed: {0}".format(ustr(error)))
+            return False
+
+        return True
+
+    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")
+        # 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")
+
+    def __verify_network_setup_service_enabled(self):
+        # Check if the custom service has already been enabled
+        cmd = ["systemctl", "is-enabled", self._network_setup_service_name]
+        try:
+            return shellutil.run_command(cmd).rstrip() == "enabled"
+        except CommandError as error:
+            msg = "{0} not enabled. Command: {1}, ExitCode: {2}\nStdout: {3}\nStderr: {4}".format(
+                self._network_setup_service_name, ' '.join(cmd), error.returncode, error.stdout, error.stderr)
+        except Exception as error:
+            msg = "Ran into error, {0} not enabled. Error: {1}".format(self._network_setup_service_name, ustr(error))
+
+        logger.verbose(msg)
+        return False
+
+    def _setup_network_setup_service(self):
+        if self.__verify_network_setup_service_enabled():
+            logger.info("Service: {0} already enabled. No change needed.".format(self._network_setup_service_name))
+            self.__log_network_setup_service_logs()
+
+        else:
+            logger.info("Service: {0} not enabled. Adding it now".format(self._network_setup_service_name))
+            # Create unit file with default values
+            self.__set_service_unit_file()
+            logger.info("Successfully added and enabled the {0}".format(self._network_setup_service_name))
+
+        # Even if service is enabled, we need to overwrite the drop-in file with the current IP in case it changed.
+        # This is to handle the case where WireIP can change midway on service restarts.
+        # Additionally, incase of auto-update this would also update the location of the new EGG file ensuring that
+        # the service is always run from the most latest agent.
+        self.__set_drop_in_file()
+
+    def __set_drop_in_file(self):
+        drop_in_file = os.path.join(self._systemd_file_path, "{0}.d".format(self._network_setup_service_name),
+                                    self._DROP_IN_ENV_FILE_NAME)
+        parent, _ = os.path.split(drop_in_file)
+        try:
+            if not os.path.exists(parent):
+                fileutil.mkdir(parent, mode=0o755)
+            fileutil.write_file(drop_in_file,
+                                self.__OVERRIDE_CONTENT.format(egg_path=self._current_agent_executable_path,
+                                                               wire_ip=self._dst_ip,
+                                                               user_id=self._uid,
+                                                               wait=self._wait))
+            logger.info("Drop-in file {0} successfully updated".format(drop_in_file))
+        except Exception:
+            self.__remove_file_without_raising(drop_in_file)
+            raise
+
+    def __set_service_unit_file(self):
+        service_unit_file = os.path.join(self._systemd_file_path, self._network_setup_service_name)
+        try:
+            fileutil.write_file(service_unit_file,
+                                self.__SERVICE_FILE_CONTENT.format(egg_path=self._current_agent_executable_path,
+                                                                   wire_ip=self._dst_ip,
+                                                                   user_id=self._uid,
+                                                                   wait=self._wait,
+                                                                   py_path=sys.executable))
+            fileutil.chmod(service_unit_file, 0o644)
+
+            # Finally enable the service. This is needed to ensure the service is started on system boot
+            cmd = ["systemctl", "enable", self._network_setup_service_name]
+            try:
+                shellutil.run_command(cmd)
+            except CommandError as error:
+                msg = "Unable to enable service: {0}; deleting service file: {1}. Command: {2}, Exit-code: {3}.\nstdout: {4}\nstderr: {5}".format(
+                    self._network_setup_service_name, service_unit_file, ' '.join(cmd), error.returncode, error.stdout,
+                    error.stderr)
+                raise Exception(msg)
+
+        except Exception:
+            self.__remove_file_without_raising(service_unit_file)
+            raise
+
+    @staticmethod
+    def __remove_file_without_raising(file_path):
+        if os.path.exists(file_path):
+            try:
+                os.remove(file_path)
+            except Exception as error:
+                logger.warn("Unable to delete file: {0}; Error: {1}".format(file_path, ustr(error)))
+
+    def __log_network_setup_service_logs(self):
+        # Get logs from journalctl - https://www.freedesktop.org/software/systemd/man/journalctl.html
+        cmd = ["journalctl", "-u", self._network_setup_service_name, "-b"]
+        cmd_success = False
+        try:
+            stdout = shellutil.run_command(cmd)
+            msg = "Logs from the {0} since system boot:\n {1}".format(self._network_setup_service_name, stdout)
+            cmd_success = True
+            logger.info(msg)
+        except CommandError as error:
+            msg = "Command: {0} failed with ExitCode: {1}\nStdout: {2}\nStderr: {3}".format(' '.join(cmd),
+                                                                                            error.returncode,
+                                                                                            error.stdout, error.stderr)
+            logger.warn(msg)
+        except Exception as error:
+            msg = "Ran into unexpected error when getting logs for {0} service. Error: {1}".format(
+                self._network_setup_service_name, ustr(error))
+            logger.warn(msg)
+
+        add_event(
+            op=WALAEventOperation.PersistFirewallRules,
+            is_success=cmd_success,
+            message=msg,
+            log_event=False)
diff --git a/azurelinuxagent/common/utils/networkutil.py b/azurelinuxagent/common/utils/networkutil.py
index d3617732..a1269ce7 100644
--- a/azurelinuxagent/common/utils/networkutil.py
+++ b/azurelinuxagent/common/utils/networkutil.py
@@ -18,6 +18,7 @@
 
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils import shellutil
+from azurelinuxagent.common.utils.shellutil import CommandError
 
 
 class RouteEntry(object):
@@ -98,7 +99,27 @@ class NetworkInterfaceCard:
         return "{{ {0} }}".format(", ".join(entries))
 
 
+class FirewallCmdDirectCommands(object):
+    # firewall-cmd --direct --passthrough ipv4 -t security -A OUTPUT -d 1.2.3.5 -p tcp -m owner --uid-owner 999 -j ACCEPT
+    # success
+    PassThrough = "--passthrough"
+
+    # firewall-cmd --direct --query-passthrough ipv4 -t security -A OUTPUT -d 1.2.3.5 -p tcp -m owner --uid-owner 9999 -j ACCEPT
+    # yes
+    # firewall-cmd --direct --query-passthrough ipv4 -t security -A OUTPUT -d 1.2.3.5 -p tcp -m owner --uid-owner 999 -j ACCEPT
+    # no
+    QueryPassThrough = "--query-passthrough"
+
+
 class AddFirewallRules(object):
+    """
+    This class is a utility class which is only meant to orchestrate adding Firewall rules (both iptables and firewalld).
+    This would also be called from a separate utility binary which would be very early up in the boot order of the VM,
+    due to which it would not have access to basic mounts like file-system.
+    Please make sure to not log anything in any function this class.
+    """
+
+    __ACCEPT_COMMAND = "-A"
 
     @staticmethod
     def _add_wait(wait, command):
@@ -106,39 +127,105 @@ class AddFirewallRules(object):
         If 'wait' is True, adds the wait option (-w) to the given iptables command line
         """
         if wait:
-            command.insert(1, "-w")
+            command.append("-w")
         return command
 
     @staticmethod
-    def get_iptables_accept_command(wait, command, destination, owner_uid):
-        cmd = ["iptables", "-t", "security", command, "OUTPUT", "-d", destination, "-p", "tcp", "-m", "owner",
-               "--uid-owner", str(owner_uid), "-j", "ACCEPT"]
+    def __get_iptables_base_command():
+        return ["iptables"]
+
+    @staticmethod
+    def __get_firewalld_base_command(command):
+        # For more documentation - https://firewalld.org/documentation/man-pages/firewall-cmd.html
+        return ["firewall-cmd", "--permanent", "--direct", command, "ipv4"]
+
+    @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)
 
     @staticmethod
-    def get_iptables_drop_command(wait, command, destination):
-        cmd = ["iptables", "-t", "security", command, "OUTPUT", "-d", destination, "-p", "tcp", "-m", "conntrack",
-               "--ctstate", "INVALID,NEW", "-j", "DROP"]
+    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 add_iptables_rules(wait, dst_ip, uid):
-        def _raise_if_empty(val, name):
-            if val == "":
-                raise Exception("{0} should not be empty".format(name))
+    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
 
-        _raise_if_empty(dst_ip, "Destination IP")
-        _raise_if_empty(uid, "User ID")
-        accept_rule = AddFirewallRules.get_iptables_accept_command(wait, "-A", dst_ip, uid)
-        try:
-            shellutil.run_command(accept_rule)
-        except Exception as e:
-            msg = "Unable to add ACCEPT firewall rule '{0}' - {1}".format(accept_rule, ustr(e))
-            raise Exception(msg)
+    @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))
+        return cmd
 
-        drop_rule = AddFirewallRules.get_iptables_drop_command(wait, "-A", dst_ip)
+    @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.__ACCEPT_COMMAND, destination,
+                                                                uid))
+        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.__ACCEPT_COMMAND, destination))
+        return cmd
+
+    @staticmethod
+    def __raise_if_empty(val, name):
+        if val == "":
+            raise Exception("{0} should not be empty".format(name))
+
+    @staticmethod
+    def __execute_cmd(cmd):
         try:
-            shellutil.run_command(drop_rule)
-        except Exception as e:
-            msg = "Unable to add DROP firewall rule '{0}' - {1}".format(drop_rule, ustr(e))
+            shellutil.run_command(cmd)
+        except CommandError as error:
+            msg = "Command {0} failed with exit-code: {1}\nStdout: {2}\nStderr: {3}".format(' '.join(cmd),
+                                                                                            error.returncode,
+                                                                                            ustr(error.stdout),
+                                                                                            ustr(error.stderr))
             raise Exception(msg)
+
+    @staticmethod
+    def add_iptables_rules(wait, dst_ip, uid):
+
+        AddFirewallRules.__raise_if_empty(dst_ip, "Destination IP")
+        AddFirewallRules.__raise_if_empty(uid, "User ID")
+
+        accept_rule = AddFirewallRules.get_iptables_accept_command(wait, AddFirewallRules.__ACCEPT_COMMAND, dst_ip, uid)
+        AddFirewallRules.__execute_cmd(accept_rule)
+
+        drop_rule = AddFirewallRules.get_iptables_drop_command(wait, AddFirewallRules.__ACCEPT_COMMAND, dst_ip)
+        AddFirewallRules.__execute_cmd(drop_rule)
+
+    @staticmethod
+    def __execute_firewalld_commands(command, dst_ip, uid):
+        AddFirewallRules.__raise_if_empty(dst_ip, "Destination IP")
+        AddFirewallRules.__raise_if_empty(uid, "User ID")
+
+        accept_cmd = AddFirewallRules.get_firewalld_accept_command(command, dst_ip, uid)
+        AddFirewallRules.__execute_cmd(accept_cmd)
+
+        drop_cmd = AddFirewallRules.get_firewalld_drop_command(command, dst_ip)
+        AddFirewallRules.__execute_cmd(drop_cmd)
+
+    @staticmethod
+    def add_firewalld_rules(dst_ip, uid):
+        # Firwalld.service fails if we set `-w` in the iptables command, so not adding it at all for firewalld commands
+        AddFirewallRules.__execute_firewalld_commands(FirewallCmdDirectCommands.PassThrough, dst_ip, uid)
+
+    @staticmethod
+    def check_firewalld_rule_applied(dst_ip, uid):
+        AddFirewallRules.__execute_firewalld_commands(FirewallCmdDirectCommands.QueryPassThrough, dst_ip, uid)
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index b8f6f761..1146f82c 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -38,6 +38,7 @@ import azurelinuxagent.common.logger as logger
 import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.utils.restutil as restutil
 import azurelinuxagent.common.utils.textutil as textutil
+from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
 from azurelinuxagent.common.cgroupapi import CGroupsApi
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 
@@ -297,6 +298,7 @@ class UpdateHandler(object):
             self._ensure_readonly_files()
             self._ensure_cgroups_initialized()
             self._ensure_extension_telemetry_state_configured_properly(protocol)
+            self._ensure_firewall_rules_persisted(dst_ip=protocol.get_endpoint())
 
             # Get all thread handlers
             telemetry_handler = get_send_telemetry_events_handler(self.protocol_util)
@@ -829,6 +831,25 @@ class UpdateHandler(object):
         except Exception as e:
             logger.warn("Error when trying to delete existing Extension events directory. Error: {0}".format(ustr(e)))
 
+    @staticmethod
+    def _ensure_firewall_rules_persisted(dst_ip):
+        is_success = False
+        logger.info("Starting setup for Persistent firewall rules")
+        try:
+            PersistFirewallRulesHandler(dst_ip=dst_ip, uid=os.getuid()).setup()
+            msg = "Persistent firewall rules setup successfully"
+            is_success = True
+            logger.info(msg)
+        except Exception as error:
+            msg = "Unable to setup the persistent firewall rules: {0}".format(ustr(error))
+            logger.error(msg)
+
+        add_event(
+            op=WALAEventOperation.PersistFirewallRules,
+            is_success=is_success,
+            message=msg,
+            log_event=False)
+
 
 class GuestAgent(object):
     def __init__(self, path=None, pkg=None, host=None):
diff --git a/bin/waagent_network_setup.py b/bin/waagent_network_setup.py
deleted file mode 100644
index f123a05a..00000000
--- a/bin/waagent_network_setup.py
+++ /dev/null
@@ -1,85 +0,0 @@
-#!/usr/bin/env python
-#
-# Azure Linux Agent
-#
-# Copyright 2015 Microsoft Corporation
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# Requires Python 2.6 and Openssl 1.0+
-#
-# Implements parts of RFC 2131, 1541, 1497 and
-# http://msdn.microsoft.com/en-us/library/cc227282%28PROT.10%29.aspx
-# http://msdn.microsoft.com/en-us/library/cc227259%28PROT.13%29.aspx
-#
-
-import sys
-import re
-from azurelinuxagent.common.future import ustr
-from azurelinuxagent.common.utils.networkutil import AddFirewallRules
-
-
-def __usage():
-    """
-    Return the usage for this script
-    """
-    return "\nUsage: {0} [-help] -dst_ip=<Wireserver IP> -uid=[User ID] [-w|--wait]\n".format(sys.argv[0])
-
-
-def _setup_firewall_rules():
-    print("Setting up firewall rules")
-    try:
-        args = sys.argv[1:]
-        dst_ip = ""
-        uid = ""
-        wait = ""
-
-        # The command to call this would look something like this -
-        # ExecStart=PYPATH /usr/sbin/waagent_network_setup.py --dst_ip=1.2.3.4 --uid=9999 [-w]
-        for arg in args:
-            dst_ip_match = re.match("^([-/]*)dst_ip=(?P<dst_ip>[\d.]{7,})", arg)
-            uid_match = re.match("^([-/]*)uid=(?P<uid>[\d]+)", arg)
-            wait_match = re.match("^([-/]*)(w|wait)$", arg)
-
-            if dst_ip_match:
-                dst_ip = dst_ip_match.group('dst_ip')
-            elif uid_match:
-                uid = uid_match.group('uid')
-            elif wait_match:
-                wait = "-w"
-            else:
-                # If any command is not recognised, just print the help command and exit
-                print(__usage())
-                sys.exit(0)
-
-        AddFirewallRules.add_iptables_rules(wait, dst_ip, uid)
-        print("Setting Firewall rules completed")
-    except Exception as error:
-        print("Unable to setup firewall rules: {0}".format(ustr(error)))
-        sys.exit(1)
-
-
-if __name__ == '__main__':
-    """
-    This file would be called by the waagent-network-setup service to bridge the gap between the system reboot and 
-    waagent.service start. The sole purpose of this file is to setup firewall rules for the VM.
-    
-    The parameters for this file would be set in the unit file for the service, which would be created dynamically the 
-    first time waagent.service starts. On every service restart, we would monitor the Wireserver IP and reset it in the 
-    unit file if it changes.
-    
-    This service would run before the network would be setup, so we need to maintain state of the previous Wireserver IP
-    to ensure we restrict access to the Wire IP.
-    """
-    _setup_firewall_rules()
-
diff --git a/setup.py b/setup.py
index be4e5c59..ae287661 100755
--- a/setup.py
+++ b/setup.py
@@ -38,9 +38,9 @@ def set_files(data_files, dest=None, src=None):
     data_files.append((dest, src))
 
 
-def set_bin_files(data_files, dest="/usr/sbin", src=None):
+def set_bin_files(data_files, dest, src=None):
     if src is None:
-        src = ["bin/waagent", "bin/waagent2.0", "bin/waagent_network_setup.py"]
+        src = ["bin/waagent", "bin/waagent2.0"]
     data_files.append((dest, src))
 
 
@@ -62,14 +62,11 @@ def set_sysv_files(data_files, dest="/etc/rc.d/init.d", src=None):
     data_files.append((dest, src))
 
 
-def set_systemd_files(data_files, dest="/lib/systemd/system", src=None):
+def set_systemd_files(data_files, dest, src=None):
     if src is None:
         src = ["init/waagent.service"]
     data_files.append((dest, src))
 
-    set_files(data_files, dest="/etc/systemd/system", src=["init/ubuntu/azure.slice",
-                                                           "init/ubuntu/azure-vmextensions.slice"])
-
 
 def set_freebsd_rc_files(data_files, dest="/etc/rc.d/", src=None):
     if src is None:
@@ -95,9 +92,12 @@ def get_data_files(name, version, fullname):  # pylint: disable=R0912
     Determine data_files according to distro name, version and init system type
     """
     data_files = []
+    osutil = get_osutil()
+    systemd_dir_path = osutil.get_systemd_unit_file_install_path()
+    agent_bin_path = osutil.get_agent_bin_path()
 
+    set_bin_files(data_files, dest=agent_bin_path)
     if name == 'redhat' or name == 'centos':  # pylint: disable=R1714
-        set_bin_files(data_files)
         set_conf_files(data_files)
         set_logrotate_files(data_files)
         set_udev_files(data_files)
@@ -105,18 +105,16 @@ def get_data_files(name, version, fullname):  # pylint: disable=R0912
             set_sysv_files(data_files)
         else:
             # redhat7.0+ use systemd
-            set_systemd_files(data_files, dest="/usr/lib/systemd/system")
+            set_systemd_files(data_files, dest=systemd_dir_path)
             if version.startswith("7.1"):
                 # TODO this is a mitigation to systemctl bug on 7.1
                 set_sysv_files(data_files)
     elif name == 'arch':
-        set_bin_files(data_files, dest="/usr/bin")
         set_conf_files(data_files, src=["config/arch/waagent.conf"])
         set_udev_files(data_files)
-        set_systemd_files(data_files, dest='/usr/lib/systemd/system',
+        set_systemd_files(data_files, dest=systemd_dir_path,
                           src=["init/arch/waagent.service"])
     elif name == 'coreos':
-        set_bin_files(data_files, dest="/usr/share/oem/bin")
         set_conf_files(data_files, dest="/usr/share/oem",
                        src=["config/coreos/waagent.conf"])
         set_logrotate_files(data_files)
@@ -124,19 +122,16 @@ def get_data_files(name, version, fullname):  # pylint: disable=R0912
         set_files(data_files, dest="/usr/share/oem",
                   src=["init/coreos/cloud-config.yml"])
     elif "Clear Linux" in fullname:
-        set_bin_files(data_files, dest="/usr/bin")
         set_conf_files(data_files, dest="/usr/share/defaults/waagent",
                        src=["config/clearlinux/waagent.conf"])
-        set_systemd_files(data_files, dest='/usr/lib/systemd/system',
+        set_systemd_files(data_files, dest=systemd_dir_path,
                           src=["init/clearlinux/waagent.service"])
     elif name == 'mariner':
-        set_bin_files(data_files, dest="/usr/bin")
         set_conf_files(data_files, dest="/etc",
                        src=["config/mariner/waagent.conf"])
-        set_systemd_files(data_files, dest='/usr/lib/systemd/system',
+        set_systemd_files(data_files, dest=systemd_dir_path,
                           src=["init/mariner/waagent.service"])
     elif name == 'ubuntu':
-        set_bin_files(data_files)
         set_conf_files(data_files, src=["config/ubuntu/waagent.conf"])
         set_logrotate_files(data_files)
         set_udev_files(data_files)
@@ -151,10 +146,9 @@ def get_data_files(name, version, fullname):  # pylint: disable=R0912
                       src=["init/ubuntu/snappy/walinuxagent.yml"])
         else:
             # Ubuntu15.04+ uses systemd
-            set_systemd_files(data_files,
+            set_systemd_files(data_files, dest=systemd_dir_path,
                               src=["init/ubuntu/walinuxagent.service"])
     elif name == 'suse' or name == 'opensuse':  # pylint: disable=R1714
-        set_bin_files(data_files)
         set_conf_files(data_files, src=["config/suse/waagent.conf"])
         set_logrotate_files(data_files)
         set_udev_files(data_files)
@@ -166,39 +160,33 @@ def get_data_files(name, version, fullname):  # pylint: disable=R0912
                            src=["init/suse/waagent"])
         else:
             # sles 12+ and openSUSE 13.2+ use systemd
-            set_systemd_files(data_files, dest='/usr/lib/systemd/system')
+            set_systemd_files(data_files, dest=systemd_dir_path)
     elif name == 'freebsd':
-        set_bin_files(data_files, dest="/usr/local/sbin")
         set_conf_files(data_files, src=["config/freebsd/waagent.conf"])
         set_freebsd_rc_files(data_files)
     elif name == 'openbsd':
-        set_bin_files(data_files, dest="/usr/local/sbin")
         set_conf_files(data_files, src=["config/openbsd/waagent.conf"])
         set_openbsd_rc_files(data_files)
     elif name == 'debian':
-        set_bin_files(data_files)
         set_conf_files(data_files, src=["config/debian/waagent.conf"])
         set_logrotate_files(data_files)
         set_udev_files(data_files, dest="/lib/udev/rules.d")
         if debian_has_systemd():
-            set_systemd_files(data_files)
+            set_systemd_files(data_files, dest=systemd_dir_path)
     elif name == 'iosxe':
-        set_bin_files(data_files)
         set_conf_files(data_files, src=["config/iosxe/waagent.conf"])
         set_logrotate_files(data_files)
         set_udev_files(data_files)
-        set_systemd_files(data_files, dest="/usr/lib/systemd/system")
+        set_systemd_files(data_files, dest=systemd_dir_path)
         if version.startswith("7.1"):
             # TODO this is a mitigation to systemctl bug on 7.1
             set_sysv_files(data_files)
     elif name == 'openwrt':
-        set_bin_files(data_files)
         set_conf_files(data_files)
         set_logrotate_files(data_files)
         set_sysv_files(data_files, dest='/etc/init.d', src=["init/openwrt/waagent"])
     else:
         # Use default setting
-        set_bin_files(data_files)
         set_conf_files(data_files)
         set_logrotate_files(data_files)
         set_udev_files(data_files)


Test patch:
diff --git a/tests/common/test_persist_firewall_rules.py b/tests/common/test_persist_firewall_rules.py
new file mode 100644
index 00000000..9379fc90
--- /dev/null
+++ b/tests/common/test_persist_firewall_rules.py
@@ -0,0 +1,275 @@
+# Copyright 2016 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.6+ and Openssl 1.0+
+#
+import contextlib
+import os
+import shutil
+import subprocess
+import sys
+
+
+from azurelinuxagent.common.future import ustr
+from azurelinuxagent.common.osutil.default import DefaultOSUtil
+from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
+from azurelinuxagent.common.utils import fileutil
+from azurelinuxagent.common.utils.networkutil import AddFirewallRules, FirewallCmdDirectCommands
+from tests.tools import AgentTestCase, MagicMock, patch
+
+
+class TestPersistFirewallRulesHandler(AgentTestCase):
+
+    original_popen = subprocess.Popen
+
+    def __init__(self, *args, **kwargs):
+        super(TestPersistFirewallRulesHandler, self).__init__(*args, **kwargs)
+        self._expected_service_name = ""
+        self._expected_service_name = ""
+        self._drop_in_file = ""
+        self._network_service_unit_file = ""
+
+    def setUp(self):
+        AgentTestCase.setUp(self)
+        # Override for mocking Popen, should be of the form - (True/False, cmd-to-execute-if-True)
+        self.__replace_popen_cmd = lambda *_: (False, "")
+        self.__executed_commands = []
+        self.__test_dst_ip = "1.2.3.4"
+        self.__test_uid = 9999
+        self.__test_wait = "-w"
+
+        self.__systemd_dir = os.path.join(self.tmp_dir, "system")
+        fileutil.mkdir(self.__systemd_dir)
+        self.__agent_bin_dir = os.path.join(self.tmp_dir, "bin")
+        fileutil.mkdir(self.__agent_bin_dir)
+
+    def tearDown(self):
+        shutil.rmtree(self.__systemd_dir, ignore_errors=True)
+        shutil.rmtree(self.__agent_bin_dir, ignore_errors=True)
+        AgentTestCase.tearDown(self)
+
+    def __mock_popen(self, cmd, *args, **kwargs):
+        self.__executed_commands.append(cmd)
+        replace_cmd, replace_with_command = self.__replace_popen_cmd(cmd)
+        if replace_cmd:
+            cmd = replace_with_command
+
+        return TestPersistFirewallRulesHandler.original_popen(cmd, *args, **kwargs)
+
+    @contextlib.contextmanager
+    def _get_persist_firewall_rules_handler(self, systemd=True):
+
+        osutil = DefaultOSUtil()
+        osutil.get_firewall_will_wait = MagicMock(return_value=self.__test_wait)
+        osutil.get_agent_bin_path = MagicMock(return_value=self.__agent_bin_dir)
+        osutil.get_systemd_unit_file_install_path = MagicMock(return_value=self.__systemd_dir)
+
+        # protected-access<W0212> Disabled: OK to access PersistFirewallRulesHandler._* from unit test for PersistFirewallRuleHandler
+        self._expected_service_name = PersistFirewallRulesHandler._AGENT_NETWORK_SETUP_NAME_FORMAT.format(  # pylint: disable=protected-access
+            osutil.get_service_name())
+
+        self._network_service_unit_file = os.path.join(self.__systemd_dir, self._expected_service_name)
+        self._drop_in_file = os.path.join(self.__systemd_dir, "{0}.d".format(self._expected_service_name),
+                                          PersistFirewallRulesHandler._DROP_IN_ENV_FILE_NAME)  # pylint: disable=protected-access
+
+        # Just for these tests, ignoring the mode of mkdir to allow non-sudo tests
+        orig_mkdir = fileutil.mkdir
+        with patch("azurelinuxagent.common.persist_firewall_rules.fileutil.mkdir",
+                   side_effect=lambda path, **mode: orig_mkdir(path)):
+            with patch("azurelinuxagent.common.persist_firewall_rules.get_osutil", return_value=osutil):
+                with patch('azurelinuxagent.common.cgroupapi.CGroupsApi.is_systemd', return_value=systemd):
+                    with patch("azurelinuxagent.common.utils.shellutil.subprocess.Popen", side_effect=self.__mock_popen):
+                        yield PersistFirewallRulesHandler(self.__test_dst_ip, self.__test_uid)
+
+    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.__executed_commands, "Firewall {0} command not found".format(cmd))
+            self.assertIn(AddFirewallRules.get_firewalld_drop_command(command=cmd,
+                                                                      destination=self.__test_dst_ip),
+                          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.__executed_commands,
+                             "Firewall {0} command found".format(cmd))
+            self.assertNotIn(AddFirewallRules.get_firewalld_drop_command(command=cmd,
+                                                                         destination=self.__test_dst_ip),
+                             self.__executed_commands, "Firewall {0} command found".format(cmd))
+
+    def __assert_systemctl_called(self, cmd="enable", validate_command_called=True):
+        systemctl_command = ["systemctl", cmd, self._expected_service_name]
+        if validate_command_called:
+            self.assertIn(systemctl_command, self.__executed_commands, "Systemctl command {0} not found".format(cmd))
+        else:
+            self.assertNotIn(systemctl_command, self.__executed_commands, "Systemctl command {0} found".format(cmd))
+
+    def __assert_firewall_cmd_running_called(self, validate_command_called=True):
+        # protected-access<W0212> Disabled: OK to access PersistFirewallRulesHandler._* from unit test for PersistFirewallRuleHandler
+        cmd = PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD    # pylint: disable=protected-access
+        if validate_command_called:
+            self.assertIn(cmd, self.__executed_commands, "Firewall state not checked")
+        else:
+            self.assertNotIn(cmd, self.__executed_commands, "Firewall state not checked")
+
+    def __assert_network_service_setup_properly(self):
+        self.__assert_systemctl_called(cmd="is-enabled", validate_command_called=True)
+        self.__assert_systemctl_called(cmd="enable", validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=False)
+        self.assertTrue(os.path.exists(self._network_service_unit_file), "Service unit file should be there")
+        self.assertTrue(os.path.exists(self._drop_in_file), "Drop in file should be there")
+
+    @staticmethod
+    def __mock_network_setup_service_enabled(cmd):
+        if "firewall-cmd" in cmd:
+            return True, ["echo", "not-running"]
+        if "systemctl" in cmd:
+            return True, ["echo", "enabled"]
+        return False, []
+
+    @staticmethod
+    def __mock_network_setup_service_disabled(cmd):
+        if "firewall-cmd" in cmd:
+            return True, ["echo", "not-running"]
+        if "systemctl" in cmd:
+            return True, ["echo", "not enabled"]
+        return False, []
+
+    def test_it_should_skip_setup_if_firewalld_already_enabled(self):
+        self.__replace_popen_cmd = lambda cmd: ("firewall-cmd" in cmd, ["echo", "running"])
+        with self._get_persist_firewall_rules_handler() as handler:
+            handler.setup()
+
+        # 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.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")
+
+    def test_it_should_skip_setup_if_agent_network_setup_service_already_enabled(self):
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_enabled
+        with self._get_persist_firewall_rules_handler() as handler:
+            handler.setup()
+
+        self.__assert_systemctl_called(cmd="is-enabled", validate_command_called=True)
+        self.__assert_systemctl_called(cmd="enabled", 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.PassThrough, validate_command_called=False)
+
+    def test_it_should_always_replace_only_drop_in_file_if_using_custom_network_service(self):
+        test_str = 'Environment="EGG={0}" "DST_IP={1}" "UID={2}" "WAIT={3}"'
+        current_exe_path = os.path.join(os.getcwd(), sys.argv[0])
+
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_disabled
+        with self._get_persist_firewall_rules_handler() as handler:
+            self.assertFalse(os.path.exists(self._drop_in_file), "Drop in file should not be there")
+            self.assertFalse(os.path.exists(self._network_service_unit_file), "Unit file should not be present")
+            handler.setup()
+
+        self.__assert_systemctl_called(cmd="is-enabled", validate_command_called=True)
+        self.__assert_firewall_called(cmd=FirewallCmdDirectCommands.PassThrough, validate_command_called=False)
+        self.assertTrue(os.path.exists(self._drop_in_file), "Drop in file should be there")
+        self.assertTrue(fileutil.findstr_in_file(self._drop_in_file,
+                                                 test_str.format(current_exe_path, self.__test_dst_ip, self.__test_uid,
+                                                                 self.__test_wait)), "DropIn file not set correctly")
+        self.assertTrue(fileutil.findstr_in_file(self._network_service_unit_file,
+                                                 test_str.format(current_exe_path, self.__test_dst_ip, self.__test_uid,
+                                                                 self.__test_wait)),
+                        "Service Unit file file not set correctly")
+
+        # Change test params
+        self.__test_dst_ip = "9.8.7.6"
+        self.__test_uid = 5555
+        self.__test_wait = ""
+        # The service should say its enabled now
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_enabled
+        with self._get_persist_firewall_rules_handler() as handler:
+            # The Drop-in file should be available on the 2nd run
+            self.assertTrue(os.path.exists(self._drop_in_file), "Drop in file should be there")
+            handler.setup()
+
+        self.assertTrue(fileutil.findstr_in_file(self._drop_in_file,
+                                                 test_str.format(current_exe_path, self.__test_dst_ip, self.__test_uid,
+                                                                 self.__test_wait)), "DropIn file not set correctly")
+        # Unit file should NOT be updated
+        self.assertFalse(fileutil.findstr_in_file(self._network_service_unit_file,
+                                                  test_str.format(current_exe_path, self.__test_dst_ip, self.__test_uid,
+                                                                  self.__test_wait)),
+                         "Service Unit file file should not be updated")
+
+    def test_it_should_use_firewalld_if_available(self):
+
+        def __mock_firewalld_running_and_not_applied(cmd):
+            # protected-access<W0212> Disabled: OK to access PersistFirewallRulesHandler._* from unit test for PersistFirewallRuleHandler
+            if cmd == PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD:   # pylint: disable=protected-access
+                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"]
+            if "firewall-cmd" in cmd:
+                return True, ["echo", "enabled"]
+            return False, []
+
+        self.__replace_popen_cmd = __mock_firewalld_running_and_not_applied
+        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.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_set_up_custom_service_if_no_firewalld(self):
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_disabled
+        with self._get_persist_firewall_rules_handler() as handler:
+            self.assertFalse(os.path.exists(self._network_service_unit_file), "Service unit file should not be there")
+            self.assertFalse(os.path.exists(self._drop_in_file), "Drop in file should not be there")
+            handler.setup()
+
+        self.__assert_network_service_setup_properly()
+
+    def test_it_should_cleanup_files_on_error(self):
+
+        orig_write_file = fileutil.write_file
+        files_to_fail = []
+
+        def mock_write_file(path, _, *__):
+            if files_to_fail[0] in path:
+                raise IOError("Invalid file: {0}".format(path))
+            return orig_write_file(path, _, *__)
+
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_disabled
+        with self._get_persist_firewall_rules_handler() as handler:
+            test_files = [self._drop_in_file, self._network_service_unit_file]
+            for file_to_fail in test_files:
+                files_to_fail = [file_to_fail]
+                with patch("azurelinuxagent.common.persist_firewall_rules.fileutil.write_file",
+                           side_effect=mock_write_file):
+                    with self.assertRaises(Exception) as context_manager:
+                        handler.setup()
+                    self.assertIn("Invalid file: {0}".format(file_to_fail), ustr(context_manager.exception))
+                    self.assertFalse(os.path.exists(file_to_fail), "File should be deleted: {0}".format(file_to_fail))
+
+                # Cleanup remaining files for test clarity
+                for test_file in test_files:
+                    try:
+                        os.remove(test_file)
+                    except Exception:
+                        pass
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 979c415f..5b91fd16 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -9,6 +9,7 @@ import json
 import os
 import shutil
 import stat
+import subprocess
 import sys
 import tempfile
 import time
@@ -24,6 +25,7 @@ from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.event import EVENTS_DIRECTORY
 from azurelinuxagent.common.exception import ProtocolError, UpdateError, ResourceGoneError
 from azurelinuxagent.common.future import ustr
+from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
 from azurelinuxagent.common.protocol.goal_state import ExtensionsConfig
 from azurelinuxagent.common.protocol.hostplugin import URI_FORMAT_GET_API_VERSIONS, HOST_PLUGIN_PORT, \
     URI_FORMAT_GET_EXTENSION_ARTIFACT, HostPluginProtocol
@@ -33,6 +35,7 @@ from azurelinuxagent.common.protocol.util import ProtocolUtil
 from azurelinuxagent.common.protocol.wire import WireProtocol
 from azurelinuxagent.common.utils import fileutil, restutil, textutil
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
+from azurelinuxagent.common.utils.networkutil import FirewallCmdDirectCommands
 from azurelinuxagent.common.version import AGENT_PKG_GLOB, AGENT_DIR_GLOB, AGENT_NAME, AGENT_DIR_PATTERN, \
     AGENT_VERSION, CURRENT_AGENT, CURRENT_VERSION
 from azurelinuxagent.ga.exthandlers import ExtHandlerInstance, HandlerEnvironment
@@ -1621,6 +1624,30 @@ Description=Slice for Azure VM Extensions""")
             self.assertIn(HandlerEnvironment.eventsFolder, content[0][HandlerEnvironment.handlerEnvironment],
                           "{0} not found in HandlerEnv file".format(HandlerEnvironment.eventsFolder))
 
+    def test_it_should_setup_firewall_rules_on_startup(self):
+        iterations = 1
+        original_popen = subprocess.Popen
+        executed_commands = []
+
+        def _mock_popen(cmd, *args, **kwargs):
+            if 'firewall-cmd' in cmd:
+                executed_commands.append(cmd)
+                cmd = ["echo", "running"]
+            return original_popen(cmd, *args, **kwargs)
+
+        with self._get_update_handler(iterations) as (update_handler, _):
+            with patch("azurelinuxagent.common.utils.shellutil.subprocess.Popen", side_effect=_mock_popen):
+                update_handler.run(debug=True)
+
+        # Firewall-cmd should only be called 3 times - 1st to check if running, 2nd & 3rd for the QueryPassThrough cmd
+        self.assertEqual(3, len(executed_commands), "The number of times firwall-cmd should be called is only 3")
+        # protected-access<W0212> Disabled: OK to access PersistFirewallRulesHandler._* from unit test for PersistFirewallRuleHandler
+        self.assertEqual(PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD, executed_commands.pop(0),  # pylint: disable=protected-access
+                         "First command should be to check if firewalld is running")
+        self.assertTrue([FirewallCmdDirectCommands.QueryPassThrough in cmd for cmd in executed_commands],
+                        "The remaining commands should only be for querying the firewall commands")
+
+
     @contextlib.contextmanager
     def _setup_test_for_ext_event_dirs_retention(self):
         try:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 854850df..d74eaf7e 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -17,7 +17,7 @@
 
 import os.path
 
-from azurelinuxagent.agent import parse_args, Agent, usage
+from azurelinuxagent.agent import parse_args, Agent, usage, AgentCommands
 from azurelinuxagent.common import conf
 from azurelinuxagent.common.utils import fileutil
 from tests.tools import AgentTestCase, data_dir, Mock, patch
@@ -84,13 +84,13 @@ class TestAgent(AgentTestCase):
 
     def test_accepts_configuration_path(self):
         conf_path = os.path.join(data_dir, "test_waagent.conf")
-        c, f, v, d, cfp, lcm = parse_args(["-configuration-path:" + conf_path])  # pylint: disable=unused-variable
+        c, f, v, d, cfp, lcm, _ = parse_args(["-configuration-path:" + conf_path])  # pylint: disable=unused-variable
         self.assertEqual(cfp, conf_path)
 
     @patch("os.path.exists", return_value=True)
     def test_checks_configuration_path(self, mock_exists):
         conf_path = "/foo/bar-baz/something.conf"
-        c, f, v, d, cfp, lcm = parse_args(["-configuration-path:"+conf_path])  # pylint: disable=unused-variable
+        c, f, v, d, cfp, lcm, _ = parse_args(["-configuration-path:"+conf_path])  # pylint: disable=unused-variable
         self.assertEqual(cfp, conf_path)
         self.assertEqual(mock_exists.call_count, 1)
 
@@ -99,12 +99,12 @@ class TestAgent(AgentTestCase):
     @patch("sys.exit", side_effect=Exception)
     def test_rejects_missing_configuration_path(self, mock_exit, mock_exists, mock_stderr):  # pylint: disable=unused-argument
         try:
-            c, f, v, d, cfp, lcm = parse_args(["-configuration-path:/foo/bar.conf"])  # pylint: disable=unused-variable
+            c, f, v, d, cfp, lcm, _ = parse_args(["-configuration-path:/foo/bar.conf"])  # pylint: disable=unused-variable
         except Exception:
             self.assertEqual(mock_exit.call_count, 1)
 
     def test_configuration_path_defaults_to_none(self):
-        c, f, v, d, cfp, lcm = parse_args([])  # pylint: disable=unused-variable
+        c, f, v, d, cfp, lcm, _ = parse_args([])  # pylint: disable=unused-variable
         self.assertEqual(cfp, None)
 
     def test_agent_accepts_configuration_path(self):
@@ -181,12 +181,12 @@ class TestAgent(AgentTestCase):
 
     def test_checks_log_collector_mode(self):
         # Specify full mode
-        c, f, v, d, cfp, lcm = parse_args(["-collect-logs", "-full"])  # pylint: disable=unused-variable
+        c, f, v, d, cfp, lcm, _ = parse_args(["-collect-logs", "-full"])  # pylint: disable=unused-variable
         self.assertEqual(c, "collect-logs")
         self.assertEqual(lcm, True)
 
         # Defaults to None if mode not specified
-        c, f, v, d, cfp, lcm = parse_args(["-collect-logs"])  # pylint: disable=unused-variable
+        c, f, v, d, cfp, lcm, _ = parse_args(["-collect-logs"])  # pylint: disable=unused-variable
         self.assertEqual(c, "collect-logs")
         self.assertEqual(lcm, False)
 
@@ -194,7 +194,7 @@ class TestAgent(AgentTestCase):
     @patch("sys.exit", side_effect=Exception)
     def test_rejects_invalid_log_collector_mode(self, mock_exit, mock_stderr):  # pylint: disable=unused-argument
         try:
-            c, f, v, d, cfp, lcm = parse_args(["-collect-logs", "-notvalid"])  # pylint: disable=unused-variable
+            c, f, v, d, cfp, lcm, _ = parse_args(["-collect-logs", "-notvalid"])  # pylint: disable=unused-variable
         except Exception:
             self.assertEqual(mock_exit.call_count, 1)
 
@@ -211,6 +211,42 @@ class TestAgent(AgentTestCase):
         full_mode = mock_log_collector.call_args_list[1][0][0]
         self.assertFalse(full_mode)
 
+    def test_it_should_parse_setup_firewall_properly(self):
+
+        test_firewall_meta = {
+            "dst_ip": "1.2.3.4",
+            "uid": "9999",
+            "wait": "-w"
+        }
+        cmd, _, _, _, _, _, firewall_metadata = parse_args(
+            ["-{0}".format(AgentCommands.SetupFirewall), "-dst_ip=1.2.3.4", "-uid=9999", "-w"])
+
+        self.assertEqual(cmd, AgentCommands.SetupFirewall)
+        self.assertEqual(firewall_metadata, test_firewall_meta)
+
+        # Defaults to None if command is different
+        test_firewall_meta = {
+            "dst_ip": None,
+            "uid": None,
+            "wait": ""
+        }
+        cmd, _, _, _, _, _, firewall_metadata = parse_args(["-{0}".format(AgentCommands.Help)])
+        self.assertEqual(cmd, AgentCommands.Help)
+        self.assertEqual(test_firewall_meta, firewall_metadata)
+
+    def test_it_should_ignore_empty_arguments(self):
+
+        test_firewall_meta = {
+            "dst_ip": "1.2.3.4",
+            "uid": "9999",
+            "wait": ""
+        }
+        cmd, _, _, _, _, _, firewall_metadata = parse_args(
+            ["-{0}".format(AgentCommands.SetupFirewall), "-dst_ip=1.2.3.4", "-uid=9999", ""])
+
+        self.assertEqual(cmd, AgentCommands.SetupFirewall)
+        self.assertEqual(firewall_metadata, test_firewall_meta)
+
     def test_agent_usage_message(self):
         message = usage()