#!/usr/bin/python3
#
# autopkgtest-buildvm-ubuntu-cloud is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2014 Canonical Ltd.
#
# Build a suitable autopkgtest VM from the Ubuntu cloud images:
# https://cloud-images.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 argparse
import tempfile
import sys
import subprocess
import shutil
import atexit
import urllib.request
import urllib.error
import os
import time
import re
from typing import List

import distro_info

sys.path.insert(0, '/usr/share/autopkgtest/lib')
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))), 'lib'))

import VirtSubproc
from autopkgtest_deps import (
    Dependency,
    Executable,
    DirectoryDependency,
    FileDependency,
    KvmDependency,
    SpaceRequirement,
    check_dependencies,
)
from autopkgtest_qemu import (
    QemuFactory,
    QemuImage,
    get_host_time_zone,
)

SPACE_UNITS = {
    'b': 1,
    'k': 1024,
    'm': 1024**2,
    'g': 1024**3,
    't': 1024**4
}

workdir = tempfile.mkdtemp(prefix='autopkgtest-buildvm-ubuntu-cloud')
atexit.register(shutil.rmtree, workdir)


def get_default_release():
    try:
        return distro_info.UbuntuDistroInfo().devel()
    except distro_info.DistroDataOutdated:
        # right after a release there is no devel series, fall back to
        # stable
        sys.stderr.write('WARNING: cannot determine development release, '
                         'falling back to latest stable\n')
        return distro_info.UbuntuDistroInfo().stable()

    return None


def parse_args():
    '''Parse CLI args'''

    # use host's apt proxy
    default_proxy = ''
    try:
        p = subprocess.check_output(
            'eval `apt-config shell p Acquire::http::Proxy`; echo $p',
            shell=True, universal_newlines=True).strip()
        if p:
            # translate proxy on localhost to QEMU IP
            default_proxy = re.sub(r'localhost|127\.0\.0\.[0-9]*', '10.0.2.2', p)
    except subprocess.CalledProcessError:
        pass

    parser = argparse.ArgumentParser(fromfile_prefix_chars='@')

    default_arch = subprocess.check_output(['dpkg', '--print-architecture'],
                                           universal_newlines=True).strip()

    parser.add_argument('-a', '--arch', default=default_arch,
                        help='Ubuntu architecture (default: %(default)s)')
    parser.add_argument('-r', '--release', metavar='CODENAME',
                        default=get_default_release(),
                        help='Ubuntu release code name (default: %(default)s)')
    parser.add_argument('-m', '--mirror', metavar='URL', default=None,
                        help='Use this mirror for apt (default: http://archive.ubuntu.com/ubuntu)')
    parser.add_argument('-p', '--proxy', metavar='URL', default=default_proxy,
                        help='Use a proxy for apt (default: %s)' % (
                            default_proxy or 'none'))
    parser.add_argument('-s', '--disk-size', default='20G',
                        help='grow image by this size (default: %(default)s)')
    parser.add_argument('-o', '--output-dir', metavar='DIR', default='.',
                        help='output directory for generated image (default: '
                        'current directory)')
    parser.add_argument('--force-use-cached', action='store_true', default=False,
                        help='Use cached image without verifying')
    parser.add_argument('-q', '--qemu-command',
                        help='qemu command (default: auto)')
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='Show VM guest and cloud-init output')
    parser.add_argument('--cloud-image-url', metavar='URL',
                        default='https://cloud-images.ubuntu.com',
                        help='cloud images URL (default: %(default)s)')
    parser.add_argument('--no-apt-upgrade', action='store_true',
                        help='Do not run apt-get dist-upgrade')
    parser.add_argument('--post-command',
                        help='Run shell command in VM after setup')
    parser.add_argument('--metadata', metavar='METADATA_FILE',
                        help='Custom metadata file for cloud-init.')
    parser.add_argument('--userdata', metavar='USERDATA_FILE',
                        help='Custom userdata file for cloud-init.')
    parser.add_argument('--timeout', metavar='SECONDS', default=3600, type=int,
                        help='Timeout for cloud-init (default: %(default)i s)')
    parser.add_argument('--compress', action='store_true', default=False,
                        help='Compress final image')
    parser.add_argument('--ram-size', type=int, default=2048,
                        help='VM RAM size in MiB (default: %(default)s)')
    parser.add_argument('--cpus', type=int, default=1,
                        help='VM number of CPUs (default: %(default)s)')

    args = parser.parse_args()
    args.output_dir = os.path.expanduser(args.output_dir)

    return args


