atex 0.5__py3-none-any.whl → 0.7__py3-none-any.whl

atex/connection/ssh.py

@@ -0,0 +1,406 @@
+"""
+Connection API implementation using the OpenSSH ssh(1) client.
+
+Any SSH options are passed via dictionaries of options, and later translated
+to '-o' client CLI options, incl. Hostname, User, Port, IdentityFile, etc.
+No "typical" ssh CLI switches are used.
+
+This allows for a nice flexibility from Python code - this module provides
+some sensible option defaults (for scripted use), but you are free to
+overwrite any options via class or function arguments (where appropriate).
+
+Note that .cmd() quotes arguments to really execute individual arguments
+as individual arguments in the remote shell, so you need to give it a proper
+iterable (like for other Connections), not a single string with spaces.
+"""
+
+import os
+import time
+import shlex
+import tempfile
+import subprocess
+from pathlib import Path
+
+from .. import util
+from . import Connection
+
+
+DEFAULT_OPTIONS = {
+    "LogLevel": "ERROR",
+    "StrictHostKeyChecking": "no",
+    "UserKnownHostsFile": "/dev/null",
+    "ConnectionAttempts": "3",
+    "ServerAliveCountMax": "4",
+    "ServerAliveInterval": "5",
+    "TCPKeepAlive": "no",
+    "EscapeChar": "none",
+    "ExitOnForwardFailure": "yes",
+}
+
+
+class SSHError(Exception):
+    pass
+
+
+class DisconnectedError(SSHError):
+    """
+    Raised when an already-connected ssh session goes away (breaks connection).
+    """
+    pass
+
+
+class NotConnectedError(SSHError):
+    """
+    Raised when an operation on ssh connection is requested, but the connection
+    is not yet open (or has been closed/disconnected).
+    """
+    pass
+
+
+class ConnectError(SSHError):
+    """
+    Raised when a to-be-opened ssh connection fails to open.
+    """
+    pass
+
+
+def _shell_cmd(command, sudo=None):
+    """
+    Make a command line for running 'command' on the target system.
+    """
+    quoted_args = (shlex.quote(arg) for arg in command)
+    if sudo:
+        return " ".join((
+            "exec", "sudo", "--no-update", "--non-interactive", "--user", sudo, "--", *quoted_args,
+        ))
+    else:
+        return " ".join(("exec", *quoted_args))
+
+
+def _options_to_cli(options):
+    """
+    Assemble an ssh(1) or sshpass(1) command line with -o options.
+    """
+    list_opts = []
+    for key, value in options.items():
+        if isinstance(value, (list, tuple, set)):
+            list_opts += (f"-o{key}={v}" for v in value)
+        else:
+            list_opts.append(f"-o{key}={value}")
+    return list_opts
+
+
+def _options_to_ssh(options, password=None, extra_cli_flags=()):
+    """
+    Assemble an ssh(1) or sshpass(1) command line with -o options.
+    """
+    cli_opts = _options_to_cli(options)
+    if password:
+        return (
+            "sshpass", "-p", password,
+            "ssh", *extra_cli_flags, "-oBatchMode=no", *cli_opts,
+            "ignored_arg",
+        )
+    else:
+        # let cli_opts override BatchMode if specified
+        return ("ssh", *extra_cli_flags, *cli_opts, "-oBatchMode=yes", "ignored_arg")
+
+
+# return a string usable for rsync -e
+def _options_to_rsync_e(options, password=None):
+    """
+    Return a string usable for the rsync -e argument.
+    """
+    cli_opts = _options_to_cli(options)
+    batch_mode = "-oBatchMode=no" if password else "-oBatchMode=yes"
+    return " ".join(("ssh", *cli_opts, batch_mode))  # no ignored_arg inside -e
+
+
+def _rsync_host_cmd(*args, options, password=None, sudo=None):
+    """
+    Assemble a rsync command line, noting that
+      - 'sshpass' must be before 'rsync', not inside the '-e' argument
+      - 'ignored_arg' must be passed by user as destination, not inside '-e'
+      - 'sudo' is part of '--rsync-path', yet another argument
+    """
+    return (
+        *(("sshpass", "-p", password) if password else ()),
+        "rsync",
+        "-e", _options_to_rsync_e(options, password=password),
+        "--rsync-path", _shell_cmd(("rsync",), sudo=sudo),
+        *args,
+    )
+
+
+class StatelessSSHConn(Connection):
+    """
+    Implements the Connection API using a ssh(1) client using "standalone"
+    (stateless) logic - connect() and disconnect() are no-op, .cmd() simply
+    executes the ssh client and .rsync() executes 'rsync -e ssh'.
+
+    Compared to ManagedSSHConn, this may be slow for many .cmd() calls,
+    but every call is stateless, there is no persistent connection.
+
+    If you need only one .cmd(), this will be faster than ManagedSSHConn.
+    """
+
+    def __init__(self, options, *, password=None, sudo=None):
+        """
+        Prepare to connect to an SSH server specified in 'options'.
+
+        If 'password' is given, spawn the ssh(1) command via 'sshpass' and
+        pass the password to it.
+
+        If 'sudo' specifies a username, call sudo(8) on the remote shell
+        to run under a different user on the remote host.
+        """
+        super().__init__()
+        self.options = DEFAULT_OPTIONS.copy()
+        self.options.update(options)
+        self.password = password
+        self.sudo = sudo
+        self._tmpdir = None
+        self._master_proc = None
+
+    def connect(self):
+        """
+        Optional, .cmd() and .rsync() work without it, but it is provided here
+        for compatibility with the Connection API.
+        """
+        # TODO: just wait until .cmd(['true']) starts responding
+        pass
+
+    def disconnect(self):
+        pass
+
+#    def alive(self):
+#        return True
+
+    # have options as kwarg to be compatible with other functions here
+    def cmd(self, command, options=None, func=util.subprocess_run, **func_args):
+        unified_options = self.options.copy()
+        if options:
+            unified_options.update(options)
+        unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
+        return func(
+            _options_to_ssh(unified_options, password=self.password),
+            skip_frames=1,
+            **func_args,
+        )
+
+    def rsync(self, *args, options=None, func=util.subprocess_run, **func_args):
+        unified_options = self.options.copy()
+        if options:
+            unified_options.update(options)
+        return func(
+            _rsync_host_cmd(
+                *args,
+                options=unified_options,
+                password=self.password,
+                sudo=self.sudo,
+            ),
+            skip_frames=1,
+            check=True,
+            stdin=subprocess.DEVNULL,
+            **func_args,
+        )
+
+
+# Note that when ControlMaster goes away (connection breaks), any ssh clients
+# connected through it will time out after a combination of
+#   ServerAliveCountMax + ServerAliveInterval + ConnectionAttempts
+# identical to the ControlMaster process.
+# Specifying different values for the clients, to make them exit faster when
+# the ControlMaster dies, has no effect. They seem to ignore the options.
+#
+# If you need to kill the clients quickly after ControlMaster disconnects,
+# you need to set up an independent polling logic (ie. every 0.1sec) that
+# checks .assert_master() and manually signals the running clients
+# when it gets DisconnectedError from it.
+
+class ManagedSSHConn(Connection):
+    """
+    Implements the Connection API using one persistently-running ssh(1) client
+    started in a 'ControlMaster' mode, with additional ssh clients using that
+    session to execute remote commands. Similarly, .rsync() uses it too.
+
+    This is much faster than StatelessSSHConn when executing multiple commands,
+    but contains a complex internal state (what if ControlMaster disconnects?).
+
+    Hence why this implementation provides extra non-standard-Connection methods
+    to manage this complexity.
+    """
+
+    # TODO: thread safety and locking via self.lock
+
+    def __init__(self, options, *, password=None, sudo=None):
+        """
+        Prepare to connect to an SSH server specified in 'options'.
+
+        If 'password' is given, spawn the ssh(1) command via 'sshpass' and
+        pass the password to it.
+
+        If 'sudo' specifies a username, call sudo(8) on the remote shell
+        to run under a different user on the remote host.
+        """
+        super().__init__()
+        self.options = DEFAULT_OPTIONS.copy()
+        self.options.update(options)
+        self.password = password
+        self.sudo = sudo
+        self._tmpdir = None
+        self._master_proc = None
+
+#    def __copy__(self):
+#        return type(self)(self.options, password=self.password)
+#
+#    def copy(self):
+#        return self.__copy__()
+
+    def assert_master(self):
+        proc = self._master_proc
+        if not proc:
+            raise NotConnectedError("SSH ControlMaster is not running")
+        # we need to consume any potential proc output for the process to
+        # actually terminate (stop being a zombie) if it crashes
+        out = proc.stdout.read()
+        code = proc.poll()
+        if code is not None:
+            self._master_proc = None
+            out = f":\n{out.decode()}" if out else ""
+            raise DisconnectedError(
+                f"SSH ControlMaster on {self._tmpdir} exited with {code}{out}",
+            )
+
+#    def alive(self):
+#        try:
+#            self.assert_master()
+#            return True
+#        except (NotConnectedError, DisconnectedError):
+#            return False
+
+    def disconnect(self):
+        proc = self._master_proc
+        if not proc:
+            return
+        proc.kill()
+        # don"t zombie forever, return EPIPE on any attempts to write to us
+        proc.stdout.close()
+        proc.wait()
+        (self._tmpdir / "control.sock").unlink(missing_ok=True)
+        self._master_proc = None
+
+    def connect(self, block=True):
+        if not self._tmpdir:
+            # _tmpdir_handle just prevents the TemporaryDirectory instance
+            # from being garbage collected (and removed on disk)
+            # TODO: create/remove it explicitly in connect/disconnect
+            #       so the removal happens immediately, even if GC delays cleaning
+            self._tmpdir_handle = tempfile.TemporaryDirectory(prefix="atex-ssh-")
+            self._tmpdir = Path(self._tmpdir_handle.name)
+
+        sock = self._tmpdir / "control.sock"
+
+        if not self._master_proc:
+            options = self.options.copy()
+            options["SessionType"] = "none"
+            options["ControlMaster"] = "yes"
+            options["ControlPath"] = sock
+            self._master_proc = util.subprocess_Popen(
+                _options_to_ssh(options),
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                cwd=str(self._tmpdir),
+            )
+            os.set_blocking(self._master_proc.stdout.fileno(), False)
+
+        proc = self._master_proc
+        if block:
+            while proc.poll() is None:
+                if sock.exists():
+                    break
+                time.sleep(0.1)
+            else:
+                code = proc.poll()
+                out = proc.stdout.read()
+                self._master_proc = None
+                # TODO: ConnectError should probably be generalized for Connection
+                raise ConnectError(
+                    f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
+                )
+        else:
+            code = proc.poll()
+            if code is not None:
+                out = proc.stdout.read()
+                self._master_proc = None
+                # TODO: ConnectError should probably be generalized for Connection
+                raise ConnectError(
+                    f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
+                )
+            elif not sock.exists():
+                raise BlockingIOError("SSH ControlMaster not yet ready")
+
+    def add_local_forward(self, *spec):
+        """
+        Add (one or more) ssh forwarding specifications as 'spec' to an
+        already-connected instance. Each specification has to follow the
+        format of ssh client's LocalForward option (see ssh_config(5)).
+        """
+        self.assert_master()
+        options = self.options.copy()
+        options["LocalForward"] = spec
+        options["ControlPath"] = self._tmpdir / "control.sock"
+        util.subprocess_run(
+            _options_to_ssh(options, extra_cli_flags=("-O", "forward")),
+            skip_frames=1,
+            check=True,
+        )
+
+    def add_remote_forward(self, *spec):
+        """
+        Add (one or more) ssh forwarding specifications as 'spec' to an
+        already-connected instance. Each specification has to follow the
+        format of ssh client's RemoteForward option (see ssh_config(5)).
+        """
+        self.assert_master()
+        options = self.options.copy()
+        options["RemoteForward"] = spec
+        options["ControlPath"] = self._tmpdir / "control.sock"
+        util.subprocess_run(
+            _options_to_ssh(options, extra_cli_flags=("-O", "forward")),
+            skip_frames=1,
+            check=True,
+        )
+
+    def cmd(self, command, options=None, func=util.subprocess_run, **func_args):
+        self.assert_master()
+        unified_options = self.options.copy()
+        if options:
+            unified_options.update(options)
+        unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
+        unified_options["ControlPath"] = self._tmpdir / "control.sock"
+        return func(
+            _options_to_ssh(unified_options, password=self.password),
+            skip_frames=1,
+            **func_args,
+        )
+
+    def rsync(self, *args, options=None, func=util.subprocess_run, **func_args):
+        self.assert_master()
+        unified_options = self.options.copy()
+        if options:
+            unified_options.update(options)
+        unified_options["ControlPath"] = self._tmpdir / "control.sock"
+        return func(
+            _rsync_host_cmd(
+                *args,
+                options=unified_options,
+                password=self.password,
+                sudo=self.sudo,
+            ),
+            skip_frames=1,
+            check=True,
+            stdin=subprocess.DEVNULL,
+            **func_args,
+        )

