Move SSH and related to netutil

Author: guzu

jobs.py

@@ -5,7 +5,7 @@
 import subprocess
 import sys
 
-from lib.commands import ssh
+from lib.netutil import ssh
 
 JOBS = {
     "main": {

lib/commands.py

@@ -1,10 +1,8 @@
 import base64
 import logging
-import shlex
 import subprocess
 
 import lib.config as config
-from lib.common import HostAddress
 
 from typing import List, Literal, Union, overload
 
@@ -17,14 +15,6 @@ def __init__(self, returncode, stdout, cmd, exception_msg):
         self.stdout = stdout
         self.cmd = cmd
 
-class SSHCommandFailed(BaseCommandFailed):
-    def __init__(self, returncode, stdout, cmd):
-        msg_end = f": {stdout}" if stdout else "."
-        super(SSHCommandFailed, self).__init__(
-            returncode, stdout, cmd,
-            f'SSH command ({cmd}) failed with return code {returncode}{msg_end}'
-        )
-
 class LocalCommandFailed(BaseCommandFailed):
     def __init__(self, returncode, stderr, cmd):
         msg_end = f": {stderr}" if stderr else "."
@@ -40,10 +30,6 @@ def __init__(self, returncode, stdout):
         self.returncode = returncode
         self.stdout = stdout
 
-class SSHResult(BaseCmdResult):
-    def __init__(self, returncode, stdout):
-        super(SSHResult, self).__init__(returncode, stdout)
-
 class LocalCommandResult(BaseCmdResult):
     def __init__(self, returncode, stdout, stderr):
         super(LocalCommandResult, self).__init__(returncode, stdout)
@@ -62,191 +48,6 @@ def _ellide_log_lines(log):
         reduced_message.append("(...)")
     return "\n{}".format("\n".join(reduced_message))
 
-def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-         background, decode, options) -> Union[SSHResult, SSHCommandFailed, str, bytes, None]:
-    opts = list(options)
-    opts.append('-o "BatchMode yes"')
-    if suppress_fingerprint_warnings:
-        # Suppress warnings and questions related to host key fingerprints
-        # because on a test network IPs get reused, VMs are reinstalled, etc.
-        # Based on https://unix.stackexchange.com/a/365976/257493
-        opts.append('-o "StrictHostKeyChecking no"')
-        opts.append('-o "LogLevel ERROR"')
-        opts.append('-o "UserKnownHostsFile /dev/null"')
-
-    if isinstance(cmd, str):
-        command = cmd
-    else:
-        command = " ".join(cmd)
-
-    ssh_cmd = f"ssh root@{hostname_or_ip} {' '.join(opts)} {shlex.quote(command)}"
-
-    # Fetch banner and remove it to avoid stdout/stderr pollution.
-    banner_res = None
-    if config.ignore_ssh_banner:
-        banner_res = subprocess.run(
-            "ssh root@%s %s '%s'" % (hostname_or_ip, ' '.join(opts), '\n'),
-            shell=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            check=False
-        )
-
-    logging.debug(f"[{hostname_or_ip}] {command}")
-    process = subprocess.Popen(
-        ssh_cmd,
-        shell=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT
-    )
-    if background:
-        return None
-
-    stdout = []
-    assert process.stdout is not None
-    for line in iter(process.stdout.readline, b''):
-        readable_line = line.decode(errors='replace').strip()
-        stdout.append(line)
-        logging.debug("> %s", readable_line)
-    _, stderr = process.communicate()
-    res = subprocess.CompletedProcess(ssh_cmd, process.returncode, b''.join(stdout), stderr)
-
-    # Get a decoded version of the output in any case, replacing potential errors
-    output_for_errors = res.stdout.decode(errors='replace').strip()
-
-    # Even if check is False, we still raise in case of return code 255, which means a SSH error.
-    if res.returncode == 255:
-        return SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, command)
-
-    output: Union[bytes, str] = res.stdout
-    if banner_res:
-        if banner_res.returncode == 255:
-            return SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), command)
-        output = output[len(banner_res.stdout):]
-
-    if decode:
-        assert isinstance(output, bytes)
-        output = output.decode()
-
-    if res.returncode and check:
-        return SSHCommandFailed(res.returncode, output_for_errors, command)
-
-    if simple_output:
-        return output.strip()
-    return SSHResult(res.returncode, output)
-
-# The actual code is in _ssh().
-# This function is kept short for shorter pytest traces upon SSH failures, which are common,
-# as pytest prints the whole function definition that raised the SSHCommandFailed exception
-@overload
-def ssh(hostname_or_ip: HostAddress, cmd: Union[str, List[str]], *, check: bool = True,
-        simple_output: Literal[True] = True,
-        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-        decode: Literal[True] = True, options: List[str] = []) -> str:
-    ...
-@overload
-def ssh(hostname_or_ip: HostAddress, cmd: Union[str, List[str]], *, check: bool = True,
-        simple_output: Literal[True] = True,
-        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-        decode: Literal[False], options: List[str] = []) -> bytes:
-    ...
-@overload
-def ssh(hostname_or_ip: HostAddress, cmd: Union[str, List[str]], *, check: bool = True,
-        simple_output: Literal[False],
-        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-        decode: bool = True, options: List[str] = []) -> SSHResult:
-    ...
-@overload
-def ssh(hostname_or_ip: HostAddress, cmd: Union[str, List[str]], *, check: bool = True,
-        simple_output: Literal[False],
-        suppress_fingerprint_warnings: bool = True, background: Literal[True],
-        decode: bool = True, options: List[str] = []) -> None:
-    ...
-@overload
-def ssh(hostname_or_ip: HostAddress, cmd: Union[str, List[str]], *, check=True,
-        simple_output: bool = True,
-        suppress_fingerprint_warnings=True, background: bool = False,
-        decode: bool = True, options: List[str] = []) \
-        -> Union[str, bytes, SSHResult, None]:
-    ...
-def ssh(hostname_or_ip, cmd, *, check=True, simple_output=True,
-        suppress_fingerprint_warnings=True,
-        background=False, decode=True, options=[]):
-    result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-                         background, decode, options)
-    if isinstance(result_or_exc, SSHCommandFailed):
-        raise result_or_exc
-    else:
-        return result_or_exc
-
-def ssh_with_result(hostname_or_ip, cmd, suppress_fingerprint_warnings=True,
-                    background=False, decode=True, options=[]) -> SSHResult:
-    result_or_exc = _ssh(hostname_or_ip, cmd, False, False, suppress_fingerprint_warnings,
-                         background, decode, options)
-    if isinstance(result_or_exc, SSHCommandFailed):
-        raise result_or_exc
-    elif isinstance(result_or_exc, SSHResult):
-        return result_or_exc
-    assert False, "unexpected type"
-
-def scp(hostname_or_ip, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
-    # local import to avoid cyclic import; lib.netutils also import lib.commands
-    from lib.netutil import wrap_ip
-
-    opts = '-o "BatchMode yes"'
-    if suppress_fingerprint_warnings:
-        # Suppress warnings and questions related to host key fingerprints
-        # because on a test network IPs get reused, VMs are reinstalled, etc.
-        # Based on https://unix.stackexchange.com/a/365976/257493
-        opts = '-o "StrictHostKeyChecking no" -o "LogLevel ERROR" -o "UserKnownHostsFile /dev/null"'
-
-    ip = wrap_ip(hostname_or_ip)
-    if local_dest:
-        src = 'root@{}:{}'.format(ip, src)
-    else:
-        dest = 'root@{}:{}'.format(ip, dest)
-
-    command = "scp {} {} {}".format(opts, src, dest)
-    res = subprocess.run(
-        command,
-        shell=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        check=False
-    )
-
-    errorcode_msg = "" if res.returncode == 0 else " - Got error code: %s" % res.returncode
-    logging.debug(f"[{hostname_or_ip}] scp: {src} => {dest}{errorcode_msg}")
-
-    if check and res.returncode:
-        raise SSHCommandFailed(res.returncode, res.stdout.decode(), command)
-
-    return res
-
-def sftp(hostname_or_ip, cmds, check=True, suppress_fingerprint_warnings=True):
-    opts = ''
-    if suppress_fingerprint_warnings:
-        # Suppress warnings and questions related to host key fingerprints
-        # because on a test network IPs get reused, VMs are reinstalled, etc.
-        # Based on https://unix.stackexchange.com/a/365976/257493
-        opts = '-o "StrictHostKeyChecking no" -o "LogLevel ERROR" -o "UserKnownHostsFile /dev/null"'
-
-    args = "sftp {} -b - root@{}".format(opts, hostname_or_ip)
-    input = bytes("\n".join(cmds), 'utf-8')
-    res = subprocess.run(
-        args,
-        input=input,
-        shell=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        check=False
-    )
-
-    if check and res.returncode:
-        raise SSHCommandFailed(res.returncode, res.stdout.decode(), "{} -- {}".format(args, cmds))
-
-    return res
-
 @overload
 def local_cmd(cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
               decode: Literal[True] = True) -> str:

lib/common.py

@@ -15,7 +15,7 @@
 
 import requests
 
-import lib.commands as commands
+import lib.netutil as netutil
 
 from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, TypeAlias, TypeVar, Union, cast, overload
 
@@ -256,7 +256,7 @@ def _param_get(host: 'lib.host.Host', xe_prefix: str, uuid: str, param_name: str
         args['param-key'] = key
     try:
         value = host.xe(f'{xe_prefix}-param-get', args)
-    except commands.SSHCommandFailed as e:
+    except netutil.SSHCommandFailed as e:
         if key and accept_unknown_key and e.stdout == "Error: Key %s not found in map" % key:
             value = None
         else:
@@ -286,7 +286,7 @@ def _param_remove(host, xe_prefix, uuid, param_name, key, accept_unknown_key=Fal
     args = {'uuid': uuid, 'param-name': param_name, 'param-key': key}
     try:
         host.xe(f'{xe_prefix}-param-remove', args)
-    except commands.SSHCommandFailed as e:
+    except netutil.SSHCommandFailed as e:
         if not accept_unknown_key or e.stdout != "Error: Key %s not found in map" % key:
             raise
 

lib/fistpoint.py

@@ -1,7 +1,7 @@
 import logging
 
-from lib.commands import SSHCommandFailed
 from lib.host import Host
+from lib.netutil import SSHCommandFailed
 
 from typing import Final
 

lib/host.py

@@ -9,7 +9,7 @@
 
 from packaging import version
 
-import lib.commands as commands
+import lib.netutil as netutil
 import lib.pif as pif
 
 from typing import TYPE_CHECKING, Dict, List, Literal, Optional, TypedDict, Union, overload
@@ -104,7 +104,7 @@ def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output:
     @overload
     def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[False],
             suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-            decode: bool = True) -> commands.SSHResult:
+            decode: bool = True) -> netutil.SSHResult:
         ...
 
     @overload
@@ -116,21 +116,21 @@ def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output:
     @overload
     def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: bool = True,
             suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True) \