def space_str_to_bytes(space_str: str) -> int:
    '''Convert disk size string to int in bytes'''
    # disk space string format from qemu-img
    space_parts = re.match(r'(\d+)([bkmgt]?$)', space_str, re.IGNORECASE)
    if not space_parts:
        raise AttributeError("Invalid disk size %s.\n" % space_str)

    size_int = int(space_parts.group(1)) * SPACE_UNITS.get(space_parts.group(2).lower(), 1)

    return size_int


def check_requirements(factory: QemuFactory, args: argparse.Namespace):
    deps: List[Dependency] = [
        DirectoryDependency(args.output_dir),
        Executable('genisoimage', 'genisoimage'),
        Executable('qemu-img', 'qemu-utils'),
        Executable(factory.qemu_command, 'qemu-system'),
        KvmDependency(if_exists=True),
    ]

    if factory.efi_code is not None and factory.efi_package:
        deps.append(FileDependency(factory.efi_code, factory.efi_package))

    deps.append(SpaceRequirement(
        args.output_dir,
        space_str_to_bytes(args.disk_size),
        fatal=True
    ))

    if not check_dependencies(deps):
        sys.exit(1)


def get_image_info(cloud_img_url, release, arch, show_warn=True):
    if release in ['trusty', 'xenial', 'bionic'] or arch in ['armhf', 'ppc64el', 'riscv64', 's390x']:
        diskname = '%s-server-cloudimg-%s.img' % (release, arch)
        image_url = '%s/%s/current/%s' % (cloud_img_url, release, diskname)
    else:
        diskname = '%s-minimal-cloudimg-%s.img' % (release, arch)
        image_url = '%s/minimal/daily/%s/current/%s' % (cloud_img_url, release, diskname)
    try:
        import distro_info
        if release in distro_info.UbuntuDistroInfo().unsupported():
            version = distro_info.UbuntuDistroInfo().version(release)
            # LTS versions are in the "YY.MM LTS" form.
            version = version.split(" ")[0]
            if release in ['trusty', 'xenial']:
                diskname = 'ubuntu-%s-server-cloudimg-%s-disk1.img' % (version, arch)
            elif release in ['bionic']:
                diskname = 'ubuntu-%s-server-cloudimg-%s.img' % (version, arch)
            else:
                diskname = 'ubuntu-%s-minimal-cloudimg-%s.img' % (version, arch)
            image_url = '%s/releases/%s/release/%s' % \
                (cloud_img_url, release, diskname)
    except ImportError:
        if show_warn:
            sys.stderr.write('WARNING: python-distro-info not installed, if %s '
                             'is end of life then downloading the image will fail\n' %
                             release)

    return diskname, image_url


