Add deploy_chrome script.

Allows developers without a chroot to deploy a Chrome package built by
trybots and uploaded to GS.  Design doc for the flow is at goo.gl/zrAHo.

Adapted from cwolfe's cros_chrome_rsync and cros_fetch_image scripts,
documented at http://www/~cwolfe/pilot.

Also adds a remote_access.py library.  This is a re-write of
src/scripts/remote_access.sh in python.

BUG=chromium-os:32575
TEST=Local testing, deploying to device.  Remote trybots.

Change-Id: Iad9cd16735d87dfa7d246c56c2a650723ff308ec
Reviewed-on: https://gerrit.chromium.org/gerrit/27741
Commit-Ready: Ryan Cui <rcui@chromium.org>
Reviewed-by: Ryan Cui <rcui@chromium.org>
Tested-by: Ryan Cui <rcui@chromium.org>
diff --git a/scripts/deploy_chrome.py b/scripts/deploy_chrome.py
new file mode 100644
index 0000000..2c23f02
--- /dev/null
+++ b/scripts/deploy_chrome.py
@@ -0,0 +1,319 @@
+#!/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.
+
+"""Script that resets your Chrome GIT checkout."""
+
+import functools
+import logging
+import optparse
+import os
+import time
+import urlparse
+
+from chromite.lib import cros_build_lib
+from chromite.lib import osutils
+from chromite.lib import remote_access as remote
+from chromite.lib import sudo
+
+
+GS_HTTP = 'https://commondatastorage.googleapis.com'
+GSUTIL_URL = '%s/chromeos-public/gsutil.tar.gz' % GS_HTTP
+GS_RETRIES = 5
+KERNEL_A_PARTITION = 2
+KERNEL_B_PARTITION = 4
+
+KILL_PROC_MAX_WAIT = 10
+POST_KILL_WAIT = 2
+
+
+# Convenience RunCommand methods
+DebugRunCommand = functools.partial(
+    cros_build_lib.RunCommand, debug_level=logging.DEBUG)
+
+DebugRunCommandCaptureOutput = functools.partial(
+    cros_build_lib.RunCommandCaptureOutput, debug_level=logging.DEBUG)
+
+DebugSudoRunCommand = functools.partial(
+    cros_build_lib.SudoRunCommand, debug_level=logging.DEBUG)
+
+
+def _TestGSLs(gs_bin):
+  """Quick test of gsutil functionality."""
+  result = DebugRunCommandCaptureOutput([gs_bin, 'ls'], error_code_ok=True)
+  return not result.returncode
+
+
+def _SetupBotoConfig(gs_bin):
+  """Make sure we can access protected bits in GS."""
+  boto_path = os.path.expanduser('~/.boto')
+  if os.path.isfile(boto_path) or _TestGSLs(gs_bin):
+    return
+
+  logging.info('Configuring gsutil. Please use your @google.com account.')
+  try:
+    cros_build_lib.RunCommand([gs_bin, 'config'], print_cmd=False)
+  finally:
+    if os.path.exists(boto_path) and not os.path.getsize(boto_path):
+      os.remove(boto_path)
+
+
+def _UrlBaseName(url):
+  """Return the last component of the URL."""
+  return url.rstrip('/').rpartition('/')[-1]
+
+
+def _ExtractChrome(src, dest):
+  osutils.SafeMakedirs(dest)
+  # Preserve permissions (-p).  This is default when running tar with 'sudo'.
+  DebugSudoRunCommand(['tar', '--checkpoint', '-xf', src],
+                        cwd=dest)
+
+
+class DeployChrome(object):
+  """Wraps the core deployment functionality."""
+  def __init__(self, options, tempdir):
+    self.tempdir = tempdir
+    self.options = options
+    self.chrome_dir = os.path.join(tempdir, 'chrome')
+    self.host = remote.RemoteAccess(options.to, tempdir, port=options.port)
+    self.start_ui_needed = False
+
+  def _FetchChrome(self):
+    """Get the chrome prebuilt tarball from GS.
+
+    Returns: Path to the fetched chrome tarball.
+    """
+    logging.info('Fetching gsutil.')
+    gsutil_tar = os.path.join(self.tempdir, 'gsutil.tar.gz')
+    cros_build_lib.RunCurl([GSUTIL_URL, '-o', gsutil_tar],
+                           debug_level=logging.DEBUG)
+    DebugRunCommand(['tar', '-xzf', gsutil_tar], cwd=self.tempdir)
+    gs_bin = os.path.join(self.tempdir, 'gsutil', 'gsutil')
+    _SetupBotoConfig(gs_bin)
+    cmd = [gs_bin, 'ls', self.options.gs_path]
+    files = DebugRunCommandCaptureOutput(cmd).output.splitlines()
+    files = [found for found in files if
+             _UrlBaseName(found).startswith('chromeos-chrome-')]
+    if not files:
+      raise Exception('No chrome package found at %s' % self.options.gs_path)
+    elif len(files) > 1:
+      # - Users should provide us with a direct link to either a stripped or
+      #   unstripped chrome package.
+      # - In the case of being provided with an archive directory, where both
+      #   stripped and unstripped chrome available, use the stripped chrome
+      #   package (comes on top after sort).
+      # - Stripped chrome pkg is chromeos-chrome-<version>.tar.gz
+      # - Unstripped chrome pkg is chromeos-chrome-<version>-unstripped.tar.gz.
+      files.sort()
+      cros_build_lib.logger.warning('Multiple chrome packages found.  Using %s',
+                                    files[0])
+
+    filename = _UrlBaseName(files[0])
+    logging.info('Fetching %s.', filename)
+    cros_build_lib.RunCommand([gs_bin, 'cp', files[0], self.tempdir],
+                              print_cmd=False)
+    chrome_path = os.path.join(self.tempdir, filename)
+    assert os.path.exists(chrome_path)
+    return chrome_path
+
+  def _ChromeFileInUse(self):
+    result = self.host.RemoteSh('lsof /opt/google/chrome/chrome',
+                                error_code_ok=True)
+    return result.returncode == 0
+
+  def _DisableRootfsVerification(self):
+    if not self.options.force:
+      logging.error('Detected that the device has rootfs verification enabled.')
+      logging.info('This script can automatically remove the rootfs '
+                   'verification, which requires that it reboot the device.')
+      logging.info('Make sure the device is in developer mode!')
+      logging.info('Skip this prompt by specifying --force.')
+      result = cros_build_lib.YesNoPrompt(
+          'no', prompt='Remove roots verification?')
+      if result == 'no':
+        cros_build_lib.Die('Need rootfs verification to be disabled. '
+                           'Aborting.')
+
+    logging.info('Removing rootfs verification from %s', self.options.to)
+    # Running in VM's cause make_dev_ssd's firmware sanity checks to fail.
+    # Use --force to bypass the checks.
+    cmd = ('/usr/share/vboot/bin/make_dev_ssd.sh --partitions %d '
+           '--remove_rootfs_verification --force')
+    for partition in (KERNEL_A_PARTITION, KERNEL_B_PARTITION):
+      self.host.RemoteSh(cmd % partition, error_code_ok=True)
+
+    # A reboot in developer mode takes a while (and has delays), so the user
+    # will have time to read and act on the USB boot instructions below.
+    logging.info('Please remember to press Ctrl-U if you are booting from USB.')
+    self.host.RemoteReboot()
+
+  def _CheckRootfsWriteable(self):
+    # /proc/mounts is in the format:
+    # <device> <dir> <type> <options>
+    result = self.host.RemoteSh('cat /proc/mounts')
+    for line in result.output.splitlines():
+      components = line.split()
+      if components[0] == '/dev/root' and components[1] == '/':
+        return 'rw' in components[3].split(',')
+    else:
+      raise Exception('Internal error - rootfs mount not found!')
+
+  def _CheckUiJobStarted(self):
+    # status output is in the format:
+    # <job_name> <status> ['process' <pid>].
+    # <status> is in the format <goal>/<state>.
+    result = self.host.RemoteSh('status ui')
+    return result.output.split()[1].split('/')[0] == 'start'
+
+  def _KillProcsIfNeeded(self):
+    if self._CheckUiJobStarted():
+      logging.info('Shutting down Chrome.')
+      self.start_ui_needed = True
+      self.host.RemoteSh('stop ui')
+
+    # Developers sometimes run session_manager manually, in which case we'll
+    # need to help shut the chrome processes down.
+    try:
+      with cros_build_lib.SubCommandTimeout(KILL_PROC_MAX_WAIT):
+        while self._ChromeFileInUse():
+          logging.warning('The chrome binary on the device is in use.')
+          logging.warning('Killing chrome and session_manager processes...\n')
+
+          self.host.RemoteSh("pkill 'chrome|session_manager'",
+                             error_code_ok=True)
+          # Wait for processes to actually terminate
+          time.sleep(POST_KILL_WAIT)
+          logging.info('Rechecking the chrome binary...')
+    except cros_build_lib.TimeoutError:
+      cros_build_lib.Die('Could not kill processes after %s seconds.  Please '
+                         'exit any running chrome processes and try again.')
+
+  def _PrepareTarget(self):
+    # Mount root partition as read/write
+    if not self._CheckRootfsWriteable():
+      logging.info('Mounting rootfs as writeable...')
+      result = self.host.RemoteSh('mount -o remount,rw /', error_code_ok=True)
+      if result.returncode:
+        self._DisableRootfsVerification()
+        logging.info('Trying again to mount rootfs as writeable...')
+        self.host.RemoteSh('mount -o remount,rw /')
+
+      if not self._CheckRootfsWriteable():
+        cros_build_lib.Die('Root partition still read-only')
+
+    # This is needed because we're doing an 'rsync --inplace' of Chrome, but
+    # makes sense to have even when going the sshfs route.
+    self._KillProcsIfNeeded()
+
+  def _Deploy(self):
+    logging.info('Copying Chrome to device.')
+    # Show the output (status) for this command.
+    self.host.Rsync('%s/' % os.path.abspath(self.chrome_dir), '/', inplace=True,
+                    debug_level=logging.INFO)
+    if self.start_ui_needed:
+      self.host.RemoteSh('start ui')
+
+  def Perform(self):
+    try:
+      logging.info('Testing connection to the device.')
+      self.host.RemoteSh('true')
+    except cros_build_lib.RunCommandError:
+      logging.error('Error connecting to the test device.')
+      raise
+
+    pkg_path = self.options.local_path
+    if self.options.gs_path:
+      pkg_path = self._FetchChrome()
+
+    logging.info('Extracting %s.', pkg_path)
+    _ExtractChrome(pkg_path, self.chrome_dir)
+
+    self._PrepareTarget()
+    self._Deploy()
+
+
+def check_gs_path(option, opt, value):
+  """Convert passed-in path to gs:// path."""
+  parsed = urlparse.urlparse(value.rstrip('/ '))
+  # pylint: disable=E1101
+  path = parsed.path.lstrip('/')
+  if parsed.hostname.startswith('sandbox.google.com'):
+    # Sandbox paths are 'storage/<bucket>/<path_to_object>', so strip out the
+    # first component.
+    storage, _, path = path.partition('/')
+    assert storage == 'storage', 'GS URL %s not in expected format.' % value
+
+  return 'gs://%s' % path
+
+
+def check_path(option, opt, value):
+  """Expand the local path"""
+  return osutils.ExpandPath(value)
+
+
+class CustomOption(optparse.Option):
+  """Subclass Option class to implement path evaluation."""
+  TYPES = optparse.Option.TYPES + ('path', 'gs_path')
+  TYPE_CHECKER = optparse.Option.TYPE_CHECKER.copy()
+  TYPE_CHECKER['path'] = check_path
+  TYPE_CHECKER['gs_path'] = check_gs_path
+
+
+def _ParseCommandLine(argv):
+  """Create the parser, parse args, and run environment-independent checks."""
+  usage = 'usage: %prog [--] [command]'
+  parser = optparse.OptionParser(usage=usage, option_class=CustomOption)
+
+  parser.add_option('--force', action='store_true', default=False,
+                    help=('Skip all prompts (i.e., for disabling of rootfs '
+                          'verification).  This may result in the target '
+                          'machine being rebooted.'))
+  parser.add_option('-g', '--gs-path', type='gs_path',
+                    help=('GS path that contains the chrome to deploy.'))
+  parser.add_option('-l', '--local-path', type='path',
+                    help='path to local chrome prebuilt package to deploy.')
+  parser.add_option('-p', '--port', type=int, default=remote.DEFAULT_SSH_PORT,
+                    help=('Port of the target device to connect to.'))
+  parser.add_option('-t', '--to',
+                    help=('The IP address of the CrOS device to deploy to.'))
+  parser.add_option('-v', '--verbose', action='store_true', default=False,
+                    help=('Show more debug output.'))
+
+  (options, args) = parser.parse_args(argv)
+
+  if not options.gs_path and not options.local_path:
+    parser.error('Need to specify either --gs-path or --local-path')
+  if options.gs_path and options.local_path:
+    parser.error('Cannot specify both --gs-path and --local-path')
+  if not options.to:
+    parser.error('Need to specify --to')
+
+  return options, args
+
+
+def _PostParseCheck(options, args):
+  """Perform some usage validation (after we've parsed the arguments
+
+  Args:
+    options/args: The options/args object returned by optparse
+  """
+  if options.local_path and not os.path.isfile(options.local_path):
+    cros_build_lib.Die('%s is not a file.', options.local_path)
+
+
+def main(argv):
+  options, args = _ParseCommandLine(argv)
+  _PostParseCheck(options, args)
+
+  # Set cros_build_lib debug level to hide RunCommand spew.
+  if options.verbose:
+    cros_build_lib.logger.setLevel(logging.DEBUG)
+  else:
+    cros_build_lib.logger.setLevel(logging.INFO)
+
+  with sudo.SudoKeepAlive():
+    with osutils.TempDirContextManager(sudo_rm=True) as tempdir:
+      deploy = DeployChrome(options, tempdir)
+      deploy.Perform()