py/toolkit/installer.py

#!/usr/bin/env python3
#
# Copyright 2014 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.

"""Factory toolkit installer.

The factory toolkit is a self-extracting shellball containing factory test
related files and this installer. This installer is invoked when the toolkit
is deployed and is responsible for installing files.
"""

from __future__ import print_function

import argparse
from contextlib import contextmanager
import getpass
import glob
import os
import shutil
import sys
import tempfile
import time

from cros.factory.test.env import paths
from cros.factory.test.test_lists import test_list_common
from cros.factory.tools import install_symlinks
from cros.factory.utils import file_utils
from cros.factory.utils import json_utils
from cros.factory.utils.process_utils import Spawn
from cros.factory.utils import sys_utils


PYTHONPATH = 'usr/local/factory/py_pkg'
INSTALLER_MODULE = 'cros.factory.toolkit.installer'
VERSION_PATH = 'usr/local/factory/TOOLKIT_VERSION'

# Short and sweet help header for the executable generated by makeself.
HELP_HEADER = """
Installs the factory toolkit, transforming a test image into a factory test
image. You can:

- Install the factory toolkit on a CrOS device that is running a test
  image.  To do this, copy install_factory_toolkit.run to the device and
  run it.  The factory tests will then come up on the next boot.

    rsync -a install_factory_toolkit.run crosdevice:/tmp
    ssh crosdevice '/tmp/install_factory_toolkit.run && sync && reboot'

- Modify a test image, turning it into a factory test image.  When you
  use the image on a device, the factory tests will come up.

    install_factory_toolkit.run chromiumos_test_image.bin
"""

HELP_HEADER_ADVANCED = """
- (advanced) Modify a mounted stateful partition, turning it into a factory
  test image.  This is equivalent to the previous command:

    mount_partition -rw chromiumos_test_image.bin 1 /mnt/stateful
    install_factory_toolkit.run /mnt/stateful
    umount /mnt/stateful

- (advanced) Unpack the factory toolkit, modify a file, and then repack it.

    # Unpack but don't actually install
    install_factory_toolkit.run --target /tmp/toolkit --noexec
    # Edit some files in /tmp/toolkit
    emacs /tmp/toolkit/whatever
    # Repack
    install_factory_toolkit.run -- --repack /tmp/toolkit \\
        --pack-into /path/to/new/install_factory_toolkit.run
"""

# The makeself-generated header comes next.  This is a little confusing,
# so explain.
HELP_HEADER_MAKESELF = """
For complete usage information and advanced operations, run
"install_factory_toolkit.run -- --help" (note the extra "--").

Following is the help message from makeself, which was used to create
this self-extracting archive.

-----
"""

SERVER_FILE_MASK = [
    # Exclude Umpire server but keep Umpire client
    '--include', 'py/umpire/__init__.*',
    '--include', 'py/umpire/common.*',
    '--include', 'py/umpire/client',
    '--include', 'py/umpire/client/**',
    '--exclude', 'py/umpire/**',
    '--exclude', 'bin/umpire',
    '--exclude', 'bin/umpired',
]