def get_cached_image(cache_dir, no_verify, cloud_img_url, release, arch):
    diskname, image_url = get_image_info(cloud_img_url, release, arch, show_warn=False)
    image_path = os.path.join(cache_dir, diskname)
    os.makedirs(cache_dir, exist_ok=True)
    if not os.path.exists(image_path):
        print('No cached image for %s/%s' % (release, arch))
        return None
    print('Found cached image %s' % image_path)

    if no_verify:
        working_path = shutil.copy(image_path, workdir)
        return working_path

    sha256sum_url = image_url.replace(diskname, 'SHA256SUMS')

    print('Checking if cached image is up to date...')
    with open(image_path, 'rb') as image_file:
        from hashlib import sha256
        data = image_file.read()
        image_sha256 = sha256(data).hexdigest()

    try:
        with urllib.request.urlopen(sha256sum_url) as resp:
            print('Got latest SHA256SUMS')
            sums = resp.read().decode('utf-8')
            for line in sums.split('\n'):
                if diskname in line:
                    good_sha256 = line.split(" ")[0]
                    if image_sha256 == good_sha256:
                        print('Cached image is up to date')
                        working_path = shutil.copy(image_path, workdir)
                        return working_path
                    else:
                        print('Cached image is not up to date')
                        return None
    except urllib.error.HTTPError as e:
        if e.code == 404:
            sys.stderr.write('\nNo SHA256SUMS found for this release/architecture\n')
        else:
            sys.stderr.write('\nDownload of SHA256SUMS failed. Try again or use '
                             '--force-use-cached to skip verification. Error: %s\n' % str(e))
        sys.exit(1)


def download_image(cloud_img_url, release, arch, cache_dir):
    diskname, image_url = get_image_info(cloud_img_url, release, arch)
    download_path = os.path.join(cache_dir, diskname)
    final_path = os.path.join(workdir, diskname)

    print('Downloading %s...' % image_url)

    is_tty = os.isatty(sys.stdout.fileno())
    download_image.lastpercent = 0

    def report(numblocks, blocksize, totalsize):
        cur_bytes = numblocks * blocksize
        if totalsize > 0:
            percent = int(cur_bytes * 100. / totalsize + .5)
        else:
            percent = 0
        if is_tty:
            if percent > download_image.lastpercent:
                download_image.lastpercent = percent
                sys.stdout.write('\r%.1f/%.1f MB (%i%%)                 ' % (
                    cur_bytes / 1000000.,
                    totalsize / 1000000.,
                    percent))
                sys.stdout.flush()
        else:
            if percent >= download_image.lastpercent + 10:
                download_image.lastpercent = percent
                sys.stdout.write('%i%% ' % percent)
                sys.stdout.flush()

    try:
        headers = urllib.request.urlretrieve(image_url, download_path, report)[1]
    except urllib.error.HTTPError as e:
        if e.code == 404:
            sys.stderr.write('\nNo image exists for this release/architecture\n')
        else:
            sys.stderr.write('\nDownload failed: %s\n' % str(e))
        sys.exit(1)
    if headers['content-type'] not in ('application/octet-stream',
                                       'text/plain'):
        sys.stderr.write('\nDownload failed!\n')
        sys.exit(1)
    print('\nDownload successful.')
    local_image = shutil.copy(download_path, final_path)

    return local_image


def resize_image(image, size):
    print('Resizing image, adding %s...' % size)
    subprocess.check_call(['qemu-img', 'resize', '-f', 'qcow2', image, '+' + size])


DEFAULT_METADATA = 'instance-id: nocloud\nlocal-hostname: autopkgtest\n'


DEFAULT_USERDATA = """#cloud-config
timezone: %(tz)s
password: ubuntu
chpasswd: { expire: False }
ssh_pwauth: True
manage_etc_hosts: True
apt:
  primary:
    - arches: default
      uri: %(mirror)s
  proxy: %(proxy)s
runcmd:
 - mount -r /dev/vdb /mnt
 - export DEBIAN_FRONTEND=noninteractive
 - env AUTOPKGTEST_SETUP_VM_UPGRADE=%(upgr)s AUTOPKGTEST_SETUP_VM_POST_COMMAND='%(postcmd)s'
     AUTOPKGTEST_APT_SOURCES_FILE='%(apt_sources)s'
     AUTOPKGTEST_KEEP_APT_SOURCES='%(keep_apt_sources)s'
     RELEASE='%(release)s'
     sh %(shopt)s /mnt/setup-testbed
 - if grep -q 'net.ifnames=0' /proc/cmdline; then ln -s /dev/null /etc/udev/rules.d/80-net-setup-link.rules; update-initramfs -u; fi
 - umount /mnt
 - fstrim -va || true
 # Ubuntu minimal images configure dpkg to avoid installing into some directories
 # from any package via the excludes file. We don't want this to happen on testbeds.
 - rm -f /etc/dpkg/dpkg.cfg.d/excludes
 - (while [ ! -e /var/lib/cloud/instance/boot-finished ]; do sleep 1; done;
    apt-get -y install ubuntu-minimal;
    apt-get -y purge %(autoremove)s cloud-init; fstrim -va || true; shutdown -P now) &
"""


