# 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.

"""Helper class for interacting with the Dev Server.

DevServer is the controlling class for all Dev Server interactions. All Dev
Server centric methods are collected here.

Using the locking methods provided in this class, multiple instances of various
tools and methods can be run concurrently with guaranteed safety.
"""

__author__ = 'dalecurtis@google.com (Dale Curtis)'

import logging
import os
import posixpath

# Autotest imports
import common

import common_util


class DevServer(object):
  """Helper class for interacting with the Dev Server.

  All methods assume password-less login for the Dev Server has been set up or
  methods are being run interactively.
  """
  AU_BASE = 'au'
  LATEST = 'LATEST'
  ROOT_UPDATE = 'update.gz'
  STATEFUL_UPDATE = 'stateful.tgz'
  TEST_IMAGE = 'chromiumos_test_image.bin'

  def __init__(self, dev_host, path, user, private_key=None, remote_host=None):
    """Initializes class variables and fixes private key permissions.

    Args:
      @param host: Address of the Dev Server.
      @param path: Images directory root on Dev Server.
      @param user: Dev Server SSH user name.
      @param private_key: Optional private key file for password-less login.
                          If the key file has any group or world permissions
                          they will be removed.
      @param remote_host: If a different hostname/ip should be needed for
                          uploading images to the dev server.
    """
    self._dev_host = dev_host
    if remote_host:
      self._remote_host = remote_host
    else:
      self._remote_host = dev_host
    self._user = user
    self._images = path
    self._private_key = private_key
    if self._private_key:
      # Check that there are no group or world permissions.
      perms = os.stat(self._private_key).st_mode
      if perms & 0o77 != 0:
        # Remove group and world permissions, keep higher order bits.
        os.chmod(self._private_key, (perms >> 6) << 6)
        logging.warning(
            'Removing group and world permissions from private key %s to make'
            ' SSH happy.', self._private_key)


  def UploadAutotestPackages(self, remote_dir, staging_dir):
    """Uploads Autotest packages from staging directory to Dev Server.

    Specifically, the autotest-pkgs directory is uploaded from the staging
    directory to the specified Dev Server.

    Args:
      @param remote_dir: Directory to upload build components into.
      @param staging_dir: Directory containing update.gz and stateful.tgz

    Raises:
      common_util.ChromeOSTestError: If any steps in the process fail to
          complete.
    """
    if os.path.isdir(os.path.join(staging_dir, 'autotest-pkgs')):
      # Upload autotest-pkgs to Dev Server.
      remote_pkgs_dir = posixpath.join(remote_dir, 'autotest')
      msg = 'Failed to upload autotest packages to Dev Server!'
      self.RemoteCopy(src='autotest-pkgs/*', dest=remote_pkgs_dir,
                      cwd=staging_dir, error_msg=msg)

  def UploadBuildComponents(self, remote_dir, staging_dir, upload_image=False):
    """Uploads various build components from staging directory to Dev Server.

    Specifically, the following components are uploaded:
      - update.gz
      - stateful.tgz
      - chromiumos_test_image.bin
      - The entire contents of the au directory. Symlinks are generated for each
        au payload as well.
      - Contents of autotest-pkgs directory.
      - Control files from autotest/server/{tests, site_tests}

    Args:
      @param remote_dir: Directory to upload build components into.
      @param staging_dir: Directory containing update.gz and stateful.tgz
      @param upload_image: Should the chromiumos_test_image.bin be uploaded?

    Raises:
      common_util.ChromeOSTestError: If any steps in the process fail to
          complete.
    """
    upload_list = [self.ROOT_UPDATE, self.STATEFUL_UPDATE]
    if upload_image:
      upload_list.append(self.TEST_IMAGE)
    else:
      # Create blank chromiumos_test_image.bin. Otherwise the Dev Server will
      # try to rebuild it unnecessarily.
      cmd = 'touch ' + posixpath.join(remote_dir, self.TEST_IMAGE)
      msg = 'Failed to create %s on Dev Server!' % self.TEST_IMAGE
      self.RemoteCommand(cmd, error_msg=msg)

    # Upload AU payloads.
    au_path = os.path.join(staging_dir, self.AU_BASE)
    if os.path.isdir(au_path):
      upload_list.append(self.AU_BASE)

      # For each AU payload, setup symlinks to the main payloads.
      cwd = os.getcwd()
      for au in os.listdir(au_path):
        os.chdir(os.path.join(staging_dir, au_path, au))
        os.symlink('../../%s' % self.TEST_IMAGE, self.TEST_IMAGE)
        os.symlink('../../%s' % self.STATEFUL_UPDATE, self.STATEFUL_UPDATE)
        os.chdir(cwd)

    msg = 'Failed to upload build components to the Dev Server!'
    self.RemoteCopy(
        ' '.join(upload_list), dest=remote_dir, cwd=staging_dir, error_msg=msg)

    self.UploadAutotestPackages(remote_dir, staging_dir)

    if os.path.isdir(os.path.join(staging_dir, 'autotest')):
      remote_server_dir = posixpath.join(remote_dir, 'server')
      cmd = 'mkdir -p ' + remote_server_dir
      msg = 'Failed to create autotest server dir on Dev Server!'
      self.RemoteCommand(cmd, error_msg=msg)

      # Upload autotest/server/{tests,site_tests} onto Dev Server.
      msg = 'Failed to upload autotest/server/{tests,site_tests} to Dev Server!'
      self.RemoteCopy(src='autotest/server/{tests,site_tests}',
                      dest=remote_server_dir, cwd=staging_dir, error_msg=msg)

  def AcquireLock(self, tag):
    """Acquires a Dev Server lock for a given tag.

    Creates a directory for the specified tag on the Dev Server, telling other
    components the resource/task represented by the tag is unavailable.

    Args:
      @param tag: Unique resource/task identifier. Use '/' for nested tags.

    Returns:
      Path to the created directory on Dev Server or None if creation failed.

    Raises:
      common_util.ChromeOSTestError: If Dev Server lock can't be acquired.
    """
    remote_dir = posixpath.join(self._images, tag)

    # Attempt to make the directory '<image dir>/<tag>' on the Dev Server. Doing
    # so tells other components that this build is being/was processed by
    # another instance. Directory creation is atomic and will return a non-zero
    # exit code if the directory already exists.
    cmd = 'mkdir ' + remote_dir
    self.RemoteCommand(cmd)
    return remote_dir

  def ReleaseLock(self, tag):
    """Releases Dev Server lock for a given tag. Removes lock directory content.

    Used to release an acquired Dev Server lock. If lock directory is not empty
    the lock will fail to release.

    Args:
      @param tag: Unique resource/task identifier. Use '/' for nested tags.

    Raises:
      common_util.ChromeOSTestError: If processing lock can't be released.
    """
    remote_dir = posixpath.join(self._images, tag)

    cmd = 'rmdir ' + remote_dir
    self.RemoteCommand(cmd)

  def UpdateLatestBuild(self, board, build):
    """Create and upload LATEST file to the Dev Server for the given build.

    If a LATEST file already exists, it's renamed to LATEST.n-1

    Args:
      @param board: Board name for this build; e.g., x86-generic-rel
      @param build: Full build string to look for; e.g., 0.8.61.0-r1cf43296-b269
    """
    try:
      latest_path = posixpath.join(self._images, board, self.LATEST)

      # Update the LATEST file and move any existing LATEST to LATEST.n-1. Use
      # cp instead of mv to prevent any race conditions elsewhere.
      cmd = '[ -f "%s" ] && cp "%s" "%s.n-1"; echo %s>"%s"' % (
          latest_path, latest_path, latest_path, build, latest_path)
      self.RemoteCommand(cmd=cmd)
    except common_util.ChromeOSTestError:
      # Log an error, but don't raise an exception. We don't want to blow away
      # all the work we did just because we can't update the LATEST file.
      logging.error('Could not update %s file for board %s, build %s.',
                    self.LATEST, board, build)

  def GetUpdateUrl(self, board, build):
    """Returns Dev Server update URL for use with memento updater.

    Args:
      @param board: Board name for this build; e.g., x86-generic-rel
      @param build: Full build string to look for; e.g., 0.8.61.0-r1cf43296-b269

    Returns:
      Properly formatted Dev Server update URL.
    """
    return 'http://%s:8080/update/%s/%s' % (self._dev_host, board, build)

  def FindMatchingBoard(self, board):
    """Returns a list of boards given a partial board name.

    Args:
      @param board: Partial board name for this build; e.g., x86-generic

    Returns:
      Returns a list of boards given a partial board and build.
    """
    cmd = 'cd %s; ls -d %s*' % (self._images, board)
    output = self.RemoteCommand(cmd, ignore_errors=True, output=True)
    if output:
      return output.splitlines()
    else:
      return []

  def FindMatchingBuild(self, board, build):
    """Returns a list of matching builds given a board and partial build.

    Args:
      @param board: Partial board name for this build; e.g., x86-generic-rel
      @param build: Partial build string to look for; e.g., 0.8.61.0

    Returns:
      Returns a list of (board, build) tuples given a partial board and build.
    """
    cmd = 'cd %s; find \$(ls -d %s*) -maxdepth 1 -type d -name "%s*"' % (
        self._images, board, build)
    results = self.RemoteCommand(cmd, output=True)
    if results:
      return [tuple(line.split('/')) for line in results.splitlines()]
    else:
      return []

  def RemoteCommand(self, cmd, **kwargs):
    """Wrapper function for executing commands on the Dev Server.

    Args:
      @param cmd: Command to execute on Dev Server.
      @param kwargs: Dicionary of optional args.

    Returns:
      Results from common_util.RunCommand()
    """
    return common_util.RemoteCommand(self._remote_host, self._user, cmd,
                                     private_key=self._private_key, **kwargs)

  def RemoteCopy(self, src, dest, **kwargs):
    """Wrapper function for copying a file to the Dev Server.

    Copies from a local source to a remote destination (on the Dev Server). See
    definition for common_util.RunCommand for complete argument definitions.

    Args:
      @param src: Local path/file.
      @param dest: Remote destination on Dev Server.
      @param kwargs: Dictionary of optional args.

    Returns:
      Results from common_util.RemoteCopy()
    """
    return common_util.RemoteCopy(self._remote_host, self._user, src, dest,
                                  private_key=self._private_key, **kwargs)

  def PrepareDevServer(self, tag, force=False):
    """Prepare Dev Server file system to recieve build components.

    Checks if the component directory for the given build is available and if
    not creates it.

    Args:
      @param tag: Unique resource/task identifier. Use '/' for nested tags.
      @param force: Force re-creation of remote_build_dir even if it already
                    exists.

    Returns:
      Tuple of (remote_build_dir, exists).
          remote_build_dir: The path on Dev Server to the remote build.
          exists: Indicates whether the directory was already present.
    """
    # Check Dev Server for the build before we begin processing.
    remote_build_dir = posixpath.join(self._images, tag)

    # If force is request, delete the existing remote build directory.
    if force:
      self.RemoteCommand('rm -rf ' + remote_build_dir)

    # Create remote directory. Will fail if it already exists.
    exists = self.RemoteCommand('[ -d %s ] && echo exists || mkdir -p %s' % (
        remote_build_dir, remote_build_dir), output=True) == 'exists'

    return remote_build_dir, exists

  def FindDevServerBuild(self, board, build):
    """Given partial build and board ids, figure out the appropriate build.

    Args:
      @param board: Partial board name for this build; e.g., x86-generic
      @param build: Partial build string to look for; e.g., 0.8.61.0

    Returns:
      Tuple of (board, build):
          board: Fully qualified board name; e.g., x86-generic-rel
          build: Fully qualified build string; e.g., 0.8.61.0-r1cf43296-b269

    Raises:
      common_util.ChromeOSTestError: If no boards, no builds, or too many builds
          are matched.
    """
    # Find matching updates on Dev Server.
    if build.lower().strip() == 'latest':
      raise NotImplementedException('FindDevServerBuild no longer supports '
                                    'the value "latest". You must pass a '
                                    'build string to look for.')
    
    builds = self.FindMatchingBuild(board, build)
    if not builds:
      raise common_util.ChromeOSTestError(
          'No builds matching %s could be found for board %s.' % (
              build, board))

    if len(builds) > 1:
      raise common_util.ChromeOSTestError(
          'The given build id is ambiguous. Disambiguate by using one of'
          ' these instead: %s' % ', '.join([b[1] for b in builds]))

    board, build = builds[0]

    return board, build

  def CloneDevServerBuild(self, board, build, tag, force=False):
    """Clone existing Dev Server build. Returns path to cloned build.

    Args:
      @param board: Fully qualified board name; e.g., x86-generic-rel
      @param build: Fully qualified build string; e.g., 0.8.61.0-r1cf43296-b269
      @param tag: Unique resource/task identifier. Use '/' for nested tags.
      @param force: Force re-creation of remote_build_dir even if it already
                    exists.

    Returns:
      The path on Dev Server to the remote build.
    """
    # Prepare the Dev Server for this build.
    remote_build_dir, exists = self.PrepareDevServer(tag, force=force)

    if not exists:
      # Make a copy of the official build, only take necessary files.
      self.RemoteCommand('cp %s %s %s %s' % (
          os.path.join(self._images, board, build, self.TEST_IMAGE),
          os.path.join(self._images, board, build, self.ROOT_UPDATE),
          os.path.join(self._images, board, build, self.STATEFUL_UPDATE),
          remote_build_dir))

    return remote_build_dir

  def GetControlFile(self, board, build, control):
    """Attempts to pull the requested control file from the Dev Server.

    Args:
      @param board: Fully qualified board name; e.g., x86-generic-rel
      @param build: Fully qualified build string; e.g., 0.8.61.0-r1cf43296-b269
      @param control: Path to control file on remote host relative to Autotest
                      root.

    Returns:
      Contents of the control file.

    Raises:
      common_util.ChromeOSTestError: If control file can't be retrieved
    """
    # Create temporary file to target via scp; close.
    return self.RemoteCommand(
        'cat %s' % posixpath.join(self._images, board, build, control),
        output=True)

  def ListAutoupdateTargets(self, board, build):
    """Returns a list of autoupdate test targets for the given board, build.

    Args:
      @param board: Fully qualified board name; e.g., x86-generic-rel
      @param build: Fully qualified build string; e.g., 0.8.61.0-r1cf43296-b269

    Returns:
      List of autoupdate test targets; e.g., ['0.14.747.0-r2bf8859c-b2927_nton']

    Raises:
      common_util.ChromeOSTestError: If control file can't be retrieved
    """
    msg = 'Unable to retrieve list of autoupdate targets!'
    return [os.path.basename(t) for t in self.RemoteCommand(
        'ls -d %s/*' % posixpath.join(self._images, board, build, self.AU_BASE),
        output=True, error_msg=msg).split()]

  def GetImage(self, board, build, staging_dir):
    """Retrieve the TEST_IMAGE for the specified board and build.

    Downloads the image using wget via the Dev Server HTTP interface. The image
    is given a new random file name in the staging directory.

    Args:
      @param board: Fully qualified board name; e.g., x86-generic-rel
      @param build: Fully qualified build string; e.g., 0.8.61.0-r1cf43296-b269
      @param staging_dir: Directory to store downloaded image in.

    Returns:
      File name of the image in the staging directory.
    """
    image_url = '%s/%s' % (
        self.GetUpdateUrl(board, build).replace('update', 'static/archive'),
        self.TEST_IMAGE)
    image_file = self.TEST_IMAGE + '.n-1'
    msg = 'Failed to retrieve the specified image from the Dev Server!'
    common_util.RunCommand(
        'wget -O %s --timeout=30 --tries=1 --no-proxy %s' % (
            image_file, image_url), cwd=staging_dir, error_msg=msg)
    return image_file