lib/moblab_vm.py - mirrors/cros/chromiumos/chromite - Git at Google

 # -*- coding: utf-8 -*-
 # Copyright 2017 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.

 """Module to manage a setup of moblab and its DUT VMs."""

 from __future__ import print_function

 import collections
 import contextlib
 import json
 import os
 import random
 import shutil
 import sys

 from chromite.lib import chroot_util
 from chromite.lib import constants
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
 from chromite.lib import osutils
 from chromite.lib import path_util
 from chromite.lib import pformat
 from chromite.lib import retry_util


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


 _CONFIG_FILE_NAME = 'moblabvm.json'

 # Keys into the config dict persisted to disk
 _CONFIG_INITIALIZED = 'initialized'
 _CONFIG_STARTED = 'started'
 _CONFIG_MOBLAB_IMAGE = 'image_path_moblab'
 _CONFIG_DUT_IMAGE = 'image_path_dut'
 _CONFIG_MOBLAB_DISK = 'disk_path_moblab'
 # Some paths are stored as relative paths in the config so that we can freely
 # move around the workspace on disk. They're always used as absolute paths once
 # loaded, because absolute paths are easier to work with.
 _CONFIG_RELATIVE_PATH_KEYS = (
     _CONFIG_MOBLAB_IMAGE,
     _CONFIG_DUT_IMAGE,
     _CONFIG_MOBLAB_DISK,
 )

 _CONFIG_MOBLAB_SSH_PORT = 'ssh_port_moblab'
 _CONFIG_DUT_SSH_PORT = 'ssh_port_dut'
 _CONFIG_NETWORK_BRIDGE = 'network_bridge'
 _CONFIG_MOBLAB_TAP_DEV = 'tap_moblab'
 _CONFIG_DUT_TAP_DEV = 'tap_dut'
 _CONFIG_MOBLAB_TAP_MAC = 'mac_tap_moblab'
 _CONFIG_DUT_TAP_MAC = 'mac_tap_dut'

 # Paths within the workspace
 _WORKSPACE_MOBLAB_DIR = 'moblab_image'
 _WORKSPACE_DUT_DIR = 'dut_image'


 class MoblabVmError(Exception):
   """Base exception for all ya folks who want to catch-all."""


 class RetriableError(MoblabVmError):
   """Some error happened where the best option is perhaps to retry."""


 class SetupError(MoblabVmError):
   """Some error happened while setting up the VMs."""


 class StartError(MoblabVmError):
   """Some error happened while setting up the VMs."""


 class MoblabVm(object):
   """Manage a setup of VMs required for a VM-only moblab setup.

   See cli/cros/cros_moblabvm.py for details about the setup.
   """

   def __init__(self, workspace_dir, chroot_dir=None):
     """Initialize us.

     Args:
       workspace_dir: An existing directory to be used for to drop files
           necessary for the setup.
       chroot_dir: Specific chroot to pass to all cros_vm commands.
     """
     self.workspace = workspace_dir
     self.chroot = chroot_dir
     self._config_path = os.path.join(self.workspace, _CONFIG_FILE_NAME)
     self._config = self._LoadConfig()

   @property
   def initialized(self):
     """Is this MoblabVm initialized? Returns bool."""
     return bool(self._config[_CONFIG_INITIALIZED])

   @property
   def running(self):
     """Is this MoblabVm running? Returns bool."""
     return (bool(self._config[_CONFIG_STARTED]) and
             bool(self._config[_CONFIG_MOBLAB_SSH_PORT]))

   @property
   def moblab_ssh_port(self):
     """The (str) SSH port to connect to the moblab VM."""
     return str(self._config[_CONFIG_MOBLAB_SSH_PORT])

   @property
   def moblab_internal_mac(self):
     """The MAC address used by moblab VM for the internal network.

     This is the network that moblab uses to connect to its DUTs.
     """
     return self._config[_CONFIG_MOBLAB_TAP_MAC]

   @property
   def dut_running(self):
     """Is the sub-DUT VM running? Returns bool.

     A moblabvm may be launched without any sub-DUTs.
     """
     return (bool(self._config[_CONFIG_STARTED]) and
             bool(self._config[_CONFIG_DUT_SSH_PORT]))

   @property
   def dut_ssh_port(self):
     """The (str) SSH port to connect to the sub-DUT VM."""
     return str(self._config[_CONFIG_DUT_SSH_PORT])

   @property
   def dut_internal_mac(self):
     """The MAC address used by sub-DUT VM for the internal network.

     This is the network that moblab uses to connect to its DUTs.
     """
     return self._config[_CONFIG_DUT_TAP_MAC]

   def Create(self, moblab_image_dir, dut_image_dir='', create_vm_images=True):
     """Create a new moblabvm setup.

     Args:
       moblab_image_dir: A directory containing the image for moblab. This
           directory should ideally contain all the files output by build_image
           because they are needed to manipulate the image properly.
       dut_image_dir: A directory containing the image for the sub-DUT. Only
           required if a DUT should be attached to the moblab VM. This directory
           should ideally contain all the files output by build_image because
           they are needed to manipulate the image properly.
       create_vm_images: If True, the source directories contain test images that
           need to be converted to vm images. If False, the source directories
           should already contain VM images to be used directly.
     """
     if self._config[_CONFIG_INITIALIZED]:
       raise SetupError('Cannot overwrite existing setup at %s' %
                        self.workspace)

     logging.notice('Initializing workspace in %s', self.workspace)
     logging.notice('This involves creating some VM images. '
                    'May take a few minutes.')

     logging.notice('Preparing moblab image...')
     moblab_dir = os.path.join(self.workspace, _WORKSPACE_MOBLAB_DIR)
     osutils.SafeMakedirsNonRoot(moblab_dir)
     if create_vm_images:
       self._config[_CONFIG_MOBLAB_IMAGE] = _CreateVMImage(
           moblab_image_dir, moblab_dir)
     else:
       self._config[_CONFIG_MOBLAB_IMAGE] = self._CopyVMImage(
           moblab_image_dir, moblab_dir)

     logging.notice('Generating moblab external disk...')
     moblab_disk = _CreateMoblabDisk(moblab_dir)
     self._config[_CONFIG_MOBLAB_DISK] = moblab_disk

     if dut_image_dir:
       logging.notice('Preparing dut image...')
       dut_dir = os.path.join(self.workspace, _WORKSPACE_DUT_DIR)
       osutils.SafeMakedirsNonRoot(dut_dir)
       if create_vm_images:
         self._config[_CONFIG_DUT_IMAGE] = _CreateVMImage(
             dut_image_dir, dut_dir)
       else:
         self._config[_CONFIG_DUT_IMAGE] = self._CopyVMImage(
             dut_image_dir, dut_dir)

     self._config[_CONFIG_INITIALIZED] = 'true'
     self._Persist()
     logging.notice('All Done!')

   def Start(self):
     """Launch the VM(s) in this setup.

     This involves grabbing some global resources on the host for networking
     between the launched VMs.
     """
     if not self._config[_CONFIG_INITIALIZED]:
       raise StartError('No moblabvm initialized at %s' % self.workspace)

     if self._config[_CONFIG_STARTED]:
       raise StartError('Can not start (partially) started VMs. '
                        'Pelase call Stop first.')

     # Start grabs system global resources. Try hard to persist even partial
     # attempts so that Stop can cleanup everything we do setup correctly.
     try:
       self._Start()
       self._config[_CONFIG_STARTED] = 'true'
     finally:
       self._Persist()

   def Stop(self):
     """Stop the VM(s) in this setup.

     Also releases global network resources.
     """
     # These should strictly be in LIFO order of setup done in Start.
     # Moreover, it should be able to cleanup after Start fails after partial
     # setup.
     if self._config[_CONFIG_DUT_SSH_PORT]:
       _StopVM(self._config[_CONFIG_DUT_SSH_PORT], chroot_path=self.chroot)
       del self._config[_CONFIG_DUT_SSH_PORT]
     if self._config[_CONFIG_MOBLAB_SSH_PORT]:
       _StopVM(self._config[_CONFIG_MOBLAB_SSH_PORT], chroot_path=self.chroot)
       del self._config[_CONFIG_MOBLAB_SSH_PORT]
     if self._config[_CONFIG_DUT_TAP_DEV]:
       _RemoveTapDeviceIgnoringErrors(self._config[_CONFIG_DUT_TAP_DEV])
       del self._config[_CONFIG_DUT_TAP_DEV]
     if self._config[_CONFIG_MOBLAB_TAP_DEV]:
       _RemoveTapDeviceIgnoringErrors(self._config[_CONFIG_MOBLAB_TAP_DEV])
       del self._config[_CONFIG_MOBLAB_TAP_DEV]
     if self._config[_CONFIG_NETWORK_BRIDGE]:
       _RemoveNetworkBridgeIgnoringErrors(self._config[_CONFIG_NETWORK_BRIDGE])
       del self._config[_CONFIG_NETWORK_BRIDGE]

     if self._config[_CONFIG_STARTED]:
       del self._config[_CONFIG_STARTED]
     self._Persist()

   def Destroy(self):
     """Completely destroy this moblabvm setup. Cleans out workspace_dir."""
     self.Stop()
     osutils.EmptyDir(self.workspace, ignore_missing=True, sudo=True)
     self._config = collections.defaultdict(str)
     # Don't self._Persist since the workspace is now clean.

   @contextlib.contextmanager
   def RunVmsContext(self):
     """A context manager to start the VM and guarantee it is stopped."""
     try:
       self.Start()
       yield
     finally:
       self.Destroy()

   @contextlib.contextmanager
   def MountedMoblabDiskContext(self):
     """A contextmanager to mount the already prepared moblab disk.

     This can be used to modify the external disk mounted by the moblab VM, while
     the VMs are not running. It is often much more performance to modify the
     disk on the host than SCP'ing large files into the VM.

     vms = moblab_vm.MoblabVm(workspace_dir)
     vms.Create(...)
     with vms.MountedMoblabDiskContext() as moblab_disk_dir:
       # Copy stuff into the directory at moblab_disk_dir
     vms.Start()
     ...
     """
     if not self.initialized:
       raise MoblabVmError('Uninitialized workspace %s. Can not mount disk.'
                           % self.workspace)
     if self.running:
       raise MoblabVmError(
           'VM at %s is already running. Stop() before mounting disk.'
           % self.workspace)

     with osutils.TempDir() as tempdir:
       osutils.MountDir(self._config[_CONFIG_MOBLAB_DISK], tempdir, 'ext4',
                        skip_mtab=True)
       try:
         yield tempdir
       finally:
         osutils.UmountDir(tempdir)

   def _LoadConfig(self):
     """Attempts to load the config from workspace, or sets up defaults."""
     config = collections.defaultdict(str)
     if os.path.isfile(self._config_path):
       with open(self._config_path, 'r') as config_file:
         config_data = json.load(config_file)
       self._ConfigRelativePathsToAbsolute(config_data,
                                           os.path.dirname(self._config_path))
       config.update(config_data)
     return config

   def _Persist(self):
     """Dump our config to disk for later MoblabVm objects."""
     # We do not want to modify config dict for the object.
     config = dict(self._config)
     self._ConfigAbsolutePathsToRelative(config,
                                         os.path.dirname(self._config_path))
     pformat.json(config, fp=self._config_path)

   def _ConfigRelativePathsToAbsolute(self, config_data, workspace_dir):
     """Converts all the relative paths loaded from config to absolute paths."""
     for key, relpath in config_data.items():
       if key not in _CONFIG_RELATIVE_PATH_KEYS:
         continue
       config_data[key] = os.path.realpath(os.path.join(workspace_dir, relpath))

   def _ConfigAbsolutePathsToRelative(self, config_data, workspace_dir):
     """Converts all absolute paths to relative paths for persisting on disk."""
     workspace_dir = os.path.realpath(workspace_dir)
     for key, abspath in config_data.items():
       if key not in _CONFIG_RELATIVE_PATH_KEYS:
         continue
       abspath = os.path.realpath(abspath)
       config_data[key] = os.path.relpath(abspath, workspace_dir)

   def _Start(self):
     """Actually start the VM(s), without persisting the updated config."""
     # Mulitple moblabvms should be able to run in parallel. At the same time, we
     # want to minimize flaky failures due to unavailable system global
     # resources. We attempt to create a bridge device a few times. Once we get a
     # hold of that, we base all other system global constants off of that to not
     # step our other moblabvms' toes.
     for _ in range(5):
       try:
         port, bridge_name = _TryCreateBridgeDevice()
         break
       except RetriableError:
         pass
     else:
       raise SetupError('Failed to create bridge device. '
                        'You seem to have many moblabvms running in parallel.')

     self._config[_CONFIG_NETWORK_BRIDGE] = bridge_name
     self._config[_CONFIG_MOBLAB_TAP_DEV] = _CreateTapDevice(port)
     self._config[_CONFIG_DUT_TAP_DEV] = _CreateTapDevice(port + 1)
     _ConnectDeviceToBridge(self._config[_CONFIG_MOBLAB_TAP_DEV],
                            self._config[_CONFIG_NETWORK_BRIDGE])
     _ConnectDeviceToBridge(self._config[_CONFIG_DUT_TAP_DEV],
                            self._config[_CONFIG_NETWORK_BRIDGE])
     _DeviceUp(self._config[_CONFIG_NETWORK_BRIDGE])

     moblab_ssh_port = port + 2
     # moblab grabs 3 extra consecutive ports for forwarding AFE and devserver to
     # the host.
     dut_ssh_port = moblab_ssh_port + 4
     self._StartMoblabVM(moblab_ssh_port, dut_ssh_port)
     if self._config[_CONFIG_DUT_IMAGE]:
       self._StartDutVM(dut_ssh_port)

   def _CopyVMImage(self, source_dir, target_dir):
     """Converts or copies VM images from source_dir to target_dir."""
     source_path = os.path.join(source_dir, constants.VM_IMAGE_BIN)
     if not os.path.isfile(source_path):
       raise SetupError('Could not find VM image at %s' % source_path)

     target_path = os.path.join(target_dir, constants.VM_IMAGE_BIN)
     shutil.copyfile(source_path, target_path)
     return target_path

   def _StartMoblabVM(self, ssh_port, max_port):
     """Starts a VM running moblab.

     Args:
       ssh_port: (int) VM SSH port to use.
       max_port: (int) ports upto this are available.
     """
     # We want a unicast, locally configured address with an obviously bogus
     # organisation id.
     tap_mac_addr = '02:00:00:99:99:01'
     self._WriteToMoblabDiskImage('private-network-macaddr.conf', [tap_mac_addr])
     _StartVM(
         self._config[_CONFIG_MOBLAB_IMAGE],
         ssh_port,
         max_port,
         self._config[_CONFIG_MOBLAB_TAP_DEV],
         tap_mac_addr,
         is_moblab=True,
         disk_path=self._config[_CONFIG_MOBLAB_DISK],
         chroot_path=self.chroot,
     )
     # Update config _after_ we've successfully launched the VM so that we don't
     # _Persist incorrect information.
     self._config[_CONFIG_MOBLAB_SSH_PORT] = ssh_port
     self._config[_CONFIG_MOBLAB_TAP_MAC] = tap_mac_addr

   def _StartDutVM(self, ssh_port):
     """Starts a VM running the sub-DUT image.

     Args:
       ssh_port: VM SSH port to use.
     """
     # We want a unicast, locally configured address with an obviously bogus
     # organisation id.
     tap_mac_addr = '02:00:00:99:99:51'
     _StartVM(
         self._config[_CONFIG_DUT_IMAGE],
         ssh_port,
         ssh_port + 1,
         self._config[_CONFIG_DUT_TAP_DEV],
         tap_mac_addr,
         is_moblab=False,
         chroot_path=self.chroot,
     )
     # Update config _after_ we've successfully launched the VM so that we don't
     # _Persist incorrect information.
     self._config[_CONFIG_DUT_SSH_PORT] = ssh_port
     self._config[_CONFIG_DUT_TAP_MAC] = tap_mac_addr

   def _WriteToMoblabDiskImage(self, target_path, lines):
     """Write a file in the moblab external disk.

     Args:
       target_path: Path within the mounted disk to write. This path will be
           overwritten.
       lines: An iterator of lines to write.
     """
     with self.MountedMoblabDiskContext() as disk_dir:
       osutils.WriteFile(os.path.join(disk_dir, target_path), lines, sudo=True)


 def _CreateVMImage(src_dir, dest_dir):
   """Creates a VM image from a given chromiumos image.

   Args:
     src_dir: Path to the directory containing (non-VM) image. Defaults to None
         to use the latest image for the board.
     dest_dir: Path to the directory where the VM image should be written.

   Returns:
     The path of the created VM image.
   """
   # image_to_vm.sh only runs in chroot, but src_dir / dest_dir may not be
   # reachable from chroot. Also, image_to_vm.sh needs all the contents of
   # src_dir to work correctly (it silently does the wrong thing if some files
   # are missing).
   # So, create a tempdir reachable from chroot, copy everything to that path,
   # create vm image there and finally move it all to dest_dir.
   with chroot_util.TempDirInChroot() as tempdir:
     logging.debug('Copying images from %s to %s.', src_dir, tempdir)
     osutils.CopyDirContents(src_dir, tempdir)
     # image_to_vm.sh doesn't let us provide arbitrary names for the input image.
     # Instead, it picks the name based on whether we pass in --test_image or not
     # (and doesn't use that flag for anything else).
     cmd = [
         path_util.ToChrootPath(
             os.path.join(constants.CROSUTILS_DIR, 'image_to_vm.sh')),
         '--from=%s' % path_util.ToChrootPath(tempdir),
         '--disk_layout=16gb-rootfs',
         '--test_image',
     ]
     try:
       cros_build_lib.run(cmd, enter_chroot=True, cwd=constants.SOURCE_ROOT)
     except cros_build_lib.RunCommandError as e:
       raise SetupError('Failed to create VM image for %s: %s' % (src_dir, e))

     # Preserve most content, although we should need only the generated VM
     # image. Other files like boot.desc might be needed elsewhere, but the
     # source images should no longer be needed.
     osutils.SafeUnlink(os.path.join(tempdir, constants.BASE_IMAGE_BIN),
                        sudo=True)
     osutils.SafeUnlink(os.path.join(tempdir, constants.TEST_IMAGE_BIN),
                        sudo=True)
     osutils.CopyDirContents(tempdir, dest_dir)
   # The exact name of the output image is hard-coded in image_to_vm.sh
   return os.path.join(dest_dir, constants.VM_IMAGE_BIN)


 def _CreateMoblabDisk(dest_dir):
   """Creates an empty moblab virtual disk at the given path.

   Args:
     dest_dir: Directory where the disk image should be created.

   Returns:
     Path to the created disk image.
   """
   dest_path = os.path.join(dest_dir, 'moblab_disk')
   cros_build_lib.run(['qemu-img', 'create', '-f', 'raw', dest_path, '64g'])
   cros_build_lib.run(['mkfs.ext4', '-F', dest_path])
   cros_build_lib.run(['e2label', dest_path, 'MOBLAB-STORAGE'])
   return dest_path


 def _TryCreateBridgeDevice():
   """Creates a bridge device named moblabvmbrXX, generating XX randomly.

   Returns:
     num, name: Counter to be used to generate unique names, name of the bridge.
   """
   # Keep this number between 1024-49151 (with padding) because:
   # - We use it to reserve networking ports, so stay within user range.
   # - We use it to name kernel devices, which has length restrictions.
   num = random.randint(10000, 30000)
   # device names can only be 16 char long. So suffix can be between 0 - 10^9.
   name = 'brmob%s' % num
   cros_build_lib.sudo_run(['ip', 'link', 'add', name, 'type', 'bridge'],
                           quiet=True)
   return num + 1, name


 def _CreateTapDevice(suffix):
   """Create a tap device using the suffix provided."""
   # device names can only be 16 char long. So suffix can be between 0 - 10^10.
   name = 'tapmob%s' % suffix
   cros_build_lib.sudo_run(['ip', 'tuntap', 'add', 'mode', 'tap', name])

   # Creating the device can take some time, and trying to create another device
   # in the meantime returns 'Device or resource busy'. So, wait for device
   # creation to complete.
   @retry_util.WithRetry(max_retry=3, sleep=0.2, exception=RetriableError)
   def _VerifyDeviceCreated(name):
     result = cros_build_lib.run(['ip', 'tuntap', 'show'], stdout=True)
     if name not in result.output:
       raise RetriableError('Device %s not found in output "%s".' %
                            (name, result.output))

   _VerifyDeviceCreated(name)
   return name


 def _ConnectDeviceToBridge(device, bridge):
   """Connects link device to network bridge."""
   cros_build_lib.sudo_run(['ip', 'link', 'set', device, 'master', bridge])


 def _DeviceUp(device):
   """Bring up the given link device."""
   cros_build_lib.sudo_run(['ip', 'link', 'set', 'dev', device, 'up'])


 def _StartVM(image_path, ssh_port, max_port, tap_dev, tap_mac_addr, is_moblab,
              disk_path=None, chroot_path=None):
   """Starts a VM instance.

   Args:
     image_path: Path to the OS image to use.
     ssh_port: (int) port to use for SSH.
     max_port: (int) ports upto this are available.
     tap_dev: Name of the tap device to use for secondary network.
     tap_mac_addr: The MAC address of the secondary network device.
     is_moblab: Whether we're starting a moblab VM.
     disk_path: Moblab disk.
     chroot_path: Location of chroot on disk to pass to cros_vm.
   """
   cmd = [
       './cros_vm', '--start', '--image-path=%s' % image_path,
       '--ssh-port=%d' % ssh_port,
       '--qemu-m', '32G',
       '--qemu-args', '-net nic,macaddr=%s' % tap_mac_addr,
       '--qemu-args', '-net tap,ifname=%s' % tap_dev
   ]
   if chroot_path:
     cmd.extend(['--chroot-path', chroot_path])
   if is_moblab:
     moblab_monitoring_port = ssh_port + 1
     afe_port = ssh_port + 2
     dev_server_port = ssh_port + 3
     assert dev_server_port < max_port, 'Exceeded maximum available ports.'
     cmd += [
         # Moblab monitoring port.
         '--qemu-hostfwd', 'tcp:127.0.0.1:%d-:9991' % moblab_monitoring_port,
         # AFE port.
         '--qemu-hostfwd', 'tcp:127.0.0.1:%d-:80' % afe_port,
         # DevServer port.
         '--qemu-hostfwd', 'tcp:127.0.0.1:%d-:8080' % dev_server_port,
         # Create a dedicated scsi controller for the external disk, and attach
         # the disk to that bus. This separates us from any scsi controllers that
         # may be created for the boot disk.
         '--qemu-args',
         '-drive id=moblabdisk,if=none,file=%s,format=raw' % disk_path,
         '--qemu-args', '-device virtio-scsi-pci,id=scsiext',
         '--qemu-args', '-device scsi-hd,bus=scsiext.0,drive=moblabdisk',
     ]
   cros_build_lib.sudo_run(cmd, cwd=constants.CHROMITE_BIN_DIR)


 def _RemoveNetworkBridgeIgnoringErrors(name):
   """Removes a previously created network bridge. Ignores errors."""
   _RunIgnoringErrors(['ip', 'link', 'del', name, 'type', 'bridge'])


 def _RemoveTapDeviceIgnoringErrors(name):
   """Removes a previously created tap device. Ignores errors."""
   _RunIgnoringErrors(['ip', 'tuntap', 'del', 'mode', 'tap', name])


 def _StopVM(ssh_port, chroot_path=None):
   """Stops a running VM instance.

   Args:
     ssh_port: (int) port to use for ssh.
     chroot_path: Location of chroot on  disk to pass to cros_vm.
   """
   logging.notice('Stopping the VM. This may take a minute.')
   cmd = [
       '%s/cros_vm' % constants.CHROMITE_BIN_DIR,
       '--stop', '--ssh-port=%d' % ssh_port,
   ]
   if chroot_path:
     cmd.extend(['--chroot-path', chroot_path])
   _RunIgnoringErrors(cmd)


 def _RunIgnoringErrors(cmd):
   """Runs the given command, ignores errors but warns user."""
   try:
     cros_build_lib.sudo_run(cmd, quiet=True)
   except cros_build_lib.RunCommandError as e:
     logging.error('Cleanup operation failed. Please run "%s" manually.',
                   ' '.join(cmd))
     logging.debug('Encountered error: %s', e)
	# -- coding: utf-8 --
	# Copyright 2017 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.

	"""Module to manage a setup of moblab and its DUT VMs."""

	from __future__ import print_function

	import collections
	import contextlib
	import json
	import os
	import random
	import shutil
	import sys

	from chromite.lib import chroot_util
	from chromite.lib import constants
	from chromite.lib import cros_build_lib
	from chromite.lib import cros_logging as logging
	from chromite.lib import osutils
	from chromite.lib import path_util
	from chromite.lib import pformat
	from chromite.lib import retry_util


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


	_CONFIG_FILE_NAME = 'moblabvm.json'

	# Keys into the config dict persisted to disk
	_CONFIG_INITIALIZED = 'initialized'
	_CONFIG_STARTED = 'started'
	_CONFIG_MOBLAB_IMAGE = 'image_path_moblab'
	_CONFIG_DUT_IMAGE = 'image_path_dut'
	_CONFIG_MOBLAB_DISK = 'disk_path_moblab'
	# Some paths are stored as relative paths in the config so that we can freely
	# move around the workspace on disk. They're always used as absolute paths once
	# loaded, because absolute paths are easier to work with.
	_CONFIG_RELATIVE_PATH_KEYS = (
	_CONFIG_MOBLAB_IMAGE,
	_CONFIG_DUT_IMAGE,
	_CONFIG_MOBLAB_DISK,
	)

	_CONFIG_MOBLAB_SSH_PORT = 'ssh_port_moblab'
	_CONFIG_DUT_SSH_PORT = 'ssh_port_dut'
	_CONFIG_NETWORK_BRIDGE = 'network_bridge'
	_CONFIG_MOBLAB_TAP_DEV = 'tap_moblab'
	_CONFIG_DUT_TAP_DEV = 'tap_dut'
	_CONFIG_MOBLAB_TAP_MAC = 'mac_tap_moblab'
	_CONFIG_DUT_TAP_MAC = 'mac_tap_dut'

	# Paths within the workspace
	_WORKSPACE_MOBLAB_DIR = 'moblab_image'
	_WORKSPACE_DUT_DIR = 'dut_image'


	class MoblabVmError(Exception):
	"""Base exception for all ya folks who want to catch-all."""


	class RetriableError(MoblabVmError):
	"""Some error happened where the best option is perhaps to retry."""


	class SetupError(MoblabVmError):
	"""Some error happened while setting up the VMs."""


	class StartError(MoblabVmError):
	"""Some error happened while setting up the VMs."""


	class MoblabVm(object):
	"""Manage a setup of VMs required for a VM-only moblab setup.

	See cli/cros/cros_moblabvm.py for details about the setup.
	"""

	def __init__(self, workspace_dir, chroot_dir=None):
	"""Initialize us.

	Args:
	workspace_dir: An existing directory to be used for to drop files
	necessary for the setup.
	chroot_dir: Specific chroot to pass to all cros_vm commands.
	"""
	self.workspace = workspace_dir
	self.chroot = chroot_dir
	self._config_path = os.path.join(self.workspace, _CONFIG_FILE_NAME)
	self._config = self._LoadConfig()

	@property
	def initialized(self):
	"""Is this MoblabVm initialized? Returns bool."""
	return bool(self._config[_CONFIG_INITIALIZED])

	@property
	def running(self):
	"""Is this MoblabVm running? Returns bool."""
	return (bool(self._config[_CONFIG_STARTED]) and
	bool(self._config[_CONFIG_MOBLAB_SSH_PORT]))

	@property
	def moblab_ssh_port(self):
	"""The (str) SSH port to connect to the moblab VM."""
	return str(self._config[_CONFIG_MOBLAB_SSH_PORT])

	@property
	def moblab_internal_mac(self):
	"""The MAC address used by moblab VM for the internal network.

	This is the network that moblab uses to connect to its DUTs.
	"""
	return self._config[_CONFIG_MOBLAB_TAP_MAC]

	@property
	def dut_running(self):
	"""Is the sub-DUT VM running? Returns bool.

	A moblabvm may be launched without any sub-DUTs.
	"""
	return (bool(self._config[_CONFIG_STARTED]) and
	bool(self._config[_CONFIG_DUT_SSH_PORT]))

	@property
	def dut_ssh_port(self):
	"""The (str) SSH port to connect to the sub-DUT VM."""
	return str(self._config[_CONFIG_DUT_SSH_PORT])

	@property
	def dut_internal_mac(self):
	"""The MAC address used by sub-DUT VM for the internal network.

	This is the network that moblab uses to connect to its DUTs.
	"""
	return self._config[_CONFIG_DUT_TAP_MAC]

	def Create(self, moblab_image_dir, dut_image_dir='', create_vm_images=True):
	"""Create a new moblabvm setup.

	Args:
	moblab_image_dir: A directory containing the image for moblab. This
	directory should ideally contain all the files output by build_image
	because they are needed to manipulate the image properly.
	dut_image_dir: A directory containing the image for the sub-DUT. Only
	required if a DUT should be attached to the moblab VM. This directory
	should ideally contain all the files output by build_image because
	they are needed to manipulate the image properly.
	create_vm_images: If True, the source directories contain test images that
	need to be converted to vm images. If False, the source directories
	should already contain VM images to be used directly.
	"""
	if self._config[_CONFIG_INITIALIZED]:
	raise SetupError('Cannot overwrite existing setup at %s' %
	self.workspace)

	logging.notice('Initializing workspace in %s', self.workspace)
	logging.notice('This involves creating some VM images. '
	'May take a few minutes.')

	logging.notice('Preparing moblab image...')
	moblab_dir = os.path.join(self.workspace, _WORKSPACE_MOBLAB_DIR)
	osutils.SafeMakedirsNonRoot(moblab_dir)
	if create_vm_images:
	self._config[_CONFIG_MOBLAB_IMAGE] = _CreateVMImage(
	moblab_image_dir, moblab_dir)
	else:
	self._config[_CONFIG_MOBLAB_IMAGE] = self._CopyVMImage(
	moblab_image_dir, moblab_dir)

	logging.notice('Generating moblab external disk...')
	moblab_disk = _CreateMoblabDisk(moblab_dir)
	self._config[_CONFIG_MOBLAB_DISK] = moblab_disk

	if dut_image_dir:
	logging.notice('Preparing dut image...')
	dut_dir = os.path.join(self.workspace, _WORKSPACE_DUT_DIR)
	osutils.SafeMakedirsNonRoot(dut_dir)
	if create_vm_images:
	self._config[_CONFIG_DUT_IMAGE] = _CreateVMImage(
	dut_image_dir, dut_dir)
	else:
	self._config[_CONFIG_DUT_IMAGE] = self._CopyVMImage(
	dut_image_dir, dut_dir)

	self._config[_CONFIG_INITIALIZED] = 'true'
	self._Persist()
	logging.notice('All Done!')

	def Start(self):
	"""Launch the VM(s) in this setup.

	This involves grabbing some global resources on the host for networking
	between the launched VMs.
	"""
	if not self._config[_CONFIG_INITIALIZED]:
	raise StartError('No moblabvm initialized at %s' % self.workspace)

	if self._config[_CONFIG_STARTED]:
	raise StartError('Can not start (partially) started VMs. '
	'Pelase call Stop first.')

	# Start grabs system global resources. Try hard to persist even partial
	# attempts so that Stop can cleanup everything we do setup correctly.
	try:
	self._Start()
	self._config[_CONFIG_STARTED] = 'true'
	finally:
	self._Persist()

	def Stop(self):
	"""Stop the VM(s) in this setup.

	Also releases global network resources.
	"""
	# These should strictly be in LIFO order of setup done in Start.
	# Moreover, it should be able to cleanup after Start fails after partial
	# setup.
	if self._config[_CONFIG_DUT_SSH_PORT]:
	_StopVM(self._config[_CONFIG_DUT_SSH_PORT], chroot_path=self.chroot)
	del self._config[_CONFIG_DUT_SSH_PORT]
	if self._config[_CONFIG_MOBLAB_SSH_PORT]:
	_StopVM(self._config[_CONFIG_MOBLAB_SSH_PORT], chroot_path=self.chroot)
	del self._config[_CONFIG_MOBLAB_SSH_PORT]
	if self._config[_CONFIG_DUT_TAP_DEV]:
	_RemoveTapDeviceIgnoringErrors(self._config[_CONFIG_DUT_TAP_DEV])
	del self._config[_CONFIG_DUT_TAP_DEV]
	if self._config[_CONFIG_MOBLAB_TAP_DEV]:
	_RemoveTapDeviceIgnoringErrors(self._config[_CONFIG_MOBLAB_TAP_DEV])
	del self._config[_CONFIG_MOBLAB_TAP_DEV]
	if self._config[_CONFIG_NETWORK_BRIDGE]:
	_RemoveNetworkBridgeIgnoringErrors(self._config[_CONFIG_NETWORK_BRIDGE])
	del self._config[_CONFIG_NETWORK_BRIDGE]

	if self._config[_CONFIG_STARTED]:
	del self._config[_CONFIG_STARTED]
	self._Persist()

	def Destroy(self):
	"""Completely destroy this moblabvm setup. Cleans out workspace_dir."""
	self.Stop()
	osutils.EmptyDir(self.workspace, ignore_missing=True, sudo=True)
	self._config = collections.defaultdict(str)
	# Don't self._Persist since the workspace is now clean.

	@contextlib.contextmanager
	def RunVmsContext(self):
	"""A context manager to start the VM and guarantee it is stopped."""
	try:
	self.Start()
	yield
	finally:
	self.Destroy()

	@contextlib.contextmanager
	def MountedMoblabDiskContext(self):
	"""A contextmanager to mount the already prepared moblab disk.

	This can be used to modify the external disk mounted by the moblab VM, while
	the VMs are not running. It is often much more performance to modify the
	disk on the host than SCP'ing large files into the VM.

	vms = moblab_vm.MoblabVm(workspace_dir)
	vms.Create(...)
	with vms.MountedMoblabDiskContext() as moblab_disk_dir:
	# Copy stuff into the directory at moblab_disk_dir
	vms.Start()
	...
	"""
	if not self.initialized:
	raise MoblabVmError('Uninitialized workspace %s. Can not mount disk.'
	% self.workspace)
	if self.running:
	raise MoblabVmError(
	'VM at %s is already running. Stop() before mounting disk.'
	% self.workspace)

	with osutils.TempDir() as tempdir:
	osutils.MountDir(self._config[_CONFIG_MOBLAB_DISK], tempdir, 'ext4',
	skip_mtab=True)
	try:
	yield tempdir
	finally:
	osutils.UmountDir(tempdir)

	def _LoadConfig(self):
	"""Attempts to load the config from workspace, or sets up defaults."""
	config = collections.defaultdict(str)
	if os.path.isfile(self._config_path):
	with open(self._config_path, 'r') as config_file:
	config_data = json.load(config_file)
	self._ConfigRelativePathsToAbsolute(config_data,
	os.path.dirname(self._config_path))
	config.update(config_data)
	return config

	def _Persist(self):
	"""Dump our config to disk for later MoblabVm objects."""
	# We do not want to modify config dict for the object.
	config = dict(self._config)
	self._ConfigAbsolutePathsToRelative(config,
	os.path.dirname(self._config_path))
	pformat.json(config, fp=self._config_path)

	def _ConfigRelativePathsToAbsolute(self, config_data, workspace_dir):
	"""Converts all the relative paths loaded from config to absolute paths."""
	for key, relpath in config_data.items():
	if key not in _CONFIG_RELATIVE_PATH_KEYS:
	continue
	config_data[key] = os.path.realpath(os.path.join(workspace_dir, relpath))

	def _ConfigAbsolutePathsToRelative(self, config_data, workspace_dir):
	"""Converts all absolute paths to relative paths for persisting on disk."""
	workspace_dir = os.path.realpath(workspace_dir)
	for key, abspath in config_data.items():
	if key not in _CONFIG_RELATIVE_PATH_KEYS:
	continue
	abspath = os.path.realpath(abspath)
	config_data[key] = os.path.relpath(abspath, workspace_dir)

	def _Start(self):
	"""Actually start the VM(s), without persisting the updated config."""
	# Mulitple moblabvms should be able to run in parallel. At the same time, we
	# want to minimize flaky failures due to unavailable system global
	# resources. We attempt to create a bridge device a few times. Once we get a
	# hold of that, we base all other system global constants off of that to not
	# step our other moblabvms' toes.
	for _ in range(5):
	try:
	port, bridge_name = _TryCreateBridgeDevice()
	break
	except RetriableError:
	pass
	else:
	raise SetupError('Failed to create bridge device. '
	'You seem to have many moblabvms running in parallel.')

	self._config[_CONFIG_NETWORK_BRIDGE] = bridge_name
	self._config[_CONFIG_MOBLAB_TAP_DEV] = _CreateTapDevice(port)
	self._config[_CONFIG_DUT_TAP_DEV] = _CreateTapDevice(port + 1)
	_ConnectDeviceToBridge(self._config[_CONFIG_MOBLAB_TAP_DEV],
	self._config[_CONFIG_NETWORK_BRIDGE])
	_ConnectDeviceToBridge(self._config[_CONFIG_DUT_TAP_DEV],
	self._config[_CONFIG_NETWORK_BRIDGE])
	_DeviceUp(self._config[_CONFIG_NETWORK_BRIDGE])

	moblab_ssh_port = port + 2
	# moblab grabs 3 extra consecutive ports for forwarding AFE and devserver to
	# the host.
	dut_ssh_port = moblab_ssh_port + 4
	self._StartMoblabVM(moblab_ssh_port, dut_ssh_port)
	if self._config[_CONFIG_DUT_IMAGE]:
	self._StartDutVM(dut_ssh_port)

	def _CopyVMImage(self, source_dir, target_dir):
	"""Converts or copies VM images from source_dir to target_dir."""
	source_path = os.path.join(source_dir, constants.VM_IMAGE_BIN)
	if not os.path.isfile(source_path):
	raise SetupError('Could not find VM image at %s' % source_path)

	target_path = os.path.join(target_dir, constants.VM_IMAGE_BIN)
	shutil.copyfile(source_path, target_path)
	return target_path

	def _StartMoblabVM(self, ssh_port, max_port):
	"""Starts a VM running moblab.

	Args:
	ssh_port: (int) VM SSH port to use.
	max_port: (int) ports upto this are available.
	"""
	# We want a unicast, locally configured address with an obviously bogus
	# organisation id.
	tap_mac_addr = '02:00:00:99:99:01'
	self._WriteToMoblabDiskImage('private-network-macaddr.conf', [tap_mac_addr])
	_StartVM(
	self._config[_CONFIG_MOBLAB_IMAGE],
	ssh_port,
	max_port,
	self._config[_CONFIG_MOBLAB_TAP_DEV],
	tap_mac_addr,
	is_moblab=True,
	disk_path=self._config[_CONFIG_MOBLAB_DISK],
	chroot_path=self.chroot,
	)
	# Update config _after_ we've successfully launched the VM so that we don't
	# _Persist incorrect information.
	self._config[_CONFIG_MOBLAB_SSH_PORT] = ssh_port
	self._config[_CONFIG_MOBLAB_TAP_MAC] = tap_mac_addr

	def _StartDutVM(self, ssh_port):
	"""Starts a VM running the sub-DUT image.

	Args:
	ssh_port: VM SSH port to use.
	"""
	# We want a unicast, locally configured address with an obviously bogus
	# organisation id.
	tap_mac_addr = '02:00:00:99:99:51'
	_StartVM(
	self._config[_CONFIG_DUT_IMAGE],
	ssh_port,
	ssh_port + 1,
	self._config[_CONFIG_DUT_TAP_DEV],
	tap_mac_addr,
	is_moblab=False,
	chroot_path=self.chroot,
	)
	# Update config _after_ we've successfully launched the VM so that we don't
	# _Persist incorrect information.
	self._config[_CONFIG_DUT_SSH_PORT] = ssh_port
	self._config[_CONFIG_DUT_TAP_MAC] = tap_mac_addr

	def _WriteToMoblabDiskImage(self, target_path, lines):
	"""Write a file in the moblab external disk.

	Args:
	target_path: Path within the mounted disk to write. This path will be
	overwritten.
	lines: An iterator of lines to write.
	"""
	with self.MountedMoblabDiskContext() as disk_dir:
	osutils.WriteFile(os.path.join(disk_dir, target_path), lines, sudo=True)


	def _CreateVMImage(src_dir, dest_dir):
	"""Creates a VM image from a given chromiumos image.

	Args:
	src_dir: Path to the directory containing (non-VM) image. Defaults to None
	to use the latest image for the board.
	dest_dir: Path to the directory where the VM image should be written.

	Returns:
	The path of the created VM image.
	"""
	# image_to_vm.sh only runs in chroot, but src_dir / dest_dir may not be
	# reachable from chroot. Also, image_to_vm.sh needs all the contents of
	# src_dir to work correctly (it silently does the wrong thing if some files
	# are missing).
	# So, create a tempdir reachable from chroot, copy everything to that path,
	# create vm image there and finally move it all to dest_dir.
	with chroot_util.TempDirInChroot() as tempdir:
	logging.debug('Copying images from %s to %s.', src_dir, tempdir)
	osutils.CopyDirContents(src_dir, tempdir)
	# image_to_vm.sh doesn't let us provide arbitrary names for the input image.
	# Instead, it picks the name based on whether we pass in --test_image or not
	# (and doesn't use that flag for anything else).
	cmd = [
	path_util.ToChrootPath(
	os.path.join(constants.CROSUTILS_DIR, 'image_to_vm.sh')),
	'--from=%s' % path_util.ToChrootPath(tempdir),
	'--disk_layout=16gb-rootfs',
	'--test_image',
	]
	try:
	cros_build_lib.run(cmd, enter_chroot=True, cwd=constants.SOURCE_ROOT)
	except cros_build_lib.RunCommandError as e:
	raise SetupError('Failed to create VM image for %s: %s' % (src_dir, e))

	# Preserve most content, although we should need only the generated VM
	# image. Other files like boot.desc might be needed elsewhere, but the
	# source images should no longer be needed.
	osutils.SafeUnlink(os.path.join(tempdir, constants.BASE_IMAGE_BIN),
	sudo=True)
	osutils.SafeUnlink(os.path.join(tempdir, constants.TEST_IMAGE_BIN),
	sudo=True)
	osutils.CopyDirContents(tempdir, dest_dir)
	# The exact name of the output image is hard-coded in image_to_vm.sh
	return os.path.join(dest_dir, constants.VM_IMAGE_BIN)


	def _CreateMoblabDisk(dest_dir):
	"""Creates an empty moblab virtual disk at the given path.

	Args:
	dest_dir: Directory where the disk image should be created.

	Returns:
	Path to the created disk image.
	"""
	dest_path = os.path.join(dest_dir, 'moblab_disk')
	cros_build_lib.run(['qemu-img', 'create', '-f', 'raw', dest_path, '64g'])
	cros_build_lib.run(['mkfs.ext4', '-F', dest_path])
	cros_build_lib.run(['e2label', dest_path, 'MOBLAB-STORAGE'])
	return dest_path


	def _TryCreateBridgeDevice():
	"""Creates a bridge device named moblabvmbrXX, generating XX randomly.

	Returns:
	num, name: Counter to be used to generate unique names, name of the bridge.
	"""
	# Keep this number between 1024-49151 (with padding) because:
	# - We use it to reserve networking ports, so stay within user range.
	# - We use it to name kernel devices, which has length restrictions.
	num = random.randint(10000, 30000)
	# device names can only be 16 char long. So suffix can be between 0 - 10^9.
	name = 'brmob%s' % num
	cros_build_lib.sudo_run(['ip', 'link', 'add', name, 'type', 'bridge'],
	quiet=True)
	return num + 1, name


	def _CreateTapDevice(suffix):
	"""Create a tap device using the suffix provided."""
	# device names can only be 16 char long. So suffix can be between 0 - 10^10.
	name = 'tapmob%s' % suffix
	cros_build_lib.sudo_run(['ip', 'tuntap', 'add', 'mode', 'tap', name])

	# Creating the device can take some time, and trying to create another device
	# in the meantime returns 'Device or resource busy'. So, wait for device
	# creation to complete.
	@retry_util.WithRetry(max_retry=3, sleep=0.2, exception=RetriableError)
	def _VerifyDeviceCreated(name):
	result = cros_build_lib.run(['ip', 'tuntap', 'show'], stdout=True)
	if name not in result.output:
	raise RetriableError('Device %s not found in output "%s".' %
	(name, result.output))

	_VerifyDeviceCreated(name)
	return name


	def _ConnectDeviceToBridge(device, bridge):
	"""Connects link device to network bridge."""
	cros_build_lib.sudo_run(['ip', 'link', 'set', device, 'master', bridge])


	def _DeviceUp(device):
	"""Bring up the given link device."""
	cros_build_lib.sudo_run(['ip', 'link', 'set', 'dev', device, 'up'])


	def _StartVM(image_path, ssh_port, max_port, tap_dev, tap_mac_addr, is_moblab,
	disk_path=None, chroot_path=None):
	"""Starts a VM instance.

	Args:
	image_path: Path to the OS image to use.
	ssh_port: (int) port to use for SSH.
	max_port: (int) ports upto this are available.
	tap_dev: Name of the tap device to use for secondary network.
	tap_mac_addr: The MAC address of the secondary network device.
	is_moblab: Whether we're starting a moblab VM.
	disk_path: Moblab disk.
	chroot_path: Location of chroot on disk to pass to cros_vm.
	"""
	cmd = [
	'./cros_vm', '--start', '--image-path=%s' % image_path,
	'--ssh-port=%d' % ssh_port,
	'--qemu-m', '32G',
	'--qemu-args', '-net nic,macaddr=%s' % tap_mac_addr,
	'--qemu-args', '-net tap,ifname=%s' % tap_dev
	]
	if chroot_path:
	cmd.extend(['--chroot-path', chroot_path])
	if is_moblab:
	moblab_monitoring_port = ssh_port + 1
	afe_port = ssh_port + 2
	dev_server_port = ssh_port + 3
	assert dev_server_port < max_port, 'Exceeded maximum available ports.'
	cmd += [
	# Moblab monitoring port.
	'--qemu-hostfwd', 'tcp:127.0.0.1:%d-:9991' % moblab_monitoring_port,
	# AFE port.
	'--qemu-hostfwd', 'tcp:127.0.0.1:%d-:80' % afe_port,
	# DevServer port.
	'--qemu-hostfwd', 'tcp:127.0.0.1:%d-:8080' % dev_server_port,
	# Create a dedicated scsi controller for the external disk, and attach
	# the disk to that bus. This separates us from any scsi controllers that
	# may be created for the boot disk.
	'--qemu-args',
	'-drive id=moblabdisk,if=none,file=%s,format=raw' % disk_path,
	'--qemu-args', '-device virtio-scsi-pci,id=scsiext',
	'--qemu-args', '-device scsi-hd,bus=scsiext.0,drive=moblabdisk',
	]
	cros_build_lib.sudo_run(cmd, cwd=constants.CHROMITE_BIN_DIR)


	def _RemoveNetworkBridgeIgnoringErrors(name):
	"""Removes a previously created network bridge. Ignores errors."""
	_RunIgnoringErrors(['ip', 'link', 'del', name, 'type', 'bridge'])


	def _RemoveTapDeviceIgnoringErrors(name):
	"""Removes a previously created tap device. Ignores errors."""
	_RunIgnoringErrors(['ip', 'tuntap', 'del', 'mode', 'tap', name])


	def _StopVM(ssh_port, chroot_path=None):
	"""Stops a running VM instance.

	Args:
	ssh_port: (int) port to use for ssh.
	chroot_path: Location of chroot on disk to pass to cros_vm.
	"""
	logging.notice('Stopping the VM. This may take a minute.')
	cmd = [
	'%s/cros_vm' % constants.CHROMITE_BIN_DIR,
	'--stop', '--ssh-port=%d' % ssh_port,
	]
	if chroot_path:
	cmd.extend(['--chroot-path', chroot_path])
	_RunIgnoringErrors(cmd)


	def _RunIgnoringErrors(cmd):
	"""Runs the given command, ignores errors but warns user."""
	try:
	cros_build_lib.sudo_run(cmd, quiet=True)
	except cros_build_lib.RunCommandError as e:
	logging.error('Cleanup operation failed. Please run "%s" manually.',
	' '.join(cmd))
	logging.debug('Encountered error: %s', e)