def build_seed(mirror, proxy, no_apt_upgrade, arch, release,
               metadata_file=None, userdata_file=None, post_command=None,
               verbose=False):
    print('Building seed image...')
    if metadata_file:
        metadata = open(metadata_file).read()
    else:
        metadata = DEFAULT_METADATA
    with open(os.path.join(workdir, 'meta-data'), 'w') as f:
        f.write(metadata)

    keep_apt_sources = False
    apt_sources_in_seed = ''
    apt_sources_basename = ''

    if os.environ.get('AUTOPKGTEST_KEEP_APT_SOURCES'):
        keep_apt_sources = True
    elif os.environ.get('AUTOPKGTEST_APT_SOURCES_FILE'):
        with open(os.environ['AUTOPKGTEST_APT_SOURCES_FILE'], 'r') as f:
            if re.search(r'^Types:', f.read(), re.MULTILINE):
                # These look like deb822 sources.
                apt_sources_basename = 'ubuntu.sources'
            else:
                apt_sources_basename = 'sources.list'

        shutil.copy(
            os.environ['AUTOPKGTEST_APT_SOURCES_FILE'],
            os.path.join(workdir, apt_sources_basename)
        )
        apt_sources_in_seed = os.path.join('/mnt', apt_sources_basename)
    elif os.environ.get('AUTOPKGTEST_APT_SOURCES'):
        apt_sources = os.environ['AUTOPKGTEST_APT_SOURCES']
        if re.search(r'^Types:', apt_sources, re.MULTILINE):
            # These look like deb822 sources.
            apt_sources_basename = 'ubuntu.sources'
        else:
            apt_sources_basename = 'sources.list'

        with open(os.path.join(workdir, apt_sources_basename), 'w') as writer:
            writer.write(apt_sources)
            writer.write('\n')
        apt_sources_in_seed = os.path.join('/mnt', apt_sources_basename)

    upgr = no_apt_upgrade and 'false' or 'true'
    if userdata_file:
        userdata = open(userdata_file).read()
    else:
        userdata = DEFAULT_USERDATA % {'proxy': proxy or '', 'mirror': mirror,
                                       'upgr': upgr,
                                       'tz': get_host_time_zone() or 'Etc/Utc',
                                       'postcmd': post_command or '',
                                       'shopt': verbose and '-x' or '',
                                       'keep_apt_sources': keep_apt_sources and 'yes' or '',
                                       'apt_sources': apt_sources_in_seed,
                                       'release': release,
                                       'autoremove': '--autoremove' if release != 'trusty' else ''}
        if arch == 'amd64':
            userdata += '''
bootcmd:
 - dpkg --add-architecture i386'''

    # preserve proxy from host
    for v in ['http_proxy', 'https_proxy', 'no_proxy']:
        if v not in os.environ:
            continue
        val = os.environ[v]
        if v != 'no_proxy':
            # translate localhost addresses
            val = val.replace('localhost', '10.0.2.2').replace(
                '127.0.0.1', '10.0.2.2')
        userdata += ''' - [ sh, -c, 'echo "%s=%s" >> /etc/environment' ]\n''' \
            % (v, val)

    with open(os.path.join(workdir, 'user-data'), 'w') as f:
        f.write(userdata)

    # find setup-testbed, copy it into the VM via the cloud-init seed iso
    for script in [os.path.join(os.path.dirname(os.path.abspath(os.path.dirname(__file__))),
                                'setup-commands', 'setup-testbed'),
                   '/usr/share/autopkgtest/setup-commands/setup-testbed']:
        if os.path.exists(script):
            shutil.copy(script, workdir)
            break
    else:
        sys.stderr.write('\nCould not find setup-testbed script\n')
        sys.exit(1)

    files = ['user-data', 'meta-data', 'setup-testbed']

    if apt_sources_in_seed:
        files.append(apt_sources_basename)

    genisoimage = subprocess.Popen(['genisoimage', '-output', 'autopkgtest.seed',
                                    '-volid', 'cidata', '-joliet', '-rock',
                                    '-quiet'] + files,
                                   cwd=workdir)
    genisoimage.communicate()
    if genisoimage.returncode != 0:
        sys.exit(1)

    return os.path.join(workdir, 'autopkgtest.seed')