atex/minitmt/__init__.py

@@ -1,115 +1,23 @@
-import os
-import random
+"""
+TODO Minitmt documentation - reference README, etc.
+"""
 
-from pathlib import Path
+import importlib as _importlib
+import pkgutil as _pkgutil
 
-# TODO: TMT_PLAN_ENVIRONMENT_FILE
+__all__ = [
+    info.name for info in _pkgutil.iter_modules(__spec__.submodule_search_locations)
+]
 
-# TODO: install rsync on the guest as part of setup
 
-# TODO: in Orchestrator, when a Provisioner becomes free, have it pick a test
-#       from the appropriate tests[platform] per the Provisioner's platform
+def __dir__():
+    return __all__
 
 
-def _random_string(length):
-    return ''.join(
-        random.choices('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', k=length),
-    )
-
-
-class Preparator:
-    """
-    Set of utilities for preparing a newly acquired/reserved machine for
-    running tests, by installing global package requirements, copying all
-    tests over, executing tmt plan 'prepare' step, etc.
-    """
-    def __init__(self, ssh_conn):
-        self.conn = ssh_conn
-
-    def copy_tests(self):
-        pass
-
-    def run_prepare_scripts(self):
-        pass
-
-    def __enter__(self):
-        self.conn.connect()
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.conn.disconnect()
-
-
-# TODO: have Executor take a finished Preparator instance as input?
-#        - for extracting copied tests location
-#        - for extracting TMT_PLAN_ENVIRONMENT_FILE location
-#        - etc.
-
-
-class Executor:
-    """
-    Helper for running one test on a remote system and processing results
-    and uploaded files by that test.
-    """
-    def __init__(self, fmf_test, ssh_conn):
-        self.fmf_test = fmf_test
-        self.conn = ssh_conn
-        self.remote_socket = self.local_socket = None
-
-    def __enter__(self):
-        # generate a (hopefully) unique test control socket name
-        # and modify the SSHConn instance to use it
-        rand_name = f'atex-control-{_random_string(50)}.sock'
-        self.local_socket = Path(os.environ.get('TMPDIR', '/tmp')) / rand_name
-        self.remote_socket = f'/tmp/{rand_name}'
-        self.conn.options['RemoteForward'] = f'{self.remote_socket} {self.local_socket}'
-        self.conn.connect()
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.conn.ssh(f'rm -f {self.remote_socket}')
-        self.local_socket.unlink()
-        self.remote_socket = self.local_socket = None
-        self.conn.disconnect()
-
-    # execute all prepares (how:install and how:shell) via ssh
-    def prepare(self):
-        # TODO: check via __some_attr (named / prefixed after our class)
-        #       whether this reserved system has been prepared already ... ?
-        #        ^^^^ in Orchestrator
-        #
-        # TODO: copy root of fmf metadata to some /var/tmp/somedir to run tests from
-        #
-        # TODO: move prepare out, possibly to class-less function,
-        #       we don't want it running over an SSHConn that would set up socket forwarding
-        #       only to tear it back down, when executed from Orchestrator for setup only
-        #
-        # TODO: install rsync
-        pass
-
-    def run_script(self, script, duration=None, shell='/bin/bash', **kwargs):
-        self.conn.ssh(shell, input=script.encode())
-
-    # run one test via ssh and parse its results on-the-fly,
-    # write out logs
-    def run_test(self, fmf_test, reporter):
-        # TODO: pass environment from test fmf metadata
-        # TODO: watch for test duration, etc. metadata
-        # TODO: logging of stdout+stderr to hidden file, doing 'ln' from it to
-        #       test-named 'testout' files
-        #       - generate hidden name suffix via:
-        #         ''.join(random.choices('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', k=20))
-        output_logfile = \
-            reporter.files_dir(fmf_test.name) / f'.test_output_{self._random_string(50)}.log'
-        output_logfile = os.open(reporter.files_dir(fmf_test.name), os.O_WRONLY | os.O_CREAT)
-        try:
-            #self.conn.ssh(
-            pass
-        finally:
-            os.close(output_logfile)
-        # TODO: create temp dir on remote via 'mktemp -d', then call
-        #       self.conn.add_remote_forward(...) with socket path inside that tmpdir
-
-# TODO: run tests by passing stdout/stderr via pre-opened fd so we don't handle it in code
-
-# TODO: read unix socket as nonblocking, check test subprocess.Popen proc status every 0.1sec
+# lazily import submodules
+def __getattr__(attr):
+    # importing a module known to exist
+    if attr in __all__:
+        return _importlib.import_module(f".{attr}", __name__)
+    else:
+        raise AttributeError(f"module '{__name__}' has no attribute '{attr}'")