class FactoryToolkitInstaller:
  """Factory toolkit installer.

  Args:
    src: Source path containing usr/ and var/.
    dest: Installation destination path. Set this to the mount point of the
          stateful partition if patching a test image.
    no_enable: True to not install the tag file.
    system_root: The path to the root of the file system. This must be left
                 as its default value except for unit testing.
    apps: The list of apps to enable/disable under factory/init/main.d/.
    active_test_list: The id of active test list for Goofy.
  """

  # Whether to sudo when rsyncing; set to False for testing.
  _sudo = True

  def __init__(self, src, dest, no_enable, non_cros=False, system_root='/',
               apps=None, active_test_list=None):
    self._src = src
    self._system_root = system_root
    if dest == self._system_root:
      self._usr_local_dest = os.path.join(dest, 'usr', 'local')

      # Make sure we're on a CrOS device.
      if not non_cros and not sys_utils.InCrOSDevice():
        sys.stderr.write(
            "ERROR: You're not on a CrOS device (for more details, please\n"
            'check sys_utils.py:InCrOSDevice), so you must specify a test\n'
            'image or a mounted stateful partition on which to install the\n'
            'factory toolkit.  Please run\n'
            '\n'
            '  install_factory_toolkit.run -- --help\n'
            '\n'
            'for help.\n'
            '\n'
            'If you want to install on a non-CrOS host,\n'
            'please run\n'
            '\n'
            '  install_factory_toolkit.run -- --non-cros \n'
            '\n')
        sys.exit(1)
      if os.getuid() != 0:
        raise Exception('You must be root to install the factory toolkit on a '
                        'CrOS device.')
    else:
      self._usr_local_dest = os.path.join(dest, 'dev_image')
      if not os.path.exists(self._usr_local_dest):
        raise Exception(
            'The destination path %s is not a stateful partition!' % dest)

    self._dest = dest
    self._usr_local_src = os.path.join(src, 'usr', 'local')
    self._no_enable = no_enable
    self._tag_file = os.path.join(self._usr_local_dest, 'factory', 'enabled')

    self._apps = apps
    self._active_test_list = active_test_list

    if not os.path.exists(self._usr_local_src):
      raise Exception(
          'This installer must be run from within the factory toolkit!')

  def WarningMessage(self, target_test_image=None):
    ret = file_utils.ReadFile(os.path.join(self._src, VERSION_PATH))
    if target_test_image:
      ret += (
          '\n'
          '\n'
          '*** You are about to patch the factory toolkit into:\n'
          '***   %s\n'
          '***' % target_test_image)
    else:
      ret += (
          '\n'
          '\n'
          '*** You are about to install the factory toolkit to:\n'
          '***   %s\n'
          '***' % self._dest)
    if self._dest == self._system_root:
      if self._no_enable:
        ret += ('\n*** Factory tests will be disabled after this process is '
                'done, but\n*** you can enable them by creating the factory '
                'enabled tag:\n***   %s\n***' % self._tag_file)
      else:
        ret += ('\n*** After this process is done, your device will start '
                'factory\n*** tests on the next reboot.\n***\n*** Factory '
                'tests can be disabled by deleting the factory enabled\n*** '
                'tag:\n***   %s\n***' % self._tag_file)
    return ret

  def _SetTagFile(self, name, path, enabled):
    """Install or remove a tag file."""
    if enabled:
      print('*** Installing %s enabled tag...' % name)
      Spawn(['touch', path], sudo=True, log=True, check_call=True)
      Spawn(['chmod', 'go+r', path], sudo=True, log=True, check_call=True)
    else:
      print('*** Removing %s enabled tag...' % name)
      Spawn(['rm', '-f', path], sudo=True, log=True, check_call=True)

  def _SetActiveTestList(self):
    """Set the active test list for Goofy."""
    if self._active_test_list is not None:
      path = os.path.join(self._usr_local_dest, 'factory',
                          test_list_common.ACTIVE_TEST_LIST_CONFIG_RELPATH)
      json_utils.DumpFile(
          path,
          test_list_common.GenerateActiveTestListConfig(self._active_test_list))

  def _EnableApp(self, app, enabled):
    """Enable / disable @app.

    In factory/init/startup, a main app is considered disabled if and only:
      1. file "factory/init/main.d/disable-@app" exists OR
      2. file "factory/init/main.d/enable-@app" doesn't exist AND
        file "factory/init/main.d/@app.sh" is not executable.

    Therefore, we enable an app by removing file "disable-@app" and creating
    file "enable-@app", and vice versa.
    """
    app_enable = os.path.join(self._usr_local_dest,
                              'factory', 'init', 'main.d', 'enable-' + app)
    app_disable = os.path.join(self._usr_local_dest,
                               'factory', 'init', 'main.d', 'disable-' + app)
    if enabled:
      print('*** Enabling {app} ***'.format(app=app))
      Spawn(['rm', '-f', app_disable], sudo=self._sudo, log=True,
            check_call=True)
      Spawn(['touch', app_enable], sudo=self._sudo, log=True, check_call=True)
    else:
      print('*** Disabling {app} ***'.format(app=app))
      Spawn(['touch', app_disable], sudo=self._sudo, log=True, check_call=True)
      Spawn(['rm', '-f', app_enable], sudo=self._sudo, log=True,
            check_call=True)

  def _EnableApps(self):
    if not self._apps:
      return

    app_list = []
    for app in self._apps:
      if app[0] == '+':
        app_list.append((app[1:], True))
      elif app[0] == '-':
        app_list.append((app[1:], False))
      else:
        raise ValueError(
            'Use +{app} to enable and -{app} to disable'.format(app=app))

    for app, enabled in app_list:
      self._EnableApp(app, enabled)

  def InstallFactorySubDir(self, sub_dirs):
    """Install only the specified directories under factory folder."""

    def _InstallOneSubDir(sub_dir_name):
      sub_dir_dest = os.path.join(self._usr_local_dest, 'factory', sub_dir_name)
      sub_dir_src = os.path.join(self._src, 'usr', 'local', 'factory',
                                 sub_dir_name)
      try:
        Spawn(['mkdir', '-p', sub_dir_dest], sudo=True, log=True,
              check_call=True)
      except OSError as e:
        print(str(e))
        return

      Spawn(['rsync', '-a', '--force', '-v',
             sub_dir_src + '/', sub_dir_dest],
            sudo=self._sudo, log=True, check_call=True)
      Spawn(['chown', '-R', 'root', sub_dir_dest],
            sudo=self._sudo, log=True, check_call=True)
      Spawn(['chmod', '-R', 'go+rX', sub_dir_dest],
            sudo=self._sudo, log=True, check_call=True)

    for sub_dir_name in sub_dirs:
      _InstallOneSubDir(sub_dir_name)

    self._SetTagFile('factory', self._tag_file, not self._no_enable)
    self._EnableApps()

  def InstallPy3Modules(self):
    if sys_utils.InChroot():
      return

    misc_dir = os.path.join(self._usr_local_dest, 'factory/misc/')
    ws4py_dir = os.path.join(misc_dir, 'ws4py-0.5.1')
    Spawn(['python3', 'setup.py', 'install'], sudo=self._sudo, log=True,
          check_call=True, cwd=ws4py_dir)

  def Install(self):
    print('*** Installing factory toolkit...')

    # --no-owner and --no-group will set owner/group to the current user/group
    # running the command. This is important if we're running with sudo, so
    # the destination will be changed to root/root instead of the user/group
    # before sudo (doesn't matter if sudo is not present). --force is also
    # necessary to allow goofy directory from prior toolkit installations to
    # be overwritten by the goofy symlink.
    print('***   %s -> %s' % (self._usr_local_src, self._usr_local_dest))
    Spawn(['rsync', '-a', '--no-owner', '--no-group', '--chmod=ugo+rX',
           '--force'] + SERVER_FILE_MASK + [self._usr_local_src + '/',
                                            self._usr_local_dest],
          sudo=self._sudo, log=True, check_output=True, cwd=self._usr_local_src)

    print('*** Ensure SSH keys file permission...')
    sshkeys_dir = os.path.join(self._usr_local_dest, 'factory/misc/sshkeys')
    sshkeys = glob.glob(os.path.join(sshkeys_dir, '*'))
    ssh_public_keys = glob.glob(os.path.join(sshkeys_dir, '*.pub'))
    ssh_private_keys = list(set(sshkeys) - set(ssh_public_keys))
    if ssh_private_keys:
      Spawn(['chmod', '600'] + ssh_private_keys, log=True, check_call=True,
            sudo=self._sudo)

    print('*** Installing symlinks...')
    install_symlinks.InstallSymlinks(
        '../factory/bin',
        os.path.join(self._usr_local_dest, 'bin'),
        install_symlinks.MODE_FULL,
        sudo=self._sudo)

    print('*** Installing python3 modules...')
    self.InstallPy3Modules()

    self._SetTagFile('factory', self._tag_file, not self._no_enable)

    self._SetActiveTestList()
    self._EnableApps()

    print('*** Installation completed.')


