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

atex/connection/ssh.py

@@ -0,0 +1,390 @@
+"""
+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(str(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
+
+    # 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 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 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/executor/__init__.py

@@ -0,0 +1,2 @@
+from . import testcontrol  # noqa: F401
+from .executor import Executor  # noqa: F401

atex/executor/duration.py

@@ -0,0 +1,60 @@
+import re
+import time
+
+
+class Duration:
+    """
+    A helper for parsing, keeping and manipulating test run time based on
+    FMF-defined 'duration' attribute.
+    """
+
+    def __init__(self, fmf_duration):
+        """
+        'fmf_duration' is the string specified as 'duration' in FMF metadata.
+        """
+        duration = self._fmf_to_seconds(fmf_duration)
+        self.end = time.monotonic() + duration
+        # keep track of only the first 'save' and the last 'restore',
+        # ignore any nested ones (as tracked by '_count')
+        self.saved = None
+        self.saved_count = 0
+
+    @staticmethod
+    def _fmf_to_seconds(string):
+        match = re.fullmatch(r"([0-9]+)([a-z]*)", string)
+        if not match:
+            raise RuntimeError(f"'duration' has invalid format: {string}")
+        length, unit = match.groups()
+        if unit == "m":
+            return int(length)*60
+        elif unit == "h":
+            return int(length)*60*60
+        elif unit == "d":
+            return int(length)*60*60*24
+        else:
+            return int(length)
+
+    def set(self, to):
+        self.end = time.monotonic() + self._fmf_to_seconds(to)
+
+    def increment(self, by):
+        self.end += self._fmf_to_seconds(by)
+
+    def decrement(self, by):
+        self.end -= self._fmf_to_seconds(by)
+
+    def save(self):
+        if self.saved_count == 0:
+            self.saved = self.end - time.monotonic()
+        self.saved_count += 1
+
+    def restore(self):
+        if self.saved_count > 1:
+            self.saved_count -= 1
+        elif self.saved_count == 1:
+            self.end = time.monotonic() + self.saved
+            self.saved_count = 0
+            self.saved = None
+
+    def out_of_time(self):
+        return time.monotonic() > self.end