# -*- coding: utf-8 -*-
# Copyright (c) 2013 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.

"""VM-related helper functions/classes."""

from __future__ import print_function

import distutils.version  # pylint: disable=import-error,no-name-in-module
import errno
import fcntl
import glob
import multiprocessing
import os
import re
import shutil
import socket
import sys
import time

from chromite.cli.cros import cros_chrome_sdk
from chromite.lib import constants
from chromite.lib import cros_build_lib
from chromite.lib import cros_logging as logging
from chromite.lib import device
from chromite.lib import image_lib
from chromite.lib import osutils
from chromite.lib import path_util
from chromite.lib import remote_access
from chromite.lib import retry_util
from chromite.utils import memoize


assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'


class VMError(device.DeviceError):
  """Exception for VM errors."""


def VMIsUpdatable(path):
  """Check if the existing VM image is updatable.

  Args:
    path: Path to the VM image.

  Returns:
    True if VM is updatable; False otherwise.
  """
  table = {p.name: p for p in image_lib.GetImageDiskPartitionInfo(path)}
  # Assume if size of the two root partitions match, the image
  # is updatable.
  return table['ROOT-B'].size == table['ROOT-A'].size


def CreateVMImage(image=None, board=None, updatable=True, dest_dir=None):
  """Returns the path of the image built to run in a VM.

  By default, the returned VM is a test image that can run full update
  testing on it. If there exists a VM image with the matching
  |updatable| setting, this method returns the path to the existing
  image. If |dest_dir| is set, it will copy/create the VM image to the
  |dest_dir|.

  Args:
    image: Path to the (non-VM) image. Defaults to None to use the latest
      image for the board.
    board: Board that the image was built with. If None, attempts to use the
      configured default board.
    updatable: Create a VM image that supports AU.
    dest_dir: If set, create/copy the VM image to |dest|; otherwise,
      use the folder where |image| resides.
  """
  if not image and not board:
    raise VMError('Cannot create VM when both image and board are None.')

  image_dir = os.path.dirname(image)
  src_path = dest_path = os.path.join(image_dir, constants.VM_IMAGE_BIN)

  if dest_dir:
    dest_path = os.path.join(dest_dir, constants.VM_IMAGE_BIN)

  exists = False
  # Do not create a new VM image if a matching image already exists.
  exists = os.path.exists(src_path) and (
      not updatable or VMIsUpdatable(src_path))

  if exists and dest_dir:
    # Copy the existing VM image to dest_dir.
    shutil.copyfile(src_path, dest_path)

  if not exists:
    # No existing VM image that we can reuse. Create a new VM image.
    logging.info('Creating %s', dest_path)
    cmd = [os.path.join(constants.CROSUTILS_DIR, 'image_to_vm.sh'),
           '--test_image']

    if image:
      cmd.append('--from=%s' % path_util.ToChrootPath(image_dir))

    if updatable:
      cmd.extend(['--disk_layout', '2gb-rootfs-updatable'])

    if board:
      cmd.extend(['--board', board])

    # image_to_vm.sh only runs in chroot, but dest_dir may not be
    # reachable from chroot. In that case, we copy it to a temporary
    # directory in chroot, and then move it to dest_dir .
    tempdir = None
    if dest_dir:
      # Create a temporary directory in chroot to store the VM
      # image. This is to avoid the case where dest_dir is not
      # reachable within chroot.
      tempdir = cros_build_lib.run(
          ['mktemp', '-d'],
          capture_output=True,
          enter_chroot=True).output.strip()
      cmd.append('--to=%s' % tempdir)

    msg = 'Failed to create the VM image'
    try:
      cros_build_lib.run(cmd, enter_chroot=True, cwd=constants.SOURCE_ROOT)
    except cros_build_lib.RunCommandError as e:
      logging.error('%s: %s', msg, e)
      if tempdir:
        osutils.RmDir(
            path_util.FromChrootPath(tempdir), ignore_missing=True)
      raise VMError(msg)

    if dest_dir:
      # Move VM from tempdir to dest_dir.
      shutil.move(
          path_util.FromChrootPath(
              os.path.join(tempdir, constants.VM_IMAGE_BIN)), dest_path)
      osutils.RmDir(path_util.FromChrootPath(tempdir), ignore_missing=True)

  if not os.path.exists(dest_path):
    raise VMError(msg)

  return dest_path


