#!/usr/bin/python3
#
# This is not a stable API; for use within autopkgtest only.
#
# Part of autopkgtest.
# autopkgtest is a tool for testing Debian binary packages.
#
# Copyright 2006-2016 Canonical Ltd.
# Copyright 2016-2020 Simon McVittie
# Copyright 2017 Martin Pitt
# Copyright 2017 Jiri Palecek
# Copyright 2017-2018 Collabora Ltd.
# Copyright 2018 Thadeu Lima de Souza Cascardo
# Copyright 2019 Michael Biebl
# Copyright 2019 Raphaël Hertzog
#
# autopkgtest-virt-qemu was developed by
# Martin Pitt <martin.pitt@ubuntu.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
# See the file CREDITS for a full list of credits information (often
# installed as /usr/share/doc/autopkgtest/CREDITS).

import errno
import fcntl
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import tempfile
import time
from typing import (
    List,
    Optional,
    Sequence,
    Union,
)

import VirtSubproc
import adtlog


def find_free_port(start: int) -> int:
    '''Find an unused port in the range [start, start+50)'''

    for p in range(start, start + 50):
        adtlog.debug('find_free_port: trying %i' % p)
        try:
            lockfile = '/tmp/autopkgtest-virt-qemu.port.%i' % p
            f = None
            try:
                f = open(lockfile, 'x')
                os.unlink(lockfile)
                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except (IOError, OSError):
                adtlog.debug('find_free_port: %i is locked' % p)
                continue
            finally:
                if f:
                    f.close()

            s = socket.create_connection(('127.0.0.1', p))
            # if that works, the port is taken
            s.close()
            continue
        except socket.error as e:
            if e.errno == errno.ECONNREFUSED:
                adtlog.debug('find_free_port: %i is free' % p)
                return p
            else:
                pass

    adtlog.debug('find_free_port: all ports are taken')
    return 0


def get_cpuflag() -> Sequence[str]:
    '''Return QEMU cpu option list suitable for host CPU'''

    try:
        with open('/proc/cpuinfo', 'r') as f:
            for line in f:
                if line.startswith('flags'):
                    words = line.split()
                    if 'vmx' in words:
                        adtlog.debug(
                            'Detected KVM capable Intel host CPU, '
                            'enabling nested KVM'
                        )
                        return ['-cpu', 'kvm64,+vmx,+lahf_lm']
                    elif 'svm' in words:  # AMD kvm
                        adtlog.debug(
                            'Detected KVM capable AMD host CPU, '
                            'enabling nested KVM'
                        )
                        # FIXME: this should really be the one below
                        # for more reproducible testbeds, but nothing
                        # except -cpu host works
                        # return ['-cpu', 'kvm64,+svm,+lahf_lm']
                        return ['-cpu', 'host']
    except IOError as e:
        adtlog.warning(
            'Cannot read /proc/cpuinfo to detect CPU flags: %s' % e
        )
        # fetching CPU flags isn't critical (only used to enable
        # nested KVM), so don't fail here

    return []


class QemuImage:
    def __init__(
        self,
        file: str,
        format: Optional[str] = None,
        readonly: bool = False,
    ) -> None:
        self.file = file
        self.overlay = None     # type: Optional[str]
        self.readonly = readonly

        if format is not None:
            self.format = format
        else:
            info = json.loads(
                VirtSubproc.check_exec([
                    'qemu-img', 'info',
                    '--output=json',
                    self.file,
                ], outp=True, timeout=5)
            )

            if 'format' not in info:
                VirtSubproc.bomb('Unable to determine format of %s' % self.file)

            self.format = str(info['format'])

    def __str__(self) -> str:
        bits = []   # type: List[str]

        if self.overlay is None:
            bits.append('file={}'.format(self.file))
        else:
            bits.append('file={}'.format(self.overlay))
            bits.append('cache=unsafe')

        bits.append('if=virtio')
        bits.append('discard=unmap')
        bits.append('format={}'.format(self.format))

        if self.readonly:
            bits.append('readonly')

        return ','.join(bits)


