py/toolkit/installer.py

#!/usr/bin/python -Bu
#
# Copyright (c) 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.
"""


import argparse
import factory_common  # pylint: disable=W0611
import os
import subprocess
import sys

from cros.factory.test import factory


class FactoryToolkitInstaller():
  """Factory toolkit installer.

  Args:
    src: Source path containing usr/ and var/.
    args: Arguments including
      dest: Installation destination path. Set this to the mount point of the
            stateful partition if patching a test image.
      patch_test_image: True if patching a test image.
  """
  def __init__(self, src, args):
    if args.patch_test_image:
      self._usr_local_dest = os.path.join(args.dest, 'dev_image', 'local')
      self._var_dest = os.path.join(args.dest, 'var_overlay')
      if (not os.path.exists(self._usr_local_dest) or
          not os.path.exists(self._var_dest)):
        raise Exception(
            'The destination path %s is not a stateful partition!' % args.dest)
    else:
      self._usr_local_dest = os.path.join(args.dest, 'usr', 'local')
      self._var_dest = os.path.join(args.dest, 'var')
      if os.getuid() != 0:
        raise Exception('Must be root to install on live machine!')
      if not os.path.exists('/etc/lsb-release'):
        raise Exception('/etc/lsb-release is missing. '
                        'Are you running this in chroot?')

    self._patch_test_image = args.patch_test_image
    self._dest = args.dest
    self._usr_local_src = os.path.join(src, 'usr', 'local')
    self._var_src = os.path.join(src, 'var')

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

  def WarningMessage(self):
    ret = (
        '\n'
        '\n'
        '*** You are about to install factory toolkit to:\n'
        '***   %s\n'
        '***' % self._dest)
    if self._dest == '/':
      ret += ('\n'
        '*** After this process is done, your device will start factory tests\n'
        '*** on the next reboot.\n'
        '***\n'
        '*** Factory tests can be disabled by deleting factory enabled tag:\n'
        '***   /usr/local/factory/enabled\n'
        '***')
    return ret

  def _Rsync(self, src, dest):
    print '***   %s -> %s' % (src, dest)
    subprocess.check_call(['rsync', '-a', src + '/', dest])

  def Install(self):
    print '*** Installing factory toolkit...'
    self._Rsync(self._usr_local_src, self._usr_local_dest)
    self._Rsync(self._var_src, self._var_dest)

    print '*** Installing factory enabled tag...'
    open(os.path.join(self._usr_local_dest, 'factory', 'enabled'), 'w').close()

    print '*** Installation completed.'


def main():
  parser = argparse.ArgumentParser(
      description='Factory toolkit installer.')
  parser.add_argument('--dest', '-d', default='/',
      help='Destination path. Mount point of stateful partition if patching '
           'a test image.')
  parser.add_argument('--patch-test-image', '-p', action='store_true',
      help='Patching a test image instead of installing to live system.')
  parser.add_argument('--yes', '-y', action='store_true',
      help="Don't ask for confirmation")
  args = parser.parse_args()

  try:
    src_root = factory.FACTORY_PATH
    for _ in xrange(3):
      src_root = os.path.dirname(src_root)
    installer = FactoryToolkitInstaller(src_root, args)
  except Exception as e:
    parser.error(e.message)

  print installer.WarningMessage()

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

  installer.Install()

if __name__ == '__main__':
  main()