class VM(device.Device):
  """Class for managing a VM."""

  SSH_PORT = 9222
  SSH_NON_KVM_CONNECT_TIMEOUT = 240
  IMAGE_FORMAT = 'raw'
  # kvm_* should match kvm_intel, kvm_amd, etc.
  NESTED_KVM_GLOB = '/sys/module/kvm_*/parameters/nested'
  FLASH_SIZE = 64*1024*1024

  # Target architecture
  ARCH_X86_64 = 'x86_64'
  ARCH_AARCH64 = 'aarch64'

  def __init__(self, opts):
    """Initialize VM.

    Args:
      opts: command line options.
    """
    super(VM, self).__init__(opts)

    self.qemu_path = opts.qemu_path
    self.qemu_img_path = opts.qemu_img_path
    self.qemu_bios_path = opts.qemu_bios_path
    self.qemu_arch = opts.qemu_arch
    self.qemu_m = opts.qemu_m
    self.qemu_cpu = opts.qemu_cpu
    if self.qemu_cpu is None:
      # TODO(pwang): replace SandyBridge to Haswell-noTSX once lab machine
      # running VMTest all migrate to GCE.
      if self.qemu_arch == VM.ARCH_X86_64:
        self.qemu_cpu = 'SandyBridge,-invpcid,-tsc-deadline'
      elif self.qemu_arch == VM.ARCH_AARCH64:
        self.qemu_cpu = 'cortex-a57'
    self.qemu_arch = opts.qemu_arch
    self.qemu_smp = opts.qemu_smp
    if self.qemu_smp == 0:
      self.qemu_smp = min(8, multiprocessing.cpu_count())
    self.qemu_hostfwd = opts.qemu_hostfwd
    self.qemu_args = opts.qemu_args

    if self.qemu_arch == VM.ARCH_X86_64:
      if opts.enable_kvm is None:
        self.enable_kvm = os.path.exists('/dev/kvm')
      else:
        self.enable_kvm = opts.enable_kvm
    else:
      self.enable_kvm = False
    self.copy_on_write = opts.copy_on_write
    # We don't need sudo access for software emulation or if /dev/kvm is
    # writeable.
    self.use_sudo = self.enable_kvm and not os.access('/dev/kvm', os.W_OK)
    self.display = opts.display
    self.image_path = opts.image_path
    self.image_format = opts.image_format

    self.device = remote_access.LOCALHOST
    self.ssh_port = self.ssh_port or opts.ssh_port or VM.SSH_PORT

    self.start = opts.start
    self.stop = opts.stop

    self.chroot_path = opts.chroot_path

    self.cache_dir = os.path.abspath(opts.cache_dir)
    assert os.path.isdir(self.cache_dir), "Cache directory doesn't exist"

    self.vm_dir = opts.vm_dir
    if not self.vm_dir:
      self.vm_dir = os.path.join(osutils.GetGlobalTempDir(),
                                 'cros_vm_%d' % self.ssh_port)
    self._CreateVMDir()

    self.pidfile = os.path.join(self.vm_dir, 'kvm.pid')
    self.kvm_monitor = os.path.join(self.vm_dir, 'kvm.monitor')
    self.kvm_pipe_in = '%s.in' % self.kvm_monitor  # to KVM
    self.kvm_pipe_out = '%s.out' % self.kvm_monitor  # from KVM
    self.kvm_serial = '%s.serial' % self.kvm_monitor
    self.flash0_file = os.path.join(self.vm_dir, 'flash0.img')
    self.flash1_file = os.path.join(self.vm_dir, 'flash1.img')

    self.copy_image_on_shutdown = False
    self.image_copy_dir = None

    # Wait 2 min for non-KVM.
    connect_timeout = (VM.SSH_CONNECT_TIMEOUT if self.enable_kvm else
                       VM.SSH_NON_KVM_CONNECT_TIMEOUT)
    self.InitRemote(connect_timeout=connect_timeout)

  def _CreateVMDir(self):
    """Safely create vm_dir."""
    if not osutils.SafeMakedirs(self.vm_dir):
      # For security, ensure that vm_dir is not a symlink, and is owned by us.
      error_str = ('VM state dir is misconfigured; please recreate: %s'
                   % self.vm_dir)
      assert os.path.isdir(self.vm_dir), error_str
      assert not os.path.islink(self.vm_dir), error_str
      assert os.stat(self.vm_dir).st_uid == os.getuid(), error_str

  def _CreatePflashFiles(self):
    """Creates parallel flash images in the temporary VM dir.

    One image is used for the UEFI firmware, the other as non-volatile
    storage for UEFI variables.

    These images will get removed on VM shutdown.

    Returns:
      Nothing
    """
    uefi_path = '/usr/share/qemu/edk2-aarch64-code.fd'
    cros_build_lib.run(
        ['dd', 'if=%s' % uefi_path, 'of=%s' % self.flash0_file,
        'count=1', 'bs=%dM' % (VM.FLASH_SIZE / (1024*1024)), 'conv=sync',
        ], dryrun=self.dryrun)
    logging.info('flash0 image created at %s.', self.flash0_file)

    with open(self.flash1_file, 'wb+') as f:
      f.truncate(VM.FLASH_SIZE)
    logging.info('flash1 image created at %s.', self.flash1_file)

  def _CreateQcow2Image(self):
    """Creates a qcow2-formatted image in the temporary VM dir.

    This image will get removed on VM shutdown.

    Returns:
      Tuple of (path to qcow2 image, format of qcow2 image)
    """
    cow_image_path = os.path.join(self.vm_dir, 'qcow2.img')
    qemu_img_args = [
        self.qemu_img_path,
        'create', '-f', 'qcow2',
        '-o', 'backing_file=%s' % self.image_path,
        cow_image_path,
    ]
    cros_build_lib.run(qemu_img_args, dryrun=self.dryrun)
    logging.info('qcow2 image created at %s.', cow_image_path)
    return cow_image_path, 'qcow2'

  def _RmVMDir(self):
    """Cleanup vm_dir."""
    osutils.RmDir(self.vm_dir, ignore_missing=True, sudo=self.use_sudo)

  @memoize.MemoizedSingleCall
  def QemuVersion(self):
    """Determine QEMU version.

    Returns:
      QEMU version.
    """
    version_str = cros_build_lib.run([self.qemu_path, '--version'],
                                     capture_output=True, dryrun=self.dryrun,
                                     encoding='utf-8').stdout
    # version string looks like one of these:
    # QEMU emulator version 2.0.0 (Debian 2.0.0+dfsg-2ubuntu1.36), Copyright (c)
    # 2003-2008 Fabrice Bellard
    #
    # QEMU emulator version 2.6.0, Copyright (c) 2003-2008 Fabrice Bellard
    #
    # qemu-x86_64 version 2.10.1
    # Copyright (c) 2003-2017 Fabrice Bellard and the QEMU Project developers
    m = re.search(r'version ([0-9.]+)', version_str)
    if not m:
      raise VMError('Unable to determine QEMU version from:\n%s.' % version_str)
    return m.group(1)

  def _CheckQemuMinVersion(self):
    """Ensure minimum QEMU version."""
    if self.dryrun:
      return
    min_qemu_version = '2.6.0'
    logging.info('QEMU version %s', self.QemuVersion())
    LooseVersion = distutils.version.LooseVersion
    if LooseVersion(self.QemuVersion()) < LooseVersion(min_qemu_version):
      raise VMError('QEMU %s is the minimum supported version. You have %s.'
                    % (min_qemu_version, self.QemuVersion()))

  def _SetQemuPath(self):
    """Find a suitable Qemu executable."""
    if self.qemu_arch == VM.ARCH_X86_64:
      qemu_exe = 'qemu-system-x86_64'
    elif self.qemu_arch == VM.ARCH_AARCH64:
      qemu_exe = 'qemu-system-aarch64'
    qemu_exe_path = os.path.join('usr/bin', qemu_exe)

    # Check SDK cache.
    if not self.qemu_path:
      qemu_dir = cros_chrome_sdk.SDKFetcher.GetCachePath(
          cros_chrome_sdk.SDKFetcher.QEMU_BIN_PATH, self.cache_dir, self.board)
      if qemu_dir:
        qemu_path = os.path.join(qemu_dir, qemu_exe_path)
        if os.path.isfile(qemu_path):
          self.qemu_path = qemu_path

    # Check chroot.
    if not self.qemu_path:
      qemu_path = os.path.join(self.chroot_path, qemu_exe_path)
      if os.path.isfile(qemu_path):
        self.qemu_path = qemu_path

    # Check system.
    if not self.qemu_path:
      logging.warning('Using system QEMU.')
      self.qemu_path = osutils.Which(qemu_exe)

    if not self.qemu_path or not os.path.isfile(self.qemu_path):
      raise VMError('QEMU not found.')

    if self.copy_on_write:
      if not self.qemu_img_path:
        # Look for qemu-img right next to qemu-system-x86_64.
        self.qemu_img_path = os.path.join(
            os.path.dirname(self.qemu_path), 'qemu-img')
      if not os.path.isfile(self.qemu_img_path):
        raise VMError('qemu-img not found. (Needed to create qcow2 image).')

    logging.debug('QEMU path: %s', self.qemu_path)
    self._CheckQemuMinVersion()

  def _GetBuiltVMImagePath(self):
    """Get path of a locally built VM image."""
    vm_image_path = os.path.join(
        constants.SOURCE_ROOT, 'src/build/images', self.board,
        'latest', constants.VM_IMAGE_BIN)
    return vm_image_path if os.path.isfile(vm_image_path) else None

  def _GetCacheVMImagePath(self):
    """Get path of a cached VM image."""
    cache_path = cros_chrome_sdk.SDKFetcher.GetCachePath(
        constants.VM_IMAGE_TAR, self.cache_dir, self.board)
    if cache_path:
      vm_image = os.path.join(cache_path, constants.VM_IMAGE_BIN)
      if os.path.isfile(vm_image):
        return vm_image
    return None

  def _SetVMImagePath(self):
    """Detect VM image path in SDK and chroot."""
    if not self.image_path:
      self.image_path = (self._GetCacheVMImagePath() or
                         self._GetBuiltVMImagePath())
    if not self.image_path:
      raise VMError('No VM image found. Use cros chrome-sdk --download-vm.')
    if not os.path.isfile(self.image_path):
      # Checks if the image path points to a directory containing the bin file.
      image_path = os.path.join(self.image_path, constants.VM_IMAGE_BIN)
      if os.path.isfile(image_path):
        self.image_path = image_path
      else:
        raise VMError('VM image does not exist: %s' % self.image_path)
    logging.debug('VM image path: %s', self.image_path)

  def _SetBoard(self):
    """Sets the board.

    Picks the first non-None board from the user-specified board,
    SDK environment variable, cros default board.

    Raises:
      DieSystemExit: If a board cannot be found.
    """
    if self.board:
      return
    sdk_board_env = os.environ.get(cros_chrome_sdk.SDKFetcher.SDK_BOARD_ENV)
    self.board = cros_build_lib.GetBoard(sdk_board_env, strict=True)

  def _WaitForSSHPort(self, sleep=5):
    """Wait for SSH port to become available."""
    class _SSHPortInUseError(Exception):
      """Exception for _CheckSSHPortBusy to throw."""

    def _CheckSSHPortBusy(ssh_port):
      """Check if the SSH port is in use."""
      sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      try:
        sock.bind((remote_access.LOCALHOST_IP, ssh_port))
      except socket.error as e:
        if e.errno == errno.EADDRINUSE:
          logging.info('SSH port %d in use...', self.ssh_port)
          raise _SSHPortInUseError()
      finally:
        sock.close()

    try:
      retry_util.RetryException(
          exception=_SSHPortInUseError,
          max_retry=10,
          functor=lambda: _CheckSSHPortBusy(self.ssh_port),
          sleep=sleep)
    except _SSHPortInUseError:
      raise VMError('SSH port %d in use' % self.ssh_port)

  def Run(self):
    """Performs an action, one of start, stop, or run a command in the VM."""
    if not self.start and not self.stop and not self.cmd:
      raise VMError('Must specify one of start, stop, or cmd.')

    if self.start:
      self.Start()
    if self.cmd:
      if not self.IsRunning():
        raise VMError('VM not running.')
      self.remote_run(self.cmd, stream_output=True)
    if self.stop:
      self.Stop()

  def _GetQemuArgs(self, image_path, image_format):
    """Returns the args to qemu used to launch the VM.

    Args:
      image_path: Path to QEMU image.
      image_format: Format of the image.
    """
    # Append 'check' to warn if the requested CPU is not fully supported.
    logging.info('CPU: %s', str(self.qemu_cpu))
    if self.qemu_arch == VM.ARCH_X86_64:
      if 'check' not in self.qemu_cpu.split(','):
        self.qemu_cpu += ',check'
    # Append 'vmx=on' if the host supports nested virtualization. It can be
    # enabled via 'vmx+' or 'vmx=on' (or similarly disabled) so just test for
    # the presence of 'vmx'. For more details, see:
    # https://www.kernel.org/doc/Documentation/virtual/kvm/nested-vmx.txt
    if 'vmx' not in self.qemu_cpu and self.enable_kvm:
      for f in glob.glob(self.NESTED_KVM_GLOB):
        if cros_build_lib.BooleanShellValue(osutils.ReadFile(f).strip(), False):
          self.qemu_cpu += ',vmx=on'
          break

    qemu_args = [self.qemu_path]
    if self.qemu_bios_path:
      if not os.path.isdir(self.qemu_bios_path):
        raise VMError('Invalid QEMU bios path: %s' % self.qemu_bios_path)
      qemu_args += ['-L', self.qemu_bios_path]

    qemu_args += [
        '-m', self.qemu_m, '-smp', str(self.qemu_smp),
        '-daemonize',
        '-pidfile', self.pidfile,
        '-chardev', 'pipe,id=control_pipe,path=%s' % self.kvm_monitor,
        '-serial', 'file:%s' % self.kvm_serial,
        '-mon', 'chardev=control_pipe',
        '-device', 'virtio-net,netdev=eth0',
        '-device', 'virtio-scsi-pci,id=scsi',
        '-device', 'virtio-rng',
        '-device', 'scsi-hd,drive=hd,bootindex=0',
        '-drive', 'if=none,id=hd,file=%s,cache=unsafe,format=%s'
        % (image_path, image_format),
    ]
    if self.qemu_arch == VM.ARCH_X86_64:
      qemu_args += [
        '-vga', 'virtio',
        '-usb', '-device', 'usb-tablet',
      ]
    if self.qemu_arch == VM.ARCH_AARCH64:
      qemu_args += [
          '-M', 'virt',
          '-pflash', self.flash0_file,
          '-pflash', self.flash1_file,
      ]
    if self.qemu_cpu:
      qemu_args += [ '-cpu', self.qemu_cpu ]
    # netdev args, including hostfwds.
    netdev_args = ('user,id=eth0,net=10.0.2.0/27,hostfwd=tcp:%s:%d-:%d'
                   % (remote_access.LOCALHOST_IP, self.ssh_port,
                      remote_access.DEFAULT_SSH_PORT))
    if self.qemu_hostfwd:
      for hostfwd in self.qemu_hostfwd:
        netdev_args += ',hostfwd=%s' % hostfwd
    qemu_args += ['-netdev', netdev_args]

    if self.qemu_args:
      for arg in self.qemu_args:
        qemu_args += arg.split()
    if self.enable_kvm:
      qemu_args += ['-enable-kvm']
    if not self.display:
      qemu_args += ['-display', 'none']

    return qemu_args

  def Start(self, retries=1):
    """Start the VM.

    Args:
      retries: Number of times to retry launching the VM if it fails to boot-up.
    """
    if not self.enable_kvm:
      logging.warning('KVM is not supported; Chrome VM will be slow')
    self._SetBoard()
    self._SetQemuPath()
    self._SetVMImagePath()
    logging.info('Pid file: %s', self.pidfile)

    for attempt in range(0, retries + 1):
      self.Stop()

      logging.debug('Start VM, attempt #%d', attempt)

      self._CreateVMDir()
      if self.qemu_arch == VM.ARCH_AARCH64:
        self._CreatePflashFiles()

      image_path = self.image_path
      image_format = self.image_format
      if self.copy_on_write:
        image_path, image_format = self._CreateQcow2Image()
      qemu_args = self._GetQemuArgs(image_path, image_format)
      # Make sure we can read these files later on by creating them as
      # ourselves.
      osutils.Touch(self.kvm_serial)
      for pipe in [self.kvm_pipe_in, self.kvm_pipe_out]:
        os.mkfifo(pipe, 0o600)
      osutils.Touch(self.pidfile)

      # Add use_sudo support to cros_build_lib.run.
      run = cros_build_lib.sudo_run if self.use_sudo else cros_build_lib.run
      run(qemu_args, dryrun=self.dryrun)
      try:
        self.WaitForBoot()
        return
      except device.DeviceError:
        if attempt == retries:
          raise
        else:
          logging.warning('Error when launching VM. Retrying...')

  def _GetVMPid(self):
    """Get the pid of the VM.

    Returns:
      pid of the VM.
    """
    if not os.path.exists(self.vm_dir):
      logging.debug('%s not present.', self.vm_dir)
      return 0

    if not os.path.exists(self.pidfile):
      logging.info('%s does not exist.', self.pidfile)
      return 0

    pid = osutils.ReadFile(self.pidfile).rstrip()
    if not pid.isdigit():
      # Ignore blank/empty files.
      if pid:
        logging.error('%s in %s is not a pid.', pid, self.pidfile)
      return 0

    return int(pid)

  def IsRunning(self):
    """Returns True if there's a running VM.

    Returns:
      True if there's a running VM.
    """
    pid = self._GetVMPid()
    if not pid:
      return False

    # Make sure the process actually exists.
    return os.path.isdir('/proc/%i' % pid)

  def SaveVMImageOnShutdown(self, output_dir):
    """Takes a VM snapshot via savevm and signals to save the VM image later.

    Args:
      output_dir: A path specifying the directory that the VM image should be
          saved to.
    """
    logging.debug('Taking VM snapshot')
    self.copy_image_on_shutdown = True
    self.image_copy_dir = output_dir
    if not self.copy_on_write:
      logging.warning(
          'Attempting to take a VM snapshot without --copy-on-write. Saved '
          'VM image may not contain the desired snapshot.')
    with open(self.kvm_pipe_in, 'w') as monitor_pipe:
      # Saving the snapshot will take an indeterminate amount of time, so also
      # send a fake command that the monitor will complain about so we can know
      # when the snapshot saving is done.
      monitor_pipe.write('savevm chromite_lib_vm_snapshot\n')
      monitor_pipe.write('thisisafakecommand\n')
    with open(self.kvm_pipe_out) as monitor_pipe:
      # Set reads to be non-blocking
      fd = monitor_pipe.fileno()
      cur_flags = fcntl.fcntl(fd, fcntl.F_GETFL)
      fcntl.fcntl(fd, fcntl.F_SETFL, cur_flags | os.O_NONBLOCK)
      # 30 second timeout.
      success = False
      start_time = time.time()
      while time.time() - start_time < 30:
        try:
          line = monitor_pipe.readline()
          if 'thisisafakecommand' in line:
            logging.debug('Finished attempting to take VM snapshot')
            success = True
            break
          logging.debug('VM monitor output: %s', line)
        except IOError:
          time.sleep(1)
      if not success:
        logging.warning('Timed out trying to take VM snapshot')

  def _KillVM(self):
    """Kill the VM process."""
    pid = self._GetVMPid()
    if pid:
      # Add use_sudo support to cros_build_lib.run.
      run = cros_build_lib.sudo_run if self.use_sudo else cros_build_lib.run
      run(['kill', '-9', str(pid)], check=False, dryrun=self.dryrun)

  def _MaybeCopyVMImage(self):
    """Saves the VM image to a location on disk if previously told to."""
    if not self.copy_image_on_shutdown:
      return
    if not self.image_copy_dir:
      logging.debug('Told to copy VM image, but no output directory set')
      return
    shutil.copy(self.image_path, os.path.join(
        self.image_copy_dir, os.path.basename(self.image_path)))

  def Stop(self):
    """Stop the VM."""
    logging.debug('Stop VM')

    self._KillVM()
    self._WaitForSSHPort()
    self._MaybeCopyVMImage()
    self._RmVMDir()

  def _WaitForProcs(self, sleep=2):
    """Wait for expected processes to launch."""
    class _TooFewPidsException(Exception):
      """Exception for _GetRunningPids to throw."""

    def _GetRunningPids(exe, numpids):
      pids = self.remote.GetRunningPids(exe, full_path=False)
      logging.info('%s pids: %s', exe, repr(pids))
      if len(pids) < numpids:
        raise _TooFewPidsException()

    def _WaitForProc(exe, numpids):
      try:
        retry_util.RetryException(
            exception=_TooFewPidsException,
            max_retry=5,
            functor=lambda: _GetRunningPids(exe, numpids),
            sleep=sleep)
      except _TooFewPidsException:
        raise VMError('_WaitForProcs failed: timed out while waiting for '
                      '%d %s processes to start.' % (numpids, exe))

    # We could also wait for session_manager, nacl_helper, etc, but chrome is
    # the long pole. We expect the parent, 2 zygotes, gpu-process,
    # utility-process, 3 renderers.
    _WaitForProc('chrome', 8)

  def WaitForBoot(self, max_retry=3, sleep=5):
    """Wait for the VM to boot up.

    Wait for ssh connection to become active, and wait for all expected chrome
    processes to be launched. Set max_retry to a lower value since we can easily
    restart the VM if something is stuck and timing out.
    """
    if not os.path.exists(self.vm_dir):
      self.Start()

    super(VM, self).WaitForBoot(sleep=sleep, max_retry=max_retry)

    # XXX: this seems to be chromeos-secific, WaitForProcs waits for chrome
    # Chrome can take a while to start with software emulation.
    # if not self.enable_kvm:
    #   self._WaitForProcs()

  @staticmethod
  def GetParser():
    """Parse a list of args.

    Args:
      argv: list of command line arguments.

    Returns:
      List of parsed opts.
    """
    parser = device.Device.GetParser()
    parser.add_argument('--start', action='store_true', default=False,
                        help='Start the VM.')
    parser.add_argument('--stop', action='store_true', default=False,
                        help='Stop the VM.')
    parser.add_argument('--image-path', type='path',
                        help='Path to VM image to launch with --start.')
    parser.add_argument('--image-format', default=VM.IMAGE_FORMAT,
                        help='Format of the VM image (raw, qcow2, ...).')
    parser.add_argument('--qemu-path', type='path',
                        help='Path of qemu binary to launch with --start.')
    parser.add_argument('--qemu-m', type=str, default='8G',
                        help='Memory argument that will be passed to qemu.')
    parser.add_argument('--qemu-smp', type=int, default='0',
                        help='SMP argument that will be passed to qemu. (0 '
                             'means auto-detection.)')
    parser.add_argument('--qemu-cpu', type=str,
                        help='CPU argument that will be passed to qemu.')
    parser.add_argument('--qemu-arch', type=str, default=VM.ARCH_X86_64,
                        choices=(VM.ARCH_AARCH64, VM.ARCH_X86_64),
                        help='VM architecture: ')
    parser.add_argument('--qemu-bios-path', type='path',
                        help='Path of directory with qemu bios files.')
    parser.add_argument('--qemu-hostfwd', action='append',
                        help='Ports to forward from the VM to the host in the '
                        'QEMU hostfwd format, eg tcp:127.0.0.1:12345-:54321 to '
                        'forward port 54321 on the VM to 12345 on the host.')
    parser.add_argument('--qemu-args', action='append',
                        help='Additional args to pass to qemu.')
    parser.add_argument('--copy-on-write', action='store_true', default=False,
                        help='Generates a temporary copy-on-write image backed '
                             'by the normal boot image. All filesystem changes '
                             'will instead be reflected in the temporary '
                             'image.')
    parser.add_argument('--qemu-img-path', type='path',
                        help='Path to qemu-img binary used to create temporary '
                             'copy-on-write images.')
    parser.add_argument('--disable-kvm', dest='enable_kvm',
                        action='store_false', default=None,
                        help='Disable KVM, use software emulation.')
    parser.add_argument('--no-display', dest='display',
                        action='store_false', default=True,
                        help='Do not display video output.')
    parser.add_argument('--ssh-port', type=int,
                        help='ssh port to communicate with VM.')
    parser.add_argument('--chroot-path', type='path',
                        default=os.path.join(constants.SOURCE_ROOT,
                                             constants.DEFAULT_CHROOT_DIR))
    parser.add_argument('--cache-dir', type='path',
                        default=path_util.GetCacheDir(),
                        help='Cache directory to use.')
    parser.add_argument('--vm-dir', type='path',
                        help='Temp VM directory to use.')
    return parser