#!/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 json
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(
        "-P",
        "--extra-package",
        action="append",
        help="Install this package in the image. Can be specified multiple times.",
    )
    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.0 / totalsize + 0.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.0, totalsize / 1000000.0, 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
packages: %(packages)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 purge %(autoremove)s cloud-init; fstrim -va || true; shutdown -P now) &
"""


def build_seed(
    mirror,
    proxy,
    no_apt_upgrade,
    arch,
    release,
    extra_packages=None,
    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 "",
            "packages": json.dumps(extra_packages),
        }
        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.extra_package,
    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,
)
