#!/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

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_qemu import Qemu, QemuImage

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


def get_default_release():
    try:
        import distro_info
        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()
    except ImportError:
        sys.stderr.write('WARNING: python-distro-info not installed, falling '
                         'back to determining default release from currently '
                         'installed release\n')
    if shutil.which('lsb_release') is not None:
        return subprocess.check_output(['lsb_release', '-cs'],
                                       universal_newlines=True).strip()

    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='http://archive.ubuntu.com/ubuntu',
                        help='Use this mirror for apt (default: %(default)s)')
    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('-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)')

    args = parser.parse_args()

    # check our dependencies
    if (
        args.qemu_command is not None and
        shutil.which(args.qemu_command) is None
    ):
        sys.stderr.write('ERROR: QEMU command %s not found\n' %
                         args.qemu_command)
        sys.exit(1)
    if shutil.which('genisoimage') is None:
        sys.stderr.write('ERROR: genisoimage not found\n')
        sys.exit(1)
    if os.path.exists('/dev/kvm') and not os.access('/dev/kvm', os.W_OK):
        sys.stderr.write('ERROR: no permission to write /dev/kvm\n')
        sys.exit(1)

    return args


def download_image(cloud_img_url, release, arch):
    if release in ['precise', 'trusty', 'xenial']:
        diskname = '%s-server-cloudimg-%s-disk1.img' % (release, arch)
    else:
        # images were renamed in Ubuntu 16.10
        diskname = '%s-server-cloudimg-%s.img' % (release, arch)
    image_url = '%s/%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)
            diskname = 'ubuntu-%s-server-cloudimg-%s.img' % (version, arch)
            image_url = '%s/releases/%s/release/%s' % \
                (cloud_img_url, release, diskname)
    except ImportError:
        sys.stderr.write('WARNING: python-distro-info not installed, if %s '
                         'is end of life then downloading the image will fail\n' %
                         release)
    local_image = 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, local_image, 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.')

    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
locale: en_US.UTF-8
timezone: %(tz)s
password: ubuntu
chpasswd: { expire: False }
ssh_pwauth: True
manage_etc_hosts: True
apt_proxy: %(proxy)s
apt_mirror: %(mirror)s
bootcmd:
 - dpkg --add-architecture i386
runcmd:
 - sed -i 's/deb-systemd-invoke/true/' /var/lib/dpkg/info/cloud-init.prerm
 - mount -r /dev/vdb /mnt
 - 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'
     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
 - (while [ ! -e /var/lib/cloud/instance/boot-finished ]; do sleep 1; done;
    apt-get -y purge cloud-init; shutdown -P now) &
"""


def build_seed(mirror, proxy, no_apt_upgrade, 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 = ''

    if os.environ.get('AUTOPKGTEST_KEEP_APT_SOURCES'):
        keep_apt_sources = True
    elif os.environ.get('AUTOPKGTEST_APT_SOURCES_FILE'):
        shutil.copy(
            os.environ['AUTOPKGTEST_APT_SOURCES_FILE'],
            os.path.join(workdir, 'sources.list')
        )
        apt_sources_in_seed = '/mnt/sources.list'
    elif os.environ.get('AUTOPKGTEST_APT_SOURCES'):
        with open(os.path.join(workdir, 'sources.list'), 'w') as writer:
            writer.write(os.environ['AUTOPKGTEST_APT_SOURCES'])
            writer.write('\n')
        apt_sources_in_seed = '/mnt/sources.list'

    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': host_tz(),
                                       'postcmd': post_command or '',
                                       'shopt': verbose and '-x' or '',
                                       'keep_apt_sources': keep_apt_sources and 'yes' or '',
                                       'apt_sources': apt_sources_in_seed}

    # 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('sources.list')

    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 host_tz():
    '''Return host timezone.

    Defaults to Etc/UTC if it cannot be determined
    '''
    if os.path.exists('/etc/timezone'):
        with open('/etc/timezone', 'rb') as f:
            for line in f:
                if line.startswith(b'#'):
                    continue
                line = line.strip()
                if line:
                    return line.decode()

    return 'Etc/UTC'


def boot_image(
    image,
    seed,
    dpkg_architecture,
    qemu_command,
    verbose,
    timeout,
):
    print('Booting image to run cloud-init...')

    qemu = Qemu(
        cpus=1,
        dpkg_architecture=dpkg_architecture,
        images=[
            QemuImage(file=image, format='qcow2'),
            QemuImage(file=seed, format='raw', readonly=True),
        ],
        overlay=False,
        qemu_command=qemu_command,
        ram_size=512,
    )

    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):
    # 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.
    print('Moving image into final destination %s' % dest)
    desttmp = dest + str(time.time())
    shutil.move(image, desttmp)
    os.rename(desttmp, dest)


#
# main
#

args = parse_args()

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

image = download_image(args.cloud_image_url, args.release, args.arch)
resize_image(image, args.disk_size)
seed = build_seed(args.mirror, args.proxy, args.no_apt_upgrade,
                  args.metadata, args.userdata, args.post_command, args.verbose)
boot_image(image, seed, args.arch, args.qemu_command,
           args.verbose, args.timeout)
install_image(image, os.path.join(args.output_dir, 'autopkgtest-%s-%s.img' % (args.release, args.arch)))