@contextmanager
def DummyContext(arg):
  """A context manager that simply yields its argument."""
  yield arg


def PrintBuildInfo(src_root):
  """Print build information."""
  info_file = os.path.join(src_root, 'REPO_STATUS')
  if not os.path.exists(info_file):
    raise OSError('Build info file not found!')
  print(file_utils.ReadFile(info_file))


def PackFactoryToolkit(src_root, output_path, initial_version):
  """Packs the files containing this script into a factory toolkit."""
  if initial_version is None:
    complete_version = '%s  repacked by %s@%s at %s\n' % (
        file_utils.ReadFile(os.path.join(src_root, VERSION_PATH)),
        getpass.getuser(), os.uname()[1], time.strftime('%Y-%m-%d %H:%M:%S'))
    initial_version = complete_version.splitlines()[0]
  else:
    complete_version = initial_version + '\n'
  modified_times = len(complete_version.splitlines()) - 1
  if modified_times == 0:
    modified_msg = ''
  else:
    modified_msg = ' (modified %d times)' % modified_times
  with tempfile.NamedTemporaryFile('w') as help_header:
    help_header.write(initial_version + '\n' +
                      HELP_HEADER + HELP_HEADER_MAKESELF)
    help_header.flush()
    cmd = [os.path.join(src_root, 'makeself.sh'), '--bzip2', '--nox11',
           '--help-header', help_header.name,
           src_root, # archive_dir
           output_path, # file_name
           initial_version + modified_msg, # label
           # startup script and args
           # We have to explicitly execute python instead of directly execute
           # INSTALLER_PATH because files under INSTALLER_PATH may not be
           # executable.
           'env', 'PYTHONPATH=' + PYTHONPATH,
           'python3', '-m', INSTALLER_MODULE, '--in-exe']
    Spawn(cmd, check_call=True, log=True)
  with file_utils.TempDirectory() as tmp_dir:
    version_path = os.path.join(tmp_dir, VERSION_PATH)
    os.makedirs(os.path.dirname(version_path))
    file_utils.WriteFile(version_path, complete_version)
    Spawn([cmd[0], '--lsm', version_path, '--append', tmp_dir, output_path],
          check_call=True, log=True)
  print('\n'
        '  Factory toolkit generated at %s.\n'
        '\n'
        '  To install factory toolkit on a live device running a test image,\n'
        '  copy this to the device and execute it as root.\n'
        '\n'
        '  Alternatively, the factory toolkit can be used to patch a test\n'
        '  image. For more information, run:\n'
        '    %s --help\n'
        '\n' % (output_path, output_path))


