| #!/usr/bin/python |
| # Copyright (c) 2012 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Library containing functions to access a remote test device.""" |
| |
| import glob |
| import logging |
| import os |
| import shutil |
| import stat |
| import time |
| |
| from chromite.lib import cros_build_lib |
| from chromite.lib import osutils |
| from chromite.lib import timeout_util |
| |
| |
| _path = os.path.dirname(os.path.realpath(__file__)) |
| TEST_PRIVATE_KEY = os.path.normpath( |
| os.path.join(_path, '../ssh_keys/testing_rsa')) |
| del _path |
| |
| REBOOT_MARKER = '/tmp/awaiting_reboot' |
| REBOOT_MAX_WAIT = 120 |
| REBOOT_SSH_CONNECT_TIMEOUT = 2 |
| REBOOT_SSH_CONNECT_ATTEMPTS = 2 |
| CHECK_INTERVAL = 5 |
| DEFAULT_SSH_PORT = 22 |
| SSH_ERROR_CODE = 255 |
| |
| |
| def CompileSSHConnectSettings(ConnectTimeout=30, ConnectionAttempts=4): |
| return ['-o', 'ConnectTimeout=%s' % ConnectTimeout, |
| '-o', 'ConnectionAttempts=%s' % ConnectionAttempts, |
| '-o', 'NumberOfPasswordPrompts=0', |
| '-o', 'Protocol=2', |
| '-o', 'ServerAliveInterval=10', |
| '-o', 'ServerAliveCountMax=3', |
| '-o', 'StrictHostKeyChecking=no', |
| '-o', 'UserKnownHostsFile=/dev/null', ] |
| |
| |
| class RemoteAccess(object): |
| """Provides access to a remote test machine.""" |
| |
| def __init__(self, remote_host, tempdir, port=DEFAULT_SSH_PORT, |
| debug_level=logging.DEBUG, interactive=True): |
| """Construct the object. |
| |
| Args: |
| remote_host: The ip or hostname of the remote test machine. The test |
| machine should be running a ChromeOS test image. |
| tempdir: A directory that RemoteAccess can use to store temporary files. |
| It's the responsibility of the caller to remove it. |
| port: The ssh port of the test machine to connect to. |
| debug_level: Logging level to use for all RunCommand invocations. |
| interactive: If set to False, pass /dev/null into stdin for the sh cmd. |
| """ |
| self.tempdir = tempdir |
| self.remote_host = remote_host |
| self.port = port |
| self.debug_level = debug_level |
| self.private_key = os.path.join(tempdir, os.path.basename(TEST_PRIVATE_KEY)) |
| self.interactive = interactive |
| shutil.copyfile(TEST_PRIVATE_KEY, self.private_key) |
| os.chmod(self.private_key, stat.S_IRUSR) |
| |
| @property |
| def target_ssh_url(self): |
| return 'root@%s' % self.remote_host |
| |
| def _GetSSHCmd(self, connect_settings=None): |
| if connect_settings is None: |
| connect_settings = CompileSSHConnectSettings() |
| |
| cmd = (['ssh', '-p', str(self.port)] + |
| connect_settings + |
| ['-i', self.private_key]) |
| if not self.interactive: |
| cmd.append('-n') |
| |
| return cmd |
| |
| def RemoteSh(self, cmd, connect_settings=None, error_code_ok=False, |
| ssh_error_ok=False, **kwargs): |
| """Run a sh command on the remote device through ssh. |
| |
| Args: |
| cmd: The command string or list to run. |
| connect_settings: The SSH connect settings to use. |
| error_code_ok: Does not throw an exception when the command exits with a |
| non-zero returncode. This does not cover the case where |
| the ssh command itself fails (return code 255). |
| See ssh_error_ok. |
| ssh_error_ok: Does not throw an exception when the ssh command itself |
| fails (return code 255). |
| **kwargs: See cros_build_lib.RunCommand documentation. |
| |
| Returns: |
| A CommandResult object. The returncode is the returncode of the command, |
| or 255 if ssh encountered an error (could not connect, connection |
| interrupted, etc.) |
| |
| Raises: |
| RunCommandError when error is not ignored through error_code_ok and |
| ssh_error_ok flags. |
| """ |
| kwargs.setdefault('debug_level', self.debug_level) |
| |
| ssh_cmd = self._GetSSHCmd(connect_settings) |
| if isinstance(cmd, basestring): |
| ssh_cmd += [self.target_ssh_url, '--', cmd] |
| else: |
| ssh_cmd += [self.target_ssh_url, '--'] + cmd |
| |
| try: |
| result = cros_build_lib.RunCommandCaptureOutput( |
| ssh_cmd, **kwargs) |
| except cros_build_lib.RunCommandError as e: |
| if ((e.result.returncode == SSH_ERROR_CODE and ssh_error_ok) or |
| (e.result.returncode and e.result.returncode != SSH_ERROR_CODE |
| and error_code_ok)): |
| result = e.result |
| else: |
| raise |
| |
| return result |
| |
| def LearnBoard(self): |
| """Grab the board reported by the remote device. |
| |
| in the case of multiple matches, uses the first one. In the case of no |
| entry, returns an empty string. |
| """ |
| result = self.RemoteSh('grep CHROMEOS_RELEASE_BOARD /etc/lsb-release') |
| # In the case of multiple matches, use the first one. |
| output = result.output.splitlines() |
| if len(output) > 1: |
| logging.debug('More than one board entry found! Using the first one.') |
| |
| return output[0].strip().partition('=')[-1] |
| |
| def _CheckIfRebooted(self): |
| """"Checks whether a remote device has rebooted successfully. |
| |
| This uses a rapidly-retried SSH connection, which will wait for at most |
| about ten seconds. If the network returns an error (e.g. host unreachable) |
| the actual delay may be shorter. |
| |
| Returns: |
| Whether the device has successfully rebooted. |
| """ |
| # In tests SSH seems to be waiting rather longer than would be expected |
| # from these parameters. These values produce a ~5 second wait. |
| connect_settings = CompileSSHConnectSettings( |
| ConnectTimeout=REBOOT_SSH_CONNECT_TIMEOUT, |
| ConnectionAttempts=REBOOT_SSH_CONNECT_ATTEMPTS) |
| cmd = "[ ! -e '%s' ]" % REBOOT_MARKER |
| result = self.RemoteSh(cmd, connect_settings=connect_settings, |
| error_code_ok=True, ssh_error_ok=True) |
| |
| errors = {0: 'Reboot complete.', |
| 1: 'Device has not yet shutdown.', |
| 255: 'Cannot connect to device; reboot in progress.'} |
| if result.returncode not in errors: |
| raise Exception('Unknown error code %s returned by %s.' |
| % (result.returncode, cmd)) |
| |
| logging.info(errors[result.returncode]) |
| return result.returncode == 0 |
| |
| def RemoteReboot(self): |
| """Reboot the remote device.""" |
| logging.info('Rebooting %s...', self.remote_host) |
| self.RemoteSh('touch %s && reboot' % REBOOT_MARKER) |
| time.sleep(CHECK_INTERVAL) |
| try: |
| timeout_util.WaitForReturnTrue(self._CheckIfRebooted, REBOOT_MAX_WAIT, |
| period=CHECK_INTERVAL) |
| except timeout_util.TimeoutError: |
| cros_build_lib.Die('Reboot has not completed after %s seconds; giving up.' |
| % (REBOOT_MAX_WAIT,)) |
| |
| def Rsync(self, src, dest, to_local=False, follow_symlinks=False, |
| inplace=False, verbose=False, sudo=False, **kwargs): |
| """Rsync a path to the remote device. |
| |
| Rsync a path to the remote device. If |to_local| is set True, it |
| rsyncs the path from the remote device to the local machine. |
| |
| Args: |
| src: The local src directory. |
| dest: The remote dest directory. |
| to_local: If set, rsync remote path to local path. |
| follow_symlinks: If set, transform symlinks into referent |
| path. Otherwise, copy symlinks as symlinks. |
| inplace: If set, cause rsync to overwrite the dest files in place. This |
| conserves space, but has some side effects - see rsync man page. |
| verbose: If set, print more verbose output during rsync file transfer. |
| sudo: If set, invoke the command via sudo. |
| **kwargs: See cros_build_lib.RunCommand documentation. |
| """ |
| kwargs.setdefault('debug_level', self.debug_level) |
| |
| ssh_cmd = ' '.join(self._GetSSHCmd()) |
| rsync_cmd = ['rsync', '--recursive', '--perms', '--verbose', |
| '--times', '--compress', '--omit-dir-times', |
| '--exclude', '.svn'] |
| rsync_cmd.append('--copy-links' if follow_symlinks else '--links') |
| # In cases where the developer sets up a ssh daemon manually on a device |
| # with a dev image, the ssh login $PATH can be incorrect, and the target |
| # rsync will not be found. So we try to provide the right $PATH here. |
| rsync_cmd += ['--rsync-path', 'PATH=/usr/local/bin:$PATH rsync'] |
| if verbose: |
| rsync_cmd.append('--progress') |
| if inplace: |
| rsync_cmd.append('--inplace') |
| |
| if to_local: |
| rsync_cmd += ['--rsh', ssh_cmd, |
| '[%s]:%s' % (self.target_ssh_url, src), dest] |
| else: |
| rsync_cmd += ['--rsh', ssh_cmd, src, |
| '[%s]:%s' % (self.target_ssh_url, dest)] |
| |
| rc_func = cros_build_lib.RunCommand |
| if sudo: |
| rc_func = cros_build_lib.SudoRunCommand |
| return rc_func(rsync_cmd, print_cmd=verbose, **kwargs) |
| |
| def RsyncToLocal(self, *args, **kwargs): |
| """Rsync a path from the remote device to the local machine.""" |
| return self.Rsync(*args, to_local=kwargs.pop('to_local', True), **kwargs) |
| |
| def Scp(self, src, dest, recursive=False, verbose=False, debug_level=None, |
| sudo=False): |
| """Scp a file or directory to the remote device. |
| |
| Args: |
| src: The local src file or directory. |
| dest: The remote dest location. |
| recursive: Whether to recursively copy entire directories. |
| verbose: If set, print more verbose output during scp file transfer. |
| debug_level: See cros_build_lib.RunCommand documentation. |
| sudo: If set, invoke the command via sudo. |
| |
| Returns: |
| A CommandResult object containing the information and return code of |
| the scp command. |
| """ |
| if not debug_level: |
| debug_level = self.debug_level |
| |
| scp_cmd = ['scp', '-p'] + CompileSSHConnectSettings(ConnectTimeout=60) |
| |
| if recursive: |
| scp_cmd.append('-r') |
| if verbose: |
| scp_cmd.append('-v') |
| |
| scp_cmd += glob.glob(src) + ['%s:%s' % (self.target_ssh_url, dest)] |
| |
| rc_func = cros_build_lib.RunCommand |
| if sudo: |
| rc_func = cros_build_lib.SudoRunCommand |
| |
| return rc_func(scp_cmd, debug_level=debug_level, print_cmd=verbose) |
| |
| def PipeToRemoteSh(self, producer_cmd, cmd, **kwargs): |
| """Run a local command and pipe it to a remote sh command over ssh. |
| |
| Args: |
| producer_cmd: Command to run locally with its results piped to |cmd|. |
| cmd: Command to run on the remote device. |
| **kwargs: See RemoteSh for documentation. |
| """ |
| result = cros_build_lib.RunCommandCaptureOutput(producer_cmd, |
| stdout_to_pipe=True, |
| print_cmd=False) |
| return self.RemoteSh(cmd, input=kwargs.pop('input', result.output), |
| **kwargs) |
| |
| |
| class RemoteDeviceHandler(object): |
| """A wrapper of RemoteDevice.""" |
| |
| def __init__(self, *args, **kwargs): |
| """Creates a RemoteDevice object.""" |
| self.device = RemoteDevice(*args, **kwargs) |
| |
| def __enter__(self): |
| """Return the temporary directory.""" |
| return self.device |
| |
| def __exit__(self, _type, _value, _traceback): |
| """Cleans up the device.""" |
| self.device.Cleanup() |
| |
| |
| class RemoteDevice(object): |
| """Handling basic SSH communication with a remote device.""" |
| |
| DEFAULT_WORK_DIR = '/tmp/remote-access' |
| |
| def __init__(self, hostname, debug_level=logging.DEBUG, work_dir=None): |
| """Initializes a RemoteDevice object. |
| |
| Args: |
| hostname: The hostname of the device. |
| debug_level: Setting debug level for logging. |
| work_dir: The default working directory on the device. |
| """ |
| self.hostname = hostname |
| # The tempdir is only for storing the rsa key on the host. |
| self.tempdir = osutils.TempDir(prefix='ssh-tmp') |
| self.agent = self._SetupSSH() |
| self.board = self.agent.LearnBoard() |
| self.debug_level = debug_level |
| # Setup a working directory on the device. |
| self.work_dir = self.DEFAULT_WORK_DIR if work_dir is None else work_dir |
| self.RunCommand(['rm', '-r', self.work_dir], error_code_ok=True) |
| self.RunCommand(['mkdir', '-p', self.work_dir]) |
| |
| def _SetupSSH(self): |
| """Setup the ssh connection with device.""" |
| return RemoteAccess(self.hostname, self.tempdir.tempdir) |
| |
| def Cleanup(self): |
| """Clean up the working/temp directories.""" |
| self.RunCommand(['rm', '-rf', self.work_dir], error_code_ok=True) |
| self.tempdir.Cleanup() |
| |
| def CopyToDevice(self, src, dest, **kwargs): |
| """Copy path to device.""" |
| self.agent.Rsync(src, dest, **kwargs) |
| |
| def CopyFromDevice(self, src, dest, **kwargs): |
| """Copy path from device.""" |
| self.agent.RsyncToLocal(src, dest, **kwargs) |
| |
| def CopyFromWorkDir(self, src, dest, **kwargs): |
| """Copy path from working directory on the device.""" |
| self.CopyFromDevice(os.path.join(self.work_dir, src), dest, **kwargs) |
| |
| def CopyToWorkDir(self, src, dest='', **kwargs): |
| """Copy path to working directory on the device.""" |
| self.CopyToDevice(src, os.path.join(self.work_dir, dest), **kwargs) |
| |
| def PipeOverSSH(self, filepath, cmd, **kwargs): |
| """Cat a file and pipe over SSH.""" |
| producer_cmd = ['cat', filepath] |
| return self.agent.PipeToRemoteSh(producer_cmd, cmd, **kwargs) |
| |
| def Reboot(self): |
| """Reboot the device.""" |
| self.agent.RemoteReboot() |
| |
| def RunCommand(self, cmd, **kwargs): |
| """Executes a shell command on the device with output captured.""" |
| kwargs.setdefault('debug_level', self.debug_level) |
| return self.agent.RemoteSh(cmd, **kwargs) |
