Improve factory toolkit command line usage

For convenience, the command line argument syntax is changed:

  ./install_factory_toolkit.run
    - Installs to live system. Stops if not running as root or in chroot
      enviroment.

  ./install_factory_toolkit.run <path/to/test_image.bin>
    - Patches the factory toolkit into a test image.

  ./install_factory_toolkit.run <path/to/stateful_partition_mount_point>
    - Checks if the path is a mount point of the stateful partition of a
      test image. Patches the factory toolkit into the path if it is.

There are two advanced arguments: --no-enable and --yes. They can still
be used by appending '--' before them. For example,

  ./install_factory_toolkit.run -- --yes <path/to/test_image.bin>

will patch the test image without asking for confirmation.

BUG=chrome-os-partner:25941
TEST=Install to a live system
TEST=Patch a test image
TEST=Install to a mount point of the stateful partition
TEST=Try to install to an artitrary path and see error message

Change-Id: Ic4fc92b732f0d9e9dde717e372f452255f410bbb
Signed-off-by: Vic (Chun-Ju) Yang <victoryang@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/186950
diff --git a/py/toolkit/installer.py b/py/toolkit/installer.py
index b90ee4c..3af3c51 100755
--- a/py/toolkit/installer.py
+++ b/py/toolkit/installer.py
@@ -13,12 +13,14 @@
 
 
 import argparse
-import factory_common  # pylint: disable=W0611
+from contextlib import contextmanager
 import os
-import subprocess
 import sys
 
+import factory_common  # pylint: disable=W0611
 from cros.factory.test import factory
+from cros.factory.tools.mount_partition import MountPartition
+from cros.factory.utils.process_utils import Spawn
 
 
 class FactoryToolkitInstaller():
@@ -26,33 +28,35 @@
 
   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.
+    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.
   """
-  def __init__(self, src, args):
-    if args.patch_test_image:
-      self._usr_local_dest = os.path.join(args.dest, 'dev_image')
-      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')
+
+  def __init__(self, src, dest, no_enable, system_root='/'):
+    self._system_root = system_root
+    if dest == self._system_root:
+      self._usr_local_dest = os.path.join(dest, 'usr', 'local')
+      self._var_dest = os.path.join(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?')
+    else:
+      self._usr_local_dest = os.path.join(dest, 'dev_image')
+      self._var_dest = os.path.join(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!' % dest)
 
-    self._patch_test_image = args.patch_test_image
-    self._dest = args.dest
+    self._dest = dest
     self._usr_local_src = os.path.join(src, 'usr', 'local')
     self._var_src = os.path.join(src, 'var')
-    self._no_enable = args.no_enable
+    self._no_enable = no_enable
     self._tag_file = os.path.join(self._usr_local_dest, 'factory', 'enabled')
 
     if (not os.path.exists(self._usr_local_src) or
@@ -60,14 +64,22 @@
       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 == '/':
+  def WarningMessage(self, target_test_image=None):
+    if target_test_image:
+      ret = (
+          '\n'
+          '\n'
+          '*** You are about to patch factory toolkit into:\n'
+          '***   %s\n'
+          '***' % target_test_image)
+    else:
+      ret = (
+          '\n'
+          '\n'
+          '*** You are about to install 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'
@@ -86,7 +98,8 @@
 
   def _Rsync(self, src, dest):
     print '***   %s -> %s' % (src, dest)
-    subprocess.check_call(['rsync', '-a', src + '/', dest])
+    Spawn(['rsync', '-a', src + '/', dest],
+          sudo=True, log=True, check_output=True)
 
   def Install(self):
     print '*** Installing factory toolkit...'
@@ -106,36 +119,49 @@
     print '*** Installation completed.'
 
 
+@contextmanager
+def DummyContext(arg):
+  """A context manager that simply yields its argument."""
+  yield arg
+
+
 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('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")
   args = parser.parse_args()
 
-  try:
+  # 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 not os.path.exists(args.dest):
+    parser.error('Destination %s does not exist!' % args.dest)
+
+  patch_test_image = os.path.isfile(args.dest)
+
+  with (MountPartition(args.dest, 1, rw=True) if patch_test_image
+        else DummyContext(args.dest)) as dest:
     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)
+    installer = FactoryToolkitInstaller(src_root, dest, args.no_enable)
 
-  print installer.WarningMessage()
+    print installer.WarningMessage(args.dest if patch_test_image else None)
 
-  if not args.yes:
-    answer = raw_input('*** Continue? [y/N] ')
-    if not answer or answer[0] not in 'yY':
-      sys.exit('Aborting.')
+    if not args.yes:
+      answer = raw_input('*** Continue? [y/N] ')
+      if not answer or answer[0] not in 'yY':
+        sys.exit('Aborting.')
 
-  installer.Install()
+    installer.Install()
 
 if __name__ == '__main__':
   main()