# Copyright 2013 The ChromiumOS Authors
# 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."""

import distutils.version  # pylint: disable=import-error,no-name-in-module
import errno
import fcntl
import glob
import logging
import multiprocessing
import os
import re
import shutil
import socket
import tempfile
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 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


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.TEST_IMAGE_BIN)

    if dest_dir:
        dest_path = os.path.join(dest_dir, constants.TEST_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 = ["./image_to_vm.sh", "--test_image"]

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

        if updatable:
            # image_to_vm.sh default, for clarity.
            cmd.extend(["--disk_layout", "usb_updatable"])
        else:
            # `cros build-image` default.
            cmd.extend(["--disk_layout", "usb"])

        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
            ).stdout.strip()
            cmd.append("--to=%s" % tempdir)

        msg = "Failed to create the VM image"
        try:
            # When enter_chroot is true the cwd needs to be src/scripts.
            cros_build_lib.run(
                cmd, enter_chroot=True, cwd=constants.CROSUTILS_DIR
            )
        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.TEST_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"

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

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

        Args:
            opts: command line options.
        """
        super().__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(
                tempfile.gettempdir(), f"cros_vm_{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
        """
        standard_uefi_paths = {
            VM.ARCH_X86_64: [
                "/usr/share/qemu/edk2-x86_64-code.fd",
                "/usr/share/OVMF/OVMF_CODE_4M.fd",
            ],
            VM.ARCH_AARCH64: [
                "/usr/share/qemu/edk2-aarch64-code.fd",
                "/usr/share/qemu-efi-aarch64/QEMU_EFI.fd",
            ],
        }
        uefi_path = None
        for p in standard_uefi_paths[self.qemu_arch]:
            if os.path.exists(p):
                uefi_path = p
                break

        if not uefi_path:
            raise VMError("EDK2 QEMU firmware not found.")

        if self.qemu_arch == VM.ARCH_AARCH64:
            # UEFI firmware for ARM64 has 64Mb pflash size hardcoded
            flash_size = 64
            cros_build_lib.run(
                [
                    "dd",
                    "if=%s" % uefi_path,
                    "of=%s" % self.flash0_file,
                    "count=1",
                    "bs=%dM" % (flash_size),
                    "conv=sync",
                ],
                dryrun=self.dryrun,
            )
        elif self.qemu_arch == VM.ARCH_X86_64:
            # X86 UEFI firmware allows up to 8Mb code + vars size combined
            flash_size = 4
            # X86 UEFI also sensitive to the pflash image size and shouldn't
            # be padded to a larger boundary.
            cros_build_lib.run(
                [
                    "dd",
                    "if=%s" % uefi_path,
                    "of=%s" % self.flash0_file,
                    "bs=4K",
                    "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(flash_size * 1024 * 1024)
        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,backing_fmt=raw" % 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"
        # Newer CrOS qemu builds provide a standalone version under libexec.
        qemu_wrapper_path = os.path.join("usr/libexec/qemu/bin", qemu_exe)
        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:
                for qemu_path in (qemu_wrapper_path, qemu_exe_path):
                    qemu_path = os.path.join(qemu_dir, qemu_path)
                    if os.path.isfile(qemu_path):
                        self.qemu_path = qemu_path
                        break

        # Check chroot.
        if not self.qemu_path:
            for qemu_path in (qemu_wrapper_path, qemu_exe_path):
                qemu_path = os.path.join(self.chroot_path, qemu_path)
                print("checking", qemu_path)
                if os.path.isfile(qemu_path):
                    self.qemu_path = qemu_path
                    break

        # 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.TEST_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.TEST_IMAGE_TAR, self.cache_dir, self.board
        )
        if cache_path:
            vm_image = os.path.join(cache_path, constants.TEST_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.TEST_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):
        """Perform 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.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,svm=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",
            "-cpu",
            self.qemu_cpu,
            "-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",
                "-drive",
                "file=%s,if=pflash,format=raw,unit=0,readonly=on" % self.flash0_file,
                "-drive",
                "file=%s,if=pflash,format=raw,unit=1" % self.flash1_file,
            ]
        if self.qemu_arch == VM.ARCH_AARCH64:
            qemu_args += [
                "-M",
                "virt",
                "-pflash",
                self.flash0_file,
                "-pflash",
                self.flash1_file,
            ]
        # 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()
            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):
        """Take a VM snapshot via savevm and signal 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", encoding="utf-8") 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, encoding="utf-8") 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=5, 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().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(parser=None):
        """Parse a list of args.

        Returns:
            commandline.ArgumentParser
        """
        parser = device.Device.GetParser(parser)
        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,
            default=None,
            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. Note that if "
            "you want to pass an argument that starts with a "
            "dash, e.g. -display you will need to enclose it "
            "in quotes and add a space at the beginning: "
            '" -display ..."',
        )
        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