#!/usr/bin/python3
#
# adt-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:
# http://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 os
import time
import socket

try:
    our_base = os.environ['AUTOPKGTEST_BASE'] + '/lib'
except KeyError:
    our_base = '/usr/share/autopkgtest/python'
sys.path.insert(1, our_base)

import VirtSubproc

workdir = tempfile.mkdtemp(prefix='adt-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 subprocess.call(['which', 'distro-info'], stdout=subprocess.PIPE) == 0:
        return subprocess.check_output(['lsb_release', '-cs'],
                                       universal_newlines=True).strip()


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

    # auto-detect apt-cacher-ng
    try:
        socket.create_connection(('localhost', 3142), 0.5)
        default_proxy = 'http://10.0.2.2:3142'
    except socket.error:
        default_proxy = ''

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

    default_arch = subprocess.check_output(['dpkg', '--print-architecture'],
                                           universal_newlines=True).strip()
    qemu_default_cmds = {'i386': 'qemu-system-i386',
                         'amd64': 'qemu-system-x86_64'}

    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; a local apt-cacher-ng '
                        'gets used automatically (default: %s)' % (
                            default_proxy or 'none'))
    parser.add_argument('-s', '--disk-size', default='4G',
                        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',
                        default=qemu_default_cmds[default_arch],
                        help='qemu command (default: %(default)s)')
    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='http://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('--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.')

    return parser.parse_args()


def download_image(cloud_img_url, release, arch):
    diskname = '%s-server-cloudimg-%s-disk1.img' % (release, arch)
    image_url = '%s/%s/current/%s' % (cloud_img_url, release, diskname)
    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()

    headers = urllib.request.urlretrieve(image_url, local_image, report)[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', 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
runcmd:
 - mount -r /dev/vdb /mnt
 - env ADT_SETUP_VM_UPGRADE=%(upgr)s sh /mnt/adt-setup-vm
 - 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):
    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)

    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()}

    # 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 adt-setup-vm, copy it into the VM via the cloud-init seed iso
    for setup_vm in [os.path.join(os.path.dirname(__file__), 'adt-setup-vm'),
                     '/usr/share/autopkgtest/adt-setup-vm']:
        if os.path.exists(setup_vm):
            shutil.copy(setup_vm, workdir)
            break
    else:
        sys.stderr.write('\nCould not find adt-setup-vm\n')
        sys.exit(1)

    genisoimage = subprocess.Popen(['genisoimage', '-output', 'adt.seed',
                                    '-volid', 'cidata', '-joliet', '-rock',
                                    '-quiet', 'user-data', 'meta-data',
                                    'adt-setup-vm'],
                                   cwd=workdir)
    genisoimage.communicate()
    if genisoimage.returncode != 0:
        sys.exit(1)

    return os.path.join(workdir, 'adt.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 l in f:
                if l.startswith(b'#'):
                    continue
                l = l.strip()
                if l:
                    return l.decode()

    return 'Etc/UTC'


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

    tty_sock = os.path.join(workdir, 'ttyS0')

    qemu = subprocess.Popen([qemu_command, '-enable-kvm', '-m', '512',
                             '-localtime', '-nographic',
                             '-monitor', 'null',
                             '-net', 'user',
                             '-net', 'nic,model=virtio',
                             '-serial', 'unix:%s,server,nowait' % tty_sock,
                             '-drive', 'file=%s,if=virtio' % image,
                             '-drive', 'file=%s,if=virtio,readonly' % seed])

    try:
        if verbose:
            tty = VirtSubproc.get_unix_socket(tty_sock)

        # wait for cloud-init to finish and VM to shutdown
        with VirtSubproc.timeout(3600, 'timed out on cloud-init'):
            while qemu.poll() is None:
                time.sleep(0.2)
                if verbose:
                    block = tty.recv(4096)
                    sys.stdout.buffer.raw.write(block)
    finally:
        if qemu.poll() is None:
            qemu.terminate()
        if qemu.wait() != 0:
            sys.stderr.write('qemu failed with status %i\n' % qemu.returncode)
            # 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()
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)
boot_image(image, seed, args.qemu_command, args.verbose)
install_image(image, os.path.join(args.output_dir, 'adt-%s-%s-cloud.img' %
                                  (args.release, args.arch)))
