add custom options to sshdriver

added the possibility to define custom option for the SSHDriver in the environment yaml

Signed-off-by: Nikolaj Rahbel <[email protected]>

Author: b2vn

CHANGES.rst

@@ -35,6 +35,7 @@ New Features in 0.5.0
   forwarding to/from the exporter.
 - The SSH connection timeout can now be globally controlled using the
   ``LG_SSH_CONNECT_TIMEOUT`` environment variable.
+- Possible to define extra option for ssh connection
 
 Bug fixes in 0.5.0
 ~~~~~~~~~~~~~~~~~~

doc/configuration.rst

@@ -1483,6 +1483,10 @@ Implements:
 
    SSHDriver:
      keyfile: example.key
+     extra_options:
+       - "-o option1"
+       - "-o options2"
+       - "-4"
 
 Arguments:
   - keyfile (str): optional, filename of private key to login into the remote system
@@ -1491,6 +1495,8 @@ Arguments:
       stdout, and an empty list as second element.
   - connection_timeout (float, default=30.0): timeout when trying to establish connection to
       target.
+  - options (list or str): optional, a list of extra options required for the ssh
+    connection, e.g. defining host key algorithms.
 
 UBootDriver
 ~~~~~~~~~~~

labgrid/driver/sshdriver.py

@@ -30,6 +30,7 @@ class SSHDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol):
     bindings = {"networkservice": "NetworkService", }
     priorities = {CommandProtocol: 10, FileTransferProtocol: 10}
     keyfile = attr.ib(default="", validator=attr.validators.instance_of(str))
+    extra_options = attr.ib(default="", validator=attr.validators.instance_of((str, list)))
     stderr_merge = attr.ib(default=False, validator=attr.validators.instance_of(bool))
     connection_timeout = attr.ib(default=float(get_ssh_connect_timeout()), validator=attr.validators.instance_of(float))
 
@@ -48,6 +49,12 @@ def on_activate(self):
         if not self.networkservice.password:
             self.ssh_prefix += ["-o", "PasswordAuthentication=no"]
 
+        if isinstance(self.extra_options, str) and len(self.extra_options) > 0:
+            self.ssh_prefix += shlex.split(self.extra_options)
+        elif isinstance(self.extra_options, list) and len(self.extra_options) > 0:
+            for eo in self.extra_options:
+                self.ssh_prefix += shlex.split(eo)
+
         self.control = self._start_own_master()
         self.ssh_prefix += ["-F", "none"]
         if self.control:

tests/test_sshdriver.py

@@ -38,6 +38,46 @@ def test_create(target, mocker):
     s = SSHDriver(target, "ssh")
     assert isinstance(s, SSHDriver)
 
+def test_can_define_ssh_options(target, mocker):
+    NetworkService(target, "service", "1.2.3.4", "root")
+    call = mocker.patch('subprocess.call')
+    call.return_value = 0
+    popen = mocker.patch('subprocess.Popen', autospec=True)
+    path = mocker.patch('os.path.exists')
+    path.return_value = True
+    instance_mock = mocker.MagicMock()
+    popen.return_value = instance_mock
+    instance_mock.wait = mocker.MagicMock(return_value=0)
+
+    tp = SSHDriver(target, "ssh")
+    tp.extra_options = "-o HostKeyAlgorithms=+ssh-rsa"
+    s = target.get_driver("SSHDriver")
+    assert "HostKeyAlgorithms=+ssh-rsa" in s.ssh_prefix
+
+def test_extra_options_can_be_a_list(target, mocker):
+    NetworkService(target, "service", "1.2.3.4", "root")
+    call = mocker.patch('subprocess.call')
+    call.return_value = 0
+    popen = mocker.patch('subprocess.Popen', autospec=True)
+    path = mocker.patch('os.path.exists')
+    path.return_value = True
+    instance_mock = mocker.MagicMock()
+    popen.return_value = instance_mock
+    instance_mock.wait = mocker.MagicMock(return_value=0)
+
+    tp = SSHDriver(target, "ssh")
+    tp.extra_options = [
+        "-o HostKeyAlgorithms=+ssh-rsa",
+        "-o HostKeyAlgorithms=+ssh-dsa",
+    ]
+    s = target.get_driver("SSHDriver")
+    assert "HostKeyAlgorithms=+ssh-rsa" in s.ssh_prefix
+    assert "HostKeyAlgorithms=+ssh-dsa" in s.ssh_prefix
+
+def test_no_options_will_not_result_in_none_prefix(ssh_driver_mocked_and_activated, mocker):
+    s = ssh_driver_mocked_and_activated
+    assert not '' in s.ssh_prefix
+
 def test_run_check(ssh_driver_mocked_and_activated, mocker):
     s = ssh_driver_mocked_and_activated
     s._run = mocker.MagicMock(return_value=(['success'], [], 0))