Repository: Azure/WALinuxAgent
Pull request #2225 (2021-04-19 23:04:59)

Use a binary file for setting up firewall rules instead of directly using service files
<!-- 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).


Patch:
diff --git a/azurelinuxagent/common/persist_firewall_rules.py b/azurelinuxagent/common/persist_firewall_rules.py
index 45611342..44a801c4 100644
--- a/azurelinuxagent/common/persist_firewall_rules.py
+++ b/azurelinuxagent/common/persist_firewall_rules.py
@@ -18,6 +18,7 @@
 import os
 import sys
 
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.event import add_event, WALAEventOperation
 from azurelinuxagent.common.future import ustr
@@ -40,25 +41,28 @@ 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}}
+ExecStart={py_path} {binary_path}
 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}"
+    __BINARY_CONTENTS = """
+# This python file was created by the Azure VM Agent. Please do not edit.
+
+import os 
+
+
+if __name__ == '__main__':
+    if os.path.exists("{egg_path}"):
+        os.system("{py_path} {egg_path} --setup-firewall --dst_ip={wire_ip} --uid={user_id} {wait}")
+    else:
+        print("{egg_path} file not found, skipping execution of firewall execution setup for this boot")
 """
 
     _AGENT_NETWORK_SETUP_NAME_FORMAT = "{0}-network-setup.service"
-    _DROP_IN_ENV_FILE_NAME = "10-environment.conf"
+    BINARY_FILE_NAME = "waagent-network-setup.py"
 
     _FIREWALLD_RUNNING_CMD = ["firewall-cmd", "--state"]
 
@@ -113,9 +117,6 @@ Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
         logger.info(
             "Firewalld service not running/unavailable, trying to set up {0}".format(self._network_setup_service_name))
 
-        if not systemd.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):
@@ -155,6 +156,12 @@ Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
         return False
 
     def _setup_network_setup_service(self):
+        # Even if service is enabled, we need to overwrite the binary 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.__setup_binary_file()
+
         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()
@@ -163,39 +170,34 @@ Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
             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()
+            # Reload systemd configurations when we setup the service for the first time to avoid systemctl warnings
+            self.__reload_systemd_conf()
             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)
+    def __setup_binary_file(self):
+        binary_file_path = os.path.join(conf.get_lib_dir(), self.BINARY_FILE_NAME)
         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))
+            fileutil.write_file(binary_file_path,
+                                self.__BINARY_CONTENTS.format(egg_path=self._current_agent_executable_path,
+                                                              wire_ip=self._dst_ip,
+                                                              user_id=self._uid,
+                                                              wait=self._wait,
+                                                              py_path=sys.executable))
+            logger.info("Successfully updated the Binary file {0} for firewall setup".format(binary_file_path))
         except Exception:
-            self.__remove_file_without_raising(drop_in_file)
+            logger.warn(
+                "Unable to setup binary file, removing the service unit file {0} to ensure its not run on system reboot".format(
+                    self.get_service_file_path()))
+            self.__remove_file_without_raising(binary_file_path)
+            self.__remove_file_without_raising(self.get_service_file_path())
             raise
 
     def __set_service_unit_file(self):
         service_unit_file = self.get_service_file_path()
+        binary_path = os.path.join(conf.get_lib_dir(), self.BINARY_FILE_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,
+                                self.__SERVICE_FILE_CONTENT.format(binary_path=binary_path,
                                                                    py_path=sys.executable))
             fileutil.chmod(service_unit_file, 0o644)
 
@@ -263,3 +265,10 @@ Environment="EGG={egg_path}" "DST_IP={wire_ip}" "UID={user_id}" "WAIT={wait}"
             is_success=(not service_failed),
             message=msg,
             log_event=False)