def boot_image(
    image,
    seed,
    factory,
    ram_size,
    cpus,
    verbose,
    timeout,
):
    print('Booting image to run cloud-init...')

    qemu = factory.new_session(
        cpus=cpus,
        images=[
            QemuImage(file=image, format='qcow2'),
            QemuImage(file=seed, format='raw', readonly=True),
        ],
        overlay=False,
        ram_size=ram_size,
    )

    try:
        if verbose:
            tty = qemu.get_console_socket()

        # wait for cloud-init to finish and VM to shutdown
        with VirtSubproc.timeout(timeout, 'timed out on cloud-init'):
            while qemu.subprocess.poll() is None:
                if verbose:
                    sys.stdout.buffer.raw.write(tty.recv(4096))
                else:
                    time.sleep(1)
    finally:
        ret = qemu.cleanup()
        assert ret is not None
        if ret != 0:
            sys.stderr.write('qemu failed with status %i\n' % ret)
            sys.exit(1)


def install_image(src, dest, compress=False):
    # We want to atomically update dest, and in case multiple instances are
    # running the last one should win. So we first copy/move it to
    # dest.<currenttime> (which might take a while for crossing file systems),
    # then atomically rename to dest.
    desttmp = dest + str(time.time())

    if compress:
        print('Compressing image for final destination %s' % dest)
        subprocess.run([
            'qemu-img', 'convert',
            '-f', 'qcow2',
            '-O', 'qcow2',
            '-c',
            '-p',
            src, desttmp,
        ], check=True)
    else:
        print('Moving image into final destination %s' % dest)
        shutil.move(src, desttmp)

    os.rename(desttmp, dest)


#
# main
#

args = parse_args()
factory = QemuFactory(
    dpkg_architecture=args.arch,
    qemu_command=args.qemu_command,
)
check_requirements(factory, args)

if args.release is None:
    sys.stderr.write('Unable to determine default Ubuntu release\n')
    sys.exit(1)

cache_home = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
cache_dir = cache_home + "/autopkgtest/ubuntu-cloud-images"
image = get_cached_image(cache_dir, args.force_use_cached, args.cloud_image_url, args.release, args.arch)

if image is None:
    image = download_image(args.cloud_image_url, args.release, args.arch, cache_dir)

resize_image(image, args.disk_size)
if args.mirror:
    mirror = args.mirror
else:
    if args.arch in ['i386', 'amd64']:
        mirror = 'http://archive.ubuntu.com/ubuntu'
    else:
        mirror = 'http://ports.ubuntu.com'
seed = build_seed(mirror, args.proxy, args.no_apt_upgrade, args.arch,
                  args.release, args.metadata, args.userdata,
                  args.post_command, args.verbose)

boot_image(image, seed, factory,
           args.ram_size, args.cpus, args.verbose, args.timeout)
install_image(image,
              os.path.join(args.output_dir,
                           'autopkgtest-%s-%s.img' % (args.release, args.arch)),
              compress=args.compress)