def ExtractOverlord(src_root, output_dir):
  output_dir = os.path.join(output_dir, 'overlord')
  try:
    os.makedirs(output_dir)
  except OSError as e:
    print(str(e))
    return

  # Copy overlord binary and resource files
  shutil.copyfile(os.path.join(src_root, 'usr/bin/overlordd'),
                  os.path.join(output_dir, 'overlordd'))
  shutil.copytree(os.path.join(src_root, 'usr/share/overlord/app'),
                  os.path.join(output_dir, 'app'))

  # Give overlordd execution permission
  os.chmod(os.path.join(output_dir, 'overlordd'), 0o755)
  print("Extracted overlord under '%s'" % output_dir)


def main():
  import logging
  logging.basicConfig(level=logging.INFO)

  # In order to determine which usage message to show, first determine
  # whether we're in the self-extracting archive.  Do this first
  # because we need it to even parse the arguments.
  if '--in-exe' in sys.argv:
    sys.argv = [x for x in sys.argv if x != '--in-exe']
    in_archive = True
  else:
    in_archive = False

  parser = argparse.ArgumentParser(
      description=HELP_HEADER + HELP_HEADER_ADVANCED,
      usage=('install_factory_toolkit.run -- [options]' if in_archive
             else None),
      formatter_class=argparse.RawDescriptionHelpFormatter)
  parser.add_argument(
      'dest', nargs='?', default='/',
      help='A test image or the mount point of the stateful partition. '
           "If omitted, install to live system, i.e. '/'.")
  parser.add_argument('--no-enable', '-n', action='store_true',
                      help="Don't enable factory tests after installing")
  parser.add_argument('--yes', '-y', action='store_true',
                      help="Don't ask for confirmation")
  parser.add_argument('--build-info', action='store_true',
                      help='Print build information and exit')
  parser.add_argument('--pack-into', metavar='NEW_TOOLKIT',
                      help='Pack the files into a new factory toolkit')
  parser.add_argument('--repack', metavar='UNPACKED_TOOLKIT',
                      help='Repack from previously unpacked toolkit')
  parser.add_argument('--version', metavar='VERSION',
                      help='String to write into TOOLKIT_VERSION when packing')

  parser.add_argument('--non-cros', dest='non_cros',
                      action='store_true',
                      help='Install on non-ChromeOS host.')


  parser.add_argument('--exe-path', dest='exe_path',
                      nargs='?', default=None,
                      help='Current self-extracting archive pathname')
  parser.add_argument('--extract-overlord', dest='extract_overlord',
                      metavar='OUTPUT_DIR', type=str, default=None,
                      help='Extract overlord from the toolkit')
  parser.add_argument('--install-dirs', nargs='+', default=None,
                      help=('Install only the specified directories under '
                            'factory folder. Can be used with --apps to '
                            'enable / disable some apps. Defaults to install '
                            'all folders.'))
  parser.add_argument('--apps', type=lambda s: s.split(','), default=None,
                      help=('Enable or disable some apps under '
                            'factory/init/main.d/. Use prefix "-" to disable, '
                            'prefix "+" to enable, and use "," to separate. '
                            'For example: --apps="-goofy,+whale_servo"'))
  parser.add_argument('--active-test-list', dest='active_test_list',
                      default=None,
                      help='Set the id of active test list for Goofy.')

  args = parser.parse_args()

  src_root = paths.FACTORY_DIR
  for unused_i in range(3):
    src_root = os.path.dirname(src_root)

  if args.extract_overlord is not None:
    ExtractOverlord(src_root, args.extract_overlord)
    return

  # --pack-into may be called directly so this must be done before changing
  # working directory to OLDPWD.
  if args.pack_into and args.repack is None:
    PackFactoryToolkit(src_root, args.pack_into, args.version)
    return

  if not in_archive:
    # If you're not in the self-extracting archive, you're not allowed to
    # do anything except the above --pack-into call.
    parser.error('Not running from install_factory_toolkit.run; '
                 'only --pack-into (without --repack) is allowed')

  # Change to original working directory in case the user specifies
  # a relative path.
  # TODO: Use USER_PWD instead when makeself is upgraded
  os.chdir(os.environ['OLDPWD'])

  if args.repack:
    if args.pack_into is None:
      parser.error('Must specify --pack-into when using --repack.')
    env = dict(os.environ,
               PYTHONPATH=os.path.join(args.repack, PYTHONPATH))
    Spawn(['python3', '-m', INSTALLER_MODULE, '--pack-into', args.pack_into],
          check_call=True, log=True, env=env)
    return

  if args.build_info:
    PrintBuildInfo(src_root)
    return

  if not os.path.exists(args.dest):
    parser.error('Destination %s does not exist!' % args.dest)

  patch_test_image = os.path.isfile(args.dest)

  with (sys_utils.MountPartition(args.dest, 1, rw=True) if patch_test_image
        else DummyContext(args.dest)) as dest:

    installer = FactoryToolkitInstaller(
        src=src_root, dest=dest, no_enable=args.no_enable,
        non_cros=args.non_cros, apps=args.apps,
        active_test_list=args.active_test_list)

    print(installer.WarningMessage(args.dest if patch_test_image else None))

    if not args.yes:
      answer = input('*** Continue? [y/N] ')
      if not answer or answer[0] not in 'yY':
        sys.exit('Aborting.')

    if args.install_dirs:
      installer.InstallFactorySubDir(args.install_dirs)
    else:
      installer.Install()

if __name__ == '__main__':
  # makself interprets "LICENSE" environment variable string as license text and
  # will prompt user to accept before installation. For factory toolkit, we
  # don't want any user interaction in installation and the license is already
  # covered by ebuild or download platform like CPFE.
  os.putenv('LICENSE', '')
  main()