+
+    def __reload_systemd_conf(self):
+        try:
+            logger.info("Executing systemctl daemon-reload for setting up {0}".format(self._network_setup_service_name))
+            shellutil.run_command(["systemctl", "daemon-reload"])
+        except Exception as exception:
+            logger.warn("Unable to reload systemctl configurations: {0}".format(ustr(exception)))
diff --git a/azurelinuxagent/pa/deprovision/default.py b/azurelinuxagent/pa/deprovision/default.py
index a33fcf28..105b6182 100644
--- a/azurelinuxagent/pa/deprovision/default.py
+++ b/azurelinuxagent/pa/deprovision/default.py
@@ -263,5 +263,6 @@ class DeprovisionHandler(object):
     @staticmethod
     def del_persist_firewall_rules(actions):
         agent_network_service_path = PersistFirewallRulesHandler.get_service_file_path()
-        actions.append(DeprovisionAction(fileutil.rm_files, [agent_network_service_path]))
-        actions.append(DeprovisionAction(fileutil.rm_dirs, ["{0}.d".format(agent_network_service_path)]))
+        actions.append(DeprovisionAction(fileutil.rm_files,
+                                         [agent_network_service_path, os.path.join(conf.get_lib_dir(),
+                                          PersistFirewallRulesHandler.BINARY_FILE_NAME)]))


Test patch:
diff --git a/tests/common/test_persist_firewall_rules.py b/tests/common/test_persist_firewall_rules.py
index a5a0b8b6..a5d832a7 100644
--- a/tests/common/test_persist_firewall_rules.py
+++ b/tests/common/test_persist_firewall_rules.py
@@ -20,11 +20,13 @@ import shutil
 import subprocess
 import sys
 
+import uuid
 
+import azurelinuxagent.common.conf as conf
 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 import fileutil, shellutil
 from azurelinuxagent.common.utils.networkutil import AddFirewallRules, FirewallCmdDirectCommands
 from tests.tools import AgentTestCase, MagicMock, patch
 
@@ -37,7 +39,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         super(TestPersistFirewallRulesHandler, self).__init__(*args, **kwargs)
         self._expected_service_name = ""
         self._expected_service_name = ""
-        self._drop_in_file = ""
+        self._binary_file = ""
         self._network_service_unit_file = ""
 
     def setUp(self):
@@ -54,9 +56,14 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         self.__agent_bin_dir = os.path.join(self.tmp_dir, "bin")
         fileutil.mkdir(self.__agent_bin_dir)
 
+        self.__tmp_conf_lib = os.path.join(self.tmp_dir, "waagent")
+        fileutil.mkdir(self.__tmp_conf_lib)
+        conf.get_lib_dir = MagicMock(return_value=self.__tmp_conf_lib)
+
     def tearDown(self):
         shutil.rmtree(self.__systemd_dir, ignore_errors=True)
         shutil.rmtree(self.__agent_bin_dir, ignore_errors=True)
+        shutil.rmtree(self.__tmp_conf_lib, ignore_errors=True)
         AgentTestCase.tearDown(self)
 
     def __mock_popen(self, cmd, *args, **kwargs):
@@ -79,8 +86,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
             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)
+        self._binary_file = os.path.join(conf.get_lib_dir(), PersistFirewallRulesHandler.BINARY_FILE_NAME)
 
         # Just for these tests, ignoring the mode of mkdir to allow non-sudo tests
         orig_mkdir = fileutil.mkdir
@@ -117,6 +123,13 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         else:
             self.assertNotIn(systemctl_command, self.__executed_commands, "Systemctl command {0} found".format(cmd))
 
+    def __assert_systemctl_reloaded(self, validate_command_called=True):
+        systemctl_reload = ["systemctl", "daemon-reload"]
+        if validate_command_called:
+            self.assertIn(systemctl_reload, self.__executed_commands, "Systemctl config not reloaded")
+        else:
+            self.assertNotIn(systemctl_reload, self.__executed_commands, "Systemctl config reloaded")
+
     def __assert_firewall_cmd_running_called(self, validate_command_called=True):
         cmd = PersistFirewallRulesHandler._FIREWALLD_RUNNING_CMD
         if validate_command_called:
@@ -127,9 +140,10 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
     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_systemctl_reloaded()
         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")
+        self.assertTrue(os.path.exists(self._binary_file), "Binary file should be there")
 
     @staticmethod
     def __mock_network_setup_service_enabled(cmd):
@@ -166,29 +180,43 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
 
         self.__assert_systemctl_called(cmd="is-enabled", validate_command_called=True)
         self.__assert_systemctl_called(cmd="enabled", validate_command_called=False)
+        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.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}"'
+
+        def _find_in_file(file_name, line_str):
+            try:
+                with open(file_name, 'r') as fh:
+                    content = fh.read()
+                    return line_str in content
+            except Exception:
+                # swallow exception
+                pass
+            return False
+
+        test_str = 'os.system("{py_path} {egg_path} --setup-firewall --dst_ip={wire_ip} --uid={user_id} {wait}")'
         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._binary_file), "Binary file should not be there")
             self.assertFalse(os.path.exists(self._network_service_unit_file), "Unit file should not be present")
             handler.setup()
 
+        orig_service_file_contents = "ExecStart={py_path} {binary_path}".format(py_path=sys.executable,
+                                                                                 binary_path=self._binary_file)
         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)),
+        self.assertTrue(os.path.exists(self._binary_file), "Binary file should be there")
+        self.assertTrue(_find_in_file(self._binary_file,
+                                      test_str.format(py_path=sys.executable, egg_path=current_exe_path,
+                                                      wire_ip=self.__test_dst_ip, user_id=self.__test_uid,
+                                                      wait=self.__test_wait)),
+                        "Binary file not set correctly")
+        self.assertTrue(_find_in_file(self._network_service_unit_file, orig_service_file_contents),
                         "Service Unit file file not set correctly")
 
         # Change test params
@@ -199,17 +227,17 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         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")
+            self.assertTrue(os.path.exists(self._binary_file), "Binary 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")
+        self.assertTrue(_find_in_file(self._binary_file,
+                                      test_str.format(py_path=sys.executable, egg_path=current_exe_path,
+                                                      wire_ip=self.__test_dst_ip, user_id=self.__test_uid,
+                                                      wait=self.__test_wait)),
+                        "Binary file not updated 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")
+        self.assertTrue(_find_in_file(self._network_service_unit_file, orig_service_file_contents),
+                        "Service Unit file file should not be updated")
 
     def test_it_should_use_firewalld_if_available(self):
 
@@ -237,7 +265,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
         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")
+            self.assertFalse(os.path.exists(self._binary_file), "Binary file should not be there")
             handler.setup()
 
         self.__assert_network_service_setup_properly()
@@ -254,7 +282,7 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
 
         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]
+            test_files = [self._binary_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",
@@ -270,3 +298,26 @@ class TestPersistFirewallRulesHandler(AgentTestCase):
                         os.remove(test_file)
                     except Exception:
                         pass
+
+    def test_it_should_execute_binary_file_successfully(self):
+        # A bare-bone test to ensure no simple syntactical errors in the binary file as its generated dynamically
+        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._binary_file), "Binary file should not be there")
+            handler.setup()
+
+            self.assertTrue(os.path.exists(self._binary_file), "Binary file not set properly")
+
+            shellutil.run_command([sys.executable, self._binary_file])
+
+    def test_it_should_not_fail_if_egg_not_found(self):
+        self.__replace_popen_cmd = TestPersistFirewallRulesHandler.__mock_network_setup_service_disabled
+        test_str = str(uuid.uuid4())
+        with patch("sys.argv", [test_str]):
+            with self._get_persist_firewall_rules_handler() as handler:
+                self.assertFalse(os.path.exists(self._binary_file), "Binary file should not be there")
+                handler.setup()
+                output = shellutil.run_command([sys.executable, self._binary_file], stderr=subprocess.STDOUT)
+                expected_str = "{0} file not found, skipping execution of firewall execution setup for this boot".format(
+                    os.path.join(os.getcwd(), test_str))
+                self.assertIn(expected_str, output, "Unexpected output")