class Qemu:
    def __init__(
        self,
        images: Sequence[Union[QemuImage, str]],
        qemu_command: str,
        cpus: int = 1,
        efi: bool = False,
        overlay: bool = False,
        overlay_dir: Optional[str] = None,
        qemu_options: Sequence[str] = (),
        ram_size: int = 1024,
        workdir: Optional[str] = None,
    ) -> None:
        """
        Constructor.

        images: Disk images for the VM. The first image is assumed to be
            the bootable, writable root filesystem, and we actually boot a
            snapshot. The remaining images are assumed to be read-only.
        qemu_command: qemu executable

        cpus: Number of vCPUs
        efi: If true, boot using OVMF/AAVMF firmware (x86/ARM only)
        overlay: If true, use a temporary overlay for first image
        overlay_dir: Store writable overlays here (default: workdir)
        qemu_options: Space-separated options for qemu
        ram_size: Amount of RAM in MiB
        workdir: Directory for temporary files (default: a random
            subdirectory of $TMPDIR)
        """

        self.cpus = cpus
        self.images = []    # type: List[QemuImage]
        self.overlay_dir = overlay_dir
        self.qemu_command = qemu_command
        self.ram_size = ram_size
        self.ssh_port = find_free_port(10022)

        if workdir is None:
            workdir = tempfile.mkdtemp(prefix='autopkgtest-qemu.')

        self.workdir = workdir      # type: Optional[str]
        os.chmod(workdir, 0o755)
        self.shareddir = os.path.join(workdir, 'shared')
        os.mkdir(self.shareddir)
        self.monitor_socket_path = os.path.join(workdir, 'monitor')
        self.ttys0_socket_path = os.path.join(workdir, 'ttyS0')
        self.ttys1_socket_path = os.path.join(workdir, 'ttyS1')

        for i, image in enumerate(images):
            if isinstance(image, QemuImage):
                self.images.append(image)
            else:
                assert isinstance(image, str)

                self.images.append(
                    QemuImage(
                        file=image,
                        format=None,
                        readonly=(i != 0),
                    )
                )

        if overlay:
            self.images[0].overlay = self.prepare_overlay(self.images[0])

        if self.ssh_port:
            adtlog.debug(
                'Forwarding local port %i to VM ssh port 22' % self.ssh_port
            )
            nic_opt = ',hostfwd=tcp:127.0.0.1:%i-:22' % self.ssh_port
        else:
            nic_opt = ''

        # start QEMU
        argv = [
            qemu_command,
            '-m', str(ram_size),
            '-smp', str(cpus),
            '-nographic',
            '-net', 'nic,model=virtio',
            '-net', 'user' + nic_opt,
            '-object', 'rng-random,filename=/dev/urandom,id=rng0',
            '-device', 'virtio-rng-pci,rng=rng0,id=rng-device0',
            '-monitor', 'unix:%s,server,nowait' % self.monitor_socket_path,
            '-serial', 'unix:%s,server,nowait' % self.ttys0_socket_path,
            '-serial', 'unix:%s,server,nowait' % self.ttys1_socket_path,
            '-virtfs',
            (
                'local,id=autopkgtest,path=%s,security_model=none,'
                'mount_tag=autopkgtest'
            ) % self.shareddir,
        ]

        for i, image in enumerate(self.images):
            argv.append('-drive')
            argv.append('index=%d,%s' % (i, image))

        if efi:
            if 'qemu-system-x86_64' in qemu_command or \
                    'qemu-system-i386' in qemu_command:
                code = '/usr/share/OVMF/OVMF_CODE.fd'
                data = '/usr/share/OVMF/OVMF_VARS.fd'
            elif 'qemu-system-aarch64' in qemu_command:
                code = '/usr/share/AAVMF/AAVMF_CODE.fd'
                data = '/usr/share/AAVMF/AAVMF_VARS.fd'
            elif 'qemu-system-arm' in qemu_command:
                code = '/usr/share/AAVMF/AAVMF32_CODE.fd'
                data = '/usr/share/AAVMF/AAVMF32_VARS.fd'
            else:
                VirtSubproc.bomb('Unknown architecture for EFI boot')

            shutil.copy(data, '%s/efivars.fd' % workdir)
            argv.append('-drive')
            argv.append('if=pflash,format=raw,read-only,file=' + code)
            argv.append('-drive')
            argv.append(
                'if=pflash,format=raw,file=%s/efivars.fd' % workdir
            )

        if os.path.exists('/dev/kvm'):
            argv.append('-enable-kvm')
            # Enable nested KVM by default on x86_64
            if (
                os.uname()[4] == 'x86_64' and
                self.qemu_command == 'qemu-system-x86_64' and
                '-cpu' not in qemu_options
            ):
                argv += get_cpuflag()

        # pass through option to qemu
        if qemu_options:
            argv.extend(qemu_options)

        self.subprocess = subprocess.Popen(
            argv,
            stdin=subprocess.DEVNULL,
            stdout=sys.stderr,
            stderr=subprocess.STDOUT,
        )   # type: Optional[subprocess.Popen[bytes]]

    @staticmethod
    def get_default_qemu_command(
        uname_m: Optional[str] = None
    ) -> str:
        uname_to_qemu_suffix = {'i[3456]86$': 'i386', '^arm': 'arm'}

        if uname_m is None:
            uname_m = os.uname()[4]

        for pattern, suffix in uname_to_qemu_suffix.items():
            if re.match(pattern, uname_m):
                return 'qemu-system-' + suffix
                break
        else:
            return 'qemu-system-' + uname_m

    @property
    def monitor_socket(self):
        return VirtSubproc.get_unix_socket(self.monitor_socket_path)

    @property
    def ttys0_socket(self):
        return VirtSubproc.get_unix_socket(self.ttys0_socket_path)

    @property
    def ttys1_socket(self):
        return VirtSubproc.get_unix_socket(self.ttys1_socket_path)

    def prepare_overlay(
        self,
        image: QemuImage,
    ) -> str:
        '''Generate a temporary overlay image'''

        # generate a temporary overlay
        if self.overlay_dir is not None:
            overlay = os.path.join(
                self.overlay_dir,
                os.path.basename(image.file) + '.overlay-%s' % time.time()
            )
        else:
            workdir = self.workdir
            assert workdir is not None
            overlay = os.path.join(workdir, 'overlay.img')

        adtlog.debug('Creating temporary overlay image in %s' % overlay)
        VirtSubproc.check_exec(
            [
                'qemu-img', 'create',
                '-f', 'qcow2',
                '-F', image.format,
                '-b', os.path.abspath(image.file),
                overlay,
            ],
            outp=True,
            timeout=300,
        )
        return overlay

    def cleanup(self) -> Optional[int]:
        ret = None

        if self.subprocess is not None:
            self.subprocess.terminate()
            ret = self.subprocess.wait()
            self.subprocess = None

        if self.workdir is not None:
            shutil.rmtree(self.workdir)
            self.workdir = None

        return ret
