lib/device.py

# Copyright 2018 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Device-related helper functions/classes."""

import argparse
import logging
import os
import subprocess

from chromite.cli.cros import cros_chrome_sdk
from chromite.lib import commandline
from chromite.lib import remote_access
from chromite.lib import retry_util


class DeviceError(Exception):
    """Exception for Device failures."""


class Device:
    """Class for managing a test device."""

    SSH_CONNECT_TIMEOUT = 30

    def __init__(self, opts) -> None:
        """Initialize Device.

        Args:
            opts: command line options.
        """
        self.device = opts.device.hostname if opts.device else None
        self.ssh_port = opts.device.port if opts.device else None
        self.should_start_vm = not self.device
        self.board = opts.board

        self.use_sudo = False
        self.cmd = opts.args[1:] if opts.cmd else None
        self.private_key = opts.private_key
        self.dryrun = opts.dryrun
        # log_level is only set if --log-level or --debug is specified.
        self.log_level = getattr(opts, "log_level", None)
        self.InitRemote()

    def InitRemote(self, connect_timeout=SSH_CONNECT_TIMEOUT) -> None:
        """Initialize remote access."""
        self.remote = remote_access.ChromiumOSDevice(
            self.device,
            port=self.ssh_port,
            connect_settings=self._ConnectSettings(
                connect_timeout=connect_timeout
            ),
            private_key=self.private_key,
            include_dev_paths=False,
        )

        self.device_addr = "ssh://%s" % self.device
        if self.ssh_port:
            self.device_addr += ":%d" % self.ssh_port

    def WaitForBoot(self, max_retry=10, sleep=5) -> None:
        """Wait for the device to boot up.

        Wait for the ssh connection to become active.
        """
        try:
            result = retry_util.RetryException(
                exception=remote_access.SSHConnectionError,
                max_retry=max_retry,
                functor=lambda: self.run(cmd=["true"]),
                sleep=sleep,
            )
        except remote_access.SSHConnectionError:
            raise DeviceError(
                "WaitForBoot timed out trying to connect to the device."
            )

        if result.returncode != 0:
            raise DeviceError("WaitForBoot failed: %s." % result.stderr)

    def chmod(self, *args, **kwargs) -> None:
        return self.remote.chmod(*args, **kwargs)

    def mkdir(self, *args, **kwargs) -> None:
        return self.remote.mkdir(*args, **kwargs)

    def run(self, cmd, stream_output=False, **kwargs):
        """Run a remote command.

        Args:
            cmd: command to run.
            stream_output: Stream output of long-running commands.
            **kwargs: additional args (see documentation for RemoteDevice.run).

        Returns:
            cros_build_lib.CompletedProcess object.
        """
        kwargs.setdefault("check", False)
        if stream_output:
            kwargs.setdefault("capture_output", False)
        else:
            kwargs.setdefault("stderr", subprocess.STDOUT)
            kwargs.setdefault("log_output", True)
        return self.remote.run(
            cmd, dryrun=self.dryrun, debug_level=logging.INFO, **kwargs
        )

    def _ConnectSettings(self, connect_timeout):
        """Increase ServerAliveCountMax and ServerAliveInterval.

        Wait 2 min before dropping the SSH connection.

        Args:
            connect_timeout: SSH ConnectTimeout setting.

        Returns:
            List of arguments to pass to SSH.
        """
        return remote_access.CompileSSHConnectSettings(
            ConnectTimeout=connect_timeout,
            ServerAliveInterval=15,
            ServerAliveCountMax=8,
        )

    @staticmethod
    def Create(opts):
        """Create either a Device or VM based on |opts.device|."""
        if not opts.device:
            from chromite.lib import vm

            return vm.VM(opts)
        return Device(opts)

    @staticmethod
    def GetParser(parser=None):
        """Parse a list of args.

        Returns:
            List of parsed opts.
        """
        parser = parser or commandline.ArgumentParser(
            description=__doc__, dryrun=True
        )
        parser.add_argument(
            "-d",
            "--device",
            type=commandline.DeviceParser(commandline.DeviceScheme.SSH),
            help="Hostname or device IP in format hostname[:port]. If not "
            "specified, a VM will be launched for the duration of the test.",
        )
        sdk_board_env = os.environ.get(cros_chrome_sdk.SDKFetcher.SDK_BOARD_ENV)
        parser.add_argument(
            "--board", default=sdk_board_env, help="Board to use."
        )
        parser.add_argument("--private-key", help="Path to ssh private key.")
        parser.add_argument(
            "--cmd", action="store_true", default=False, help="Run a command."
        )
        parser.add_argument(
            "args", nargs=argparse.REMAINDER, help="Command to run."
        )
        return parser