#!/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)