-            -> Union[str, bytes, commands.SSHResult, None]:
+            -> Union[str, bytes, netutil.SSHResult, None]:
         ...
 
     def ssh(self, cmd, *, check=True, simple_output=True, suppress_fingerprint_warnings=True,
             background=False, decode=True):
-        return commands.ssh(self.hostname_or_ip, cmd, check=check, simple_output=simple_output,
-                            suppress_fingerprint_warnings=suppress_fingerprint_warnings,
-                            background=background, decode=decode)
+        return netutil.ssh(self.hostname_or_ip, cmd, check=check, simple_output=simple_output,
+                           suppress_fingerprint_warnings=suppress_fingerprint_warnings,
+                           background=background, decode=decode)
 
-    def ssh_with_result(self, cmd) -> commands.SSHResult:
+    def ssh_with_result(self, cmd) -> netutil.SSHResult:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
-        return commands.ssh_with_result(self.hostname_or_ip, cmd)
+        return netutil.ssh_with_result(self.hostname_or_ip, cmd)
 
     def scp(self, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
-        return commands.scp(
+        return netutil.scp(
             self.hostname_or_ip, src, dest, check=check,
             suppress_fingerprint_warnings=suppress_fingerprint_warnings, local_dest=local_dest
         )
@@ -142,11 +142,11 @@ def xe(self, action: str, args: Dict[str, Union[str, bool]] = {}, *, check: bool
 
     @overload
     def xe(self, action: str, args: Dict[str, Union[str, bool]] = {}, *, check: bool = ...,
-           simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> commands.SSHResult:
+           simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> netutil.SSHResult:
         ...
 
     def xe(self, action, args={}, *, check=True, simple_output=True, minimal=False, force=False) \
-            -> Union[str, commands.SSHResult]:
+            -> Union[str, netutil.SSHResult]:
         maybe_param_minimal = ['--minimal'] if minimal else []
         maybe_param_force = ['--force'] if force else []
 
@@ -167,7 +167,7 @@ def stringify(key, value):
             check=check,
             simple_output=simple_output
         )
-        assert isinstance(result, (str, commands.SSHResult))
+        assert isinstance(result, (str, netutil.SSHResult))
 
         return result
 
@@ -426,7 +426,7 @@ def restart_toolstack(self, verify=False):
     def is_enabled(self) -> bool:
         try:
             return strtobool(self.param_get('enabled'))
-        except commands.SSHCommandFailed:
+        except netutil.SSHCommandFailed:
             # If XAPI is not ready yet, or the host is down, this will throw. We return False in that case.
             return False
 
@@ -436,7 +436,7 @@ def has_updates(self):
             self.ssh(['yum', 'check-update'])
             # returned 0, else there would have been a SSHCommandFailed
             return False
-        except commands.SSHCommandFailed as e:
+        except netutil.SSHCommandFailed as e:
             if e.returncode == 100:
                 return True
             else:
@@ -458,7 +458,7 @@ def get_last_yum_history_tid(self):
         """
         try:
             history_str = self.ssh(['yum', 'history', 'list', '--noplugins'])
-        except commands.SSHCommandFailed:
+        except netutil.SSHCommandFailed:
             # yum history list fails if the list is empty, and it's also not possible to rollback
             # to before the first transaction, so "0" would not be appropriate as last transaction.
             # To workaround this, create transactions: install and remove a small package.
@@ -545,7 +545,7 @@ def reboot(self, verify=False):
         logging.info("Reboot host %s" % self)
         try:
             self.ssh(['reboot'])
-        except commands.SSHCommandFailed as e:
+        except netutil.SSHCommandFailed as e:
             # ssh connection may get killed by the reboot and terminate with an error code
             if "closed by remote host" not in e.stdout:
                 raise

lib/installer.py

@@ -2,8 +2,8 @@
 import time
 import xml.etree.ElementTree as ET
 
-from lib.commands import SSHCommandFailed, ssh
 from lib.common import wait_for
+from lib.netutil import SSHCommandFailed, ssh
 
 class AnswerFile:
     def __init__(self, kind, /):