From 5902f05265756a49c4c6b8baa1c99fccc67a2554 Mon Sep 17 00:00:00 2001
From: Daniel Vincze <dvincze@cloudbasesolutions.com>
Date: Wed, 31 Jul 2024 17:43:57 +0300
Subject: [PATCH] Refactor and fix SSH service restart on OSMount worker SSH
 env configuration

On newer Ubuntu (starting with 24.04) Server images, `sshd` service name has
been changed to `ssh`. This patch handles this rename and refactors it into
the base OSMount class, in order for it to be picked up for all Linux OSMount
minions.
---
 coriolis/osmorphing/osmount/base.py            | 12 ++++++++++--
 coriolis/osmorphing/osmount/redhat.py          |  5 -----
 coriolis/osmorphing/osmount/ubuntu.py          |  5 -----
 coriolis/tests/osmorphing/osmount/test_base.py |  8 ++++++--
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/coriolis/osmorphing/osmount/base.py b/coriolis/osmorphing/osmount/base.py
index 0ce5d405..e521ff6c 100644
--- a/coriolis/osmorphing/osmount/base.py
+++ b/coriolis/osmorphing/osmount/base.py
@@ -17,7 +17,6 @@
 from coriolis import exception
 from coriolis import utils
 
-
 LOG = logging.getLogger(__name__)
 
 MAJOR_COLUMN_INDEX = 4
@@ -96,7 +95,16 @@ def setup(self):
             self._connect()
 
     def _allow_ssh_env_vars(self):
-        pass
+        self._exec_cmd('sudo sed -i -e "\$aAcceptEnv *" /etc/ssh/sshd_config')
+        try:
+            utils.restart_service(self._ssh, "sshd")
+        except exception.CoriolisException:
+            # Newer Ubuntu distros renamed this service to `ssh`
+            LOG.warning(
+                "Could not restart service sshd. Attempting to restart ssh "
+                "service")
+            utils.restart_service(self._ssh, 'ssh')
+        return True
 
     def _exec_cmd(self, cmd, timeout=None):
         if not timeout:
diff --git a/coriolis/osmorphing/osmount/redhat.py b/coriolis/osmorphing/osmount/redhat.py
index 211ec46a..4d9e2ffd 100644
--- a/coriolis/osmorphing/osmount/redhat.py
+++ b/coriolis/osmorphing/osmount/redhat.py
@@ -23,8 +23,3 @@ def setup(self):
         self._exec_cmd("sudo -E yum install -y lvm2 psmisc")
         self._exec_cmd("sudo modprobe dm-mod")
         self._exec_cmd("sudo rm -f /etc/lvm/devices/system.devices")
-
-    def _allow_ssh_env_vars(self):
-        self._exec_cmd('sudo sed -i -e "\$aAcceptEnv *" /etc/ssh/sshd_config')
-        utils.restart_service(self._ssh, "sshd")
-        return True
diff --git a/coriolis/osmorphing/osmount/ubuntu.py b/coriolis/osmorphing/osmount/ubuntu.py
index fc8c32c6..ad7b3ec6 100644
--- a/coriolis/osmorphing/osmount/ubuntu.py
+++ b/coriolis/osmorphing/osmount/ubuntu.py
@@ -36,8 +36,3 @@ def setup(self):
             "install lvm2 psmisc -y")
 
         self._exec_cmd("sudo modprobe dm-mod")
-
-    def _allow_ssh_env_vars(self):
-        self._exec_cmd('sudo sed -i -e "\$aAcceptEnv *" /etc/ssh/sshd_config')
-        utils.restart_service(self._ssh, "sshd")
-        return True
diff --git a/coriolis/tests/osmorphing/osmount/test_base.py b/coriolis/tests/osmorphing/osmount/test_base.py
index e9c727d5..eb8c3caa 100644
--- a/coriolis/tests/osmorphing/osmount/test_base.py
+++ b/coriolis/tests/osmorphing/osmount/test_base.py
@@ -100,9 +100,13 @@ def test__connect(self, mock_wait_for_port_connectivity, mock_ssh_client):
                 self.conn_info['ip'], 22)
         )
 
-    def test_setup(self):
+    @mock.patch.object(base.BaseSSHOSMountTools, '_allow_ssh_env_vars')
+    @mock.patch.object(base.BaseSSHOSMountTools, '_connect')
+    def test_setup(self, mock_connect, mock_allow_ssh_vars):
         self.base_os_mount_tools.setup()
-        self.ssh.close_assert_called_once()
+        mock_allow_ssh_vars.return_value = True
+        self.ssh.close.assert_called_once_with()
+        mock_connect.assert_called_once_with()
 
     @mock.patch.object(base.utils, 'exec_ssh_cmd')
     def test__exec_cmd(self, mock_exec_ssh_cmd):