Validate that process is actually a patching operation before terminating

src/extension/src/ProcessHandler.py

@@ -199,9 +199,27 @@ def is_process_running(self, pid):
                 # According to "man 2 kill" possible error values are (EINVAL, EPERM, ESRCH) Thus considering this as an error
                 return False
 
+    def is_process_patching_operation(self, pid):
+        try:
+            # Patching operation cmdline will look like:
+            # /usr/bin/python3.10 /var/lib/waagent/Microsoft.CPlat.Core.LinuxPatchExtension-1.6.69/MsftLinuxPatchCore.py <args>
+            with self.env_layer.file_system.open("/proc/{0}/cmdline".format(str(pid)), mode="rb") as cmdline_file:
+                cmdline = cmdline_file.read()
+                cmdline = cmdline.replace(b'\x00', b' ').decode("utf-8")  # cmdline is null-separated, so replacing nulls with spaces for easier parsing and logging
+
+                if Constants.CORE_CODE_FILE_NAME in cmdline:
+                    self.logger.log_verbose("Process is a patching operation. [PID={0}]".format(str(pid)))
+                    return True
+                else:
+                    self.logger.log_debug("Process is not a patching operation. [PID={0}] [CmdLine={1}]".format(str(pid), cmdline))
+                    return False
+        except Exception as error:
+            self.logger.log_debug("Error checking if process is a patching operation. [PID={0}] [Error={1}]".format(str(pid), repr(error)))
+            return True  # If we cannot determine, assume it is a patching operation to be safe
+
     def kill_process(self, pid):
         try:
-            if self.is_process_running(pid):
+            if self.is_process_running(pid) and self.is_process_patching_operation(pid):
                 self.logger.log("Terminating process: [PID={0}]".format(str(pid)))
                 os.kill(pid, signal.SIGTERM)
         except OSError as error:

src/extension/tests/Test_ProcessHandler.py

@@ -15,8 +15,10 @@
 # Requires Python 2.7+
 
 import os
+import shutil
 import subprocess
 import sys
+import tempfile
 import unittest
 from extension.src.Constants import Constants
 from extension.src.file_handlers.ExtOutputStatusHandler import ExtOutputStatusHandler
@@ -39,16 +41,22 @@ def setUp(self):
         self.json_file_handler = runtime.json_file_handler
         self.env_layer = runtime.env_layer
         dir_path = os.path.join(os.path.pardir, "tests", "helpers")
+        self.test_dir = tempfile.mkdtemp()
+        self.proc_cmdline_path = os.path.join(self.test_dir, "proc_cmdline")
         self.ext_output_status_handler = ExtOutputStatusHandler(self.logger, self.utility, self.json_file_handler, dir_path)
         self.process = subprocess.Popen(["echo", "Hello World!"], shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
     def tearDown(self):
         VirtualTerminal().print_lowlight("\n----------------- tear down test runner -----------------")
+        shutil.rmtree(self.test_dir)
         self.process.kill()
 
     def mock_is_process_running_to_return_true(self, pid):
         return True
 
+    def mock_is_process_patching_operation_to_return_true(self, pid):
+        return True
+
     def mock_os_kill_to_raise_exception(self, pid, sig):
         raise OSError
 
@@ -64,6 +72,12 @@ def mock_run_command_output_for_python_not_found(self, cmd, no_output=False, chk
     def mock_get_python_cmd(self):
         return "python"
 
+    def mock_file_system_open_to_return_proc_cmdline(self, path, mode):
+        return open(self.proc_cmdline_path, mode=mode)
+
+    def mock_file_system_open_raises_exception(self, path, mode):
+        raise FileNotFoundError
+
     def mock_run_command_to_set_auto_assess_shell_file_permission(self, cmd, no_output=False, chk_err=False):
         return 0, "permissions set"
 
@@ -133,6 +147,8 @@ def test_kill_process(self):
         # setting mocks
         is_process_running_backup = ProcessHandler.is_process_running
         ProcessHandler.is_process_running = self.mock_is_process_running_to_return_true
+        is_process_patching_operation_backup = ProcessHandler.is_process_patching_operation
+        ProcessHandler.is_process_patching_operation = self.mock_is_process_patching_operation_to_return_true
         os_kill_backup = os.kill
         os.kill = self.mock_os_kill_to_raise_exception
 
@@ -143,6 +159,7 @@ def test_kill_process(self):
 
         # reseting mocks
         ProcessHandler.is_process_running = is_process_running_backup
+        ProcessHandler.is_process_patching_operation = is_process_patching_operation_backup
         os.kill = os_kill_backup
 
     def test_get_python_cmd(self):
@@ -205,6 +222,42 @@ def test_start_daemon(self):
         subprocess.Popen = subprocess_popen_backup
         process_handler.env_layer.run_command_output = run_command_output_backup
         ExtEnvHandler.get_temp_folder = ext_env_handler_get_temp_folder_backup
+
+    def test_is_process_patching_operation(self):
+        # setting mocks
+        backup_file_system_open = self.env_layer.file_system.open
+        self.env_layer.file_system.open = self.mock_file_system_open_to_return_proc_cmdline
+
+        # process is patching operation
+        is_patching_operation_content = b'/usr/bin/python3.10\x00/var/lib/waagent/Microsoft.CPlat.Core.LinuxPatchExtension-1.6.69/MsftLinuxPatchCore.py\x00--seqno\x001234\x00--configfile\x00/var/lib/waagent/Microsoft.CPlat.Core.LinuxPatchExtension-1.6.69/config/1234/config.json\x00--envfile\x00/var/lib/waagent/Microsoft.CPlat.Core.LinuxPatchExtension-1.6.69/env/1234/env.json\x00--outputfile\x00/var/lib/waagent/Microsoft.CPlat.Core.LinuxPatchExtension-1.6.69/output/1234/output.json'
+        with open(self.proc_cmdline_path, "wb") as f:
+            f.write(is_patching_operation_content)
+
+        process_handler = ProcessHandler(self.logger, self.env_layer, self.ext_output_status_handler)
+        pid = 1234
+
+        self.assertTrue(process_handler.is_process_patching_operation(pid))
+
+        # process is not patching operation
+        is_not_patching_operation_content = b'/usr/bin/python3.10\x00/someother/path/some_other_script.py\x00--somearg\x00somevalue'
+        with open(self.proc_cmdline_path, "wb") as f:
+            f.write(is_not_patching_operation_content)
+
+        process_handler = ProcessHandler(self.logger, self.env_layer, self.ext_output_status_handler)
+        pid = 1234
+
+        self.assertFalse(process_handler.is_process_patching_operation(pid))
+
+        # process cmdline exception
+        self.env_layer.file_system.open = self.mock_file_system_open_raises_exception
+
+        process_handler = ProcessHandler(self.logger, self.env_layer, self.ext_output_status_handler)
+        pid = 1234
+
+        self.assertTrue(process_handler.is_process_patching_operation(pid))
+
+        # resetting mocks
+        self.env_layer.file_system.open = backup_file_system_open
 
 
 if __name__ == '__main__':