#!/usr/bin/python3

# autopkgtest-build-qemu is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# Copyright (C) 2016-2020 Antonio Terceiro <terceiro@debian.org>.
# Copyright (C) 2019 Sébastien Delafond
# Copyright (C) 2019-2020 Simon McVittie
# Copyright (C) 2020 Christian Kastner
#
# Build a QEMU image for using with autopkgtest
#
# 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 logging
import os
import re
import shlex
import subprocess
import sys
from contextlib import suppress
from tempfile import TemporaryDirectory
from typing import Any

logger = logging.getLogger("autopkgtest-build-qemu")

DATA_PATHS = (
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "/usr/share/autopkgtest",
)

for p in reversed(DATA_PATHS):
    sys.path.insert(0, os.path.join(p, "lib"))


from autopkgtest_deps import (
    Executable,
    FileDependency,
    KvmDependency,
    check_dependencies,
)
from autopkgtest_qemu import QemuFactory

DEBIAN_KERNELS = {
    "armhf": "linux-image-armmp-lpae",
    "hppa": "linux-image-parisc",
    "i386": "linux-image-686-pae",
    "ppc64": "linux-image-powerpc64",
    "ppc64el": "linux-image-powerpc64le",
}


def can_run_arch(this: str, other: str) -> bool:
    if this == other:
        return True

    if this == "amd64":
        return other == "i386"

    return False


class UsageError(Exception):
    pass


class BuildQemu:
    def __init__(self) -> None:
        pass

    def run(self) -> None:
        default_arch = subprocess.check_output(
            ["dpkg", "--print-architecture"], text=True
        ).strip()
        default_mirror = "http://deb.debian.org/debian"

        parser = argparse.ArgumentParser()

        parser.add_argument(
            "--architecture",
            "--arch",
            default="",
            help=f"dpkg architecture name [default: {default_arch}]",
        )
        parser.add_argument(
            "--apt-proxy",
            default="",
            metavar="http://PROXY:PORT",
            help="Set apt proxy [default: auto]",
        )
        parser.add_argument(
            "--mirror",
            default="",
            metavar="URL",
            help=(
                "Debian or Debian derivative mirror " + f"[default: {default_mirror}]"
            ),
        )
        parser.add_argument(
            "--script",
            default="",
            dest="user_script",
            help="Run an extra customization script",
        )
        parser.add_argument(
            "--size",
            default="",
            help="Set image size [default: 25G]",
        )
        parser.add_argument(
            "--boot",
            default="auto",
            choices=("auto", "bios", "efi", "ieee1275", "none"),
            help=(
                "Set up a bootloader for this boot protocol "
                "[auto|bios|efi|ieee1275|none; default: auto]"
            ),
        )
        parser.add_argument(
            "--init",
            default="auto",
            choices=("auto", "systemd", "sysv-rc", "openrc"),
            help=(
                "Boot using this init system "
                "[auto|systemd|sysv-rc|openrc; default: auto]"
            ),
        )
        parser.add_argument(
            "--efi",
            dest="boot",
            action="store_const",
            const="efi",
            help="Alias for --boot=efi",
        )
        parser.add_argument(
            "--keyring",
            default=None,
            metavar="PATH",
            help="Keyring for validating a Debian derivative's apt repository",
        )
        parser.add_argument(
            "release",
            metavar="RELEASE",
            help="An apt suite or codename available from MIRROR",
        )
        parser.add_argument(
            "image",
            metavar="IMAGE",
            help="Filename of qcow2 image to create",
        )
        parser.add_argument(
            "_mirror",
            default=None,
            metavar="MIRROR",
            nargs="?",
            help="Deprecated, use --mirror instead",
        )
        parser.add_argument(
            "_architecture",
            default=None,
            metavar="ARCHITECTURE",
            nargs="?",
            help="Deprecated, use --architecture instead",
        )
        parser.add_argument(
            "_user_script",
            default=None,
            metavar="SCRIPT",
            nargs="?",
            help="Deprecated, use --script instead",
        )
        parser.add_argument(
            "_size",
            default=None,
            metavar="SIZE",
            nargs="?",
            help="Deprecated, use --size instead",
        )

        args = parser.parse_args()
        image = os.path.abspath(args.image)
        env: dict[str, str] = {}

        if args._mirror is not None:
            if args.mirror:
                parser.error(
                    "--mirror and 3rd positional argument cannot both be specified"
                )
            else:
                args.mirror = args._mirror

        if args._architecture is not None:
            if args.architecture:
                parser.error(
                    "--architecture and 4th positional argument cannot both "
                    "be specified"
                )
            else:
                args.architecture = args._architecture

        if args._user_script is not None:
            if args.user_script:
                parser.error(
                    "--script and 5th positional argument cannot both be specified"
                )
            else:
                args.user_script = args._user_script

        if args._size is not None:
            if args.size:
                parser.error(
                    "--size and 6th positional argument cannot both be specified"
                )
            else:
                args.size = args._size

        # zerofree is in /usr/sbin, make sure we have that in the PATH
        # before checking for required tools
        path = os.environ["PATH"]

        if "/usr/sbin" not in path.split(":"):
            os.environ["PATH"] = path + ":/usr/sbin:/sbin"

        factory = QemuFactory(
            boot=args.boot,
            dpkg_architecture=args.architecture,
        )
        deps = [
            Executable("fakemachine", "fakemachine", if_not_root=True),
            Executable(
                "/lib/systemd/systemd-resolved",
                "systemd-resolved",
                # Only needed if we are going to run under fakemachine
                if_not_root=True,
            ),
            Executable("qemu-img", "qemu-utils"),
            Executable(factory.qemu_command, "qemu-system", fatal=False),
            Executable("vmdb2", "vmdb2"),
            Executable("zerofree", "zerofree"),
            # Only needed if we are going to run under fakemachine, but
            # root presumably has access to /dev/kvm anyway, so there's
            # no real need to conditionalize it
            KvmDependency(if_exists=True),
        ]

        if factory.efi_code and factory.efi_package:
            deps.append(
                FileDependency(
                    factory.efi_code,
                    factory.efi_package,
                    fatal=False,
                )
            )

        if not args.architecture:
            args.architecture = default_arch
        elif not can_run_arch(default_arch, args.architecture):
            qemu_arch = QemuFactory.qemu_arch_for_dpkg_arch(args.architecture)
            # qemu-user-binfmt isn't enough: the interpreter needs to be
            # runnable inside a chroot
            deps.extend(
                [
                    Executable(f"qemu-{qemu_arch}", "qemu-user"),
                    Executable(
                        "update-binfmts",
                        "binfmt-support",
                        if_not_systemd=True,
                    ),
                ]
            )

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

        if not args.mirror:
            import distro_info

            if args.release in set(distro_info.DebianDistroInfo().unsupported()) - set(
                distro_info.DebianDistroInfo().lts_supported()
            ):
                default_mirror = "http://archive.debian.org/debian"

            args.mirror = default_mirror

        if not args.size:
            args.size = "25G"

        if not args.apt_proxy:
            args.apt_proxy = os.getenv(
                "AUTOPKGTEST_APT_PROXY",
                os.getenv("ADT_APT_PROXY", ""),
            )

        if not args.apt_proxy:
            args.apt_proxy = subprocess.check_output(
                'eval "$(apt-config shell p Acquire::http::Proxy)"; echo "$p"',
                shell=True,
                text=True,
            ).strip()

        if not args.apt_proxy:
            proxy_command = subprocess.check_output(
                'eval "$(apt-config shell p Acquire::http::Proxy-Auto-Detect)"; echo "$p"',
                shell=True,
                text=True,
            ).strip()

            if proxy_command:
                args.apt_proxy = subprocess.check_output(
                    proxy_command,
                    shell=True,
                    text=True,
                ).strip()

        if args.apt_proxy:
            env["AUTOPKGTEST_SETUP_APT_PROXY"] = args.apt_proxy

            # Set http_proxy for the initial debootstrap
            if args.apt_proxy == "DIRECT":
                with suppress(KeyError):
                    del os.environ["http_proxy"]
            else:
                env["http_proxy"] = args.apt_proxy

            # Translate proxy address on localhost to one that can be
            # accessed from the running VM
            env["AUTOPKGTEST_APT_PROXY"] = re.sub(
                r"localhost|127\.0\.0\.[0-9]*",
                "10.0.2.2",
                args.apt_proxy,
            )

            # Unprivileged builds also run inside a VM (fakemachine).
            if os.getuid() != 0:
                env["AUTOPKGTEST_SETUP_APT_PROXY"] = env["AUTOPKGTEST_APT_PROXY"]
                if args.apt_proxy != "DIRECT":
                    env["http_proxy"] = env["AUTOPKGTEST_APT_PROXY"]

        script = ""

        for d in DATA_PATHS:
            s = os.path.join(d, "setup-commands", "setup-testbed")

            if os.access(s, os.R_OK):
                script = s
                break

        if args.user_script:
            args.user_script = os.path.abspath(args.user_script)
            logger.info("Using customization script %s...", args.user_script)

        if args.architecture == default_arch:
            override_arch = None
        else:
            override_arch = args.architecture

        if args.boot == "auto" and factory.boot == "none":
            raise UsageError(
                "Unable to guess an appropriate boot protocol, use --boot to specify"
            )

        if args.init != "auto":
            env["AUTOPKGTEST_SETUP_INIT_SYSTEM"] = args.init

        with TemporaryDirectory() as temp:
            vmdb2_config = os.path.join(temp, "vmdb2.yaml")

            self.write_vmdb2_config(
                vmdb2_config,
                boot=factory.boot,
                kernel=self.choose_kernel(args.mirror, args.architecture),
                mirror=args.mirror,
                override_keyring=args.keyring,
                override_arch=override_arch,
                release=args.release,
                script=script,
                size=args.size,
                user_script=args.user_script,
            )

            try:
                argv = [
                    "vmdb2",
                    "--verbose",
                    "--image=" + image + ".raw",
                    "--log=" + image + ".log",
                    vmdb2_config,
                ]

                if os.getuid() != 0:
                    argv = self.get_fakemachine_argv(
                        argv=argv,
                        env=env,
                        expose=[script, args.user_script, image],
                        temp=temp,
                    )
                else:
                    os.environ.update(env)

                try:
                    subprocess.check_call(argv)
                except Exception:
                    subprocess.call(
                        [
                            "cat",
                            image + ".log",
                        ]
                    )
                    raise

                subprocess.check_call(
                    [
                        "qemu-img",
                        "convert",
                        "-f",
                        "raw",
                        "-O",
                        "qcow2",
                        image + ".raw",
                        image + ".new",
                    ]
                )
                # Replace a potentially existing image as atomically as
                # possible
                os.rename(image + ".new", image)
            finally:
                with suppress(FileNotFoundError):
                    os.unlink(image + ".new")

                with suppress(FileNotFoundError):
                    os.unlink(image + ".raw")

                with suppress(FileNotFoundError):
                    os.unlink(image + ".log")

    def get_fakemachine_argv(
        self,
        *,
        argv: list[str],
        env: dict[str, str],
        expose: list[str],
        temp: str,
    ) -> list[str]:
        # fakemachine doesn't really work well with arguments
        # that might contain shell metacharacters, so wrap it
        # in a script when passing it into fakemachine.
        wrapper = os.path.join(temp, "vmdb2-script")

        with open(wrapper, "w") as writer:
            writer.write("#!/bin/sh\n")
            writer.write("set -x\n")

            writer.writelines(
                f"export {k}={shlex.quote(v)}\n" for k, v in sorted(env.items())
            )

            writer.write(shlex.join(argv))
            writer.write("\n")

        os.chmod(wrapper, 0o700)

        fakemachine_argv = [
            "fakemachine",
            "-v",
            temp,
        ]
        volumes = {temp}

        for f in expose:
            if f:
                volume = os.path.dirname(os.path.abspath(f))

                if volume not in volumes:
                    volumes.add(volume)
                    fakemachine_argv.append("-v")
                    fakemachine_argv.append(volume)

        fakemachine_argv.append(wrapper)
        return fakemachine_argv

    def write_vmdb2_config(
        self,
        path: str,
        *,
        boot: str,
        kernel: str,
        mirror: str,
        override_keyring: str | None,
        override_arch: str | None,
        release: str,
        script: str,
        size: str,
        user_script: str,
    ):
        steps: list[dict[str, Any]] = []
        steps.append({"mkimg": "{{ image }}", "size": size})

        if boot == "bios":
            steps.append({"mklabel": "msdos", "device": "{{ image }}"})
        else:
            steps.append({"mklabel": "gpt", "device": "{{ image }}"})

        if boot == "efi":
            root_start = "128MiB"
            steps.append(
                {
                    "mkpart": "primary",
                    "device": "{{ image }}",
                    "start": "0%",
                    "end": root_start,
                    "tag": "efi",
                }
            )
        elif boot == "ieee1275":
            root_start = "10MiB"
            steps.append(
                {
                    "mkpart": "primary",
                    "device": "{{ image }}",
                    "start": "0%",
                    "end": root_start,
                    "tag": "prep",
                }
            )
        else:
            root_start = "0%"

        steps.append(
            {
                "mkpart": "primary",
                "device": "{{ image }}",
                "start": root_start,
                "end": "100%",
                "tag": "root",
            },
        )

        steps.append({"kpartx": "{{ image }}"})
        mkfs = {"mkfs": "ext4", "partition": "root"}
        if release in (
            "jessie",
            "stretch",
            "buster",
            "bullseye",
            "bookworm",
            "trusty",
            "xenial",
            "bionic",
            "focal",
            "jammy",
            "kinetic",
            "lunar",
        ):
            mkfs["options"] = "-O ^large_dir,^metadata_csum_seed"
            if release in ("trusty", "jessie"):
                mkfs["options"] += ",^metadata_csum"
            if release in ("jessie", "stretch", "buster", "bullseye"):
                mkfs["options"] += ",^orphan_file"
        steps.append(mkfs)
        steps.append({"mount": "root"})

        debootstrap: dict[str, Any] = {}

        debootstrap["debootstrap"] = release
        if override_arch is not None:
            debootstrap["arch"] = override_arch

        debootstrap["mirror"] = mirror
        debootstrap["target"] = "root"
        debootstrap["include"] = ["python3-minimal"]

        if override_keyring is not None:
            debootstrap["keyring"] = override_keyring

        steps.append(debootstrap)

        steps.append(
            {
                "apt": "install",
                "packages": [kernel],
                "tag": "root",
            },
        )

        if boot == "efi":
            steps.append({"mkfs": "vfat", "partition": "efi"})
            steps.append(
                {
                    "mount": "efi",
                    "dirname": "boot/efi",
                    "mount-on": "root",
                }
            )
            steps.append(
                {
                    "grub": "uefi",
                    "tag": "root",
                    "efi": "efi",
                    "console": "serial",
                }
            )
        elif boot == "ieee1275":
            steps.append(
                {
                    "grub": "ieee1275",
                    "tag": "root",
                    "prep": "prep",
                    "console": "serial",
                },
            )
        elif boot == "bios":
            steps.append(
                {
                    "grub": "bios",
                    "tag": "root",
                    "console": "serial",
                },
            )

        steps.append(
            {
                "chroot": "root",
                "shell": (
                    "passwd --delete root\n"
                    "useradd --home-dir /home/user --create-home user\n"
                    "passwd --delete user\n"
                    "echo host > /etc/hostname\n"
                    "echo '127.0.1.1\thost' >> /etc/hosts"
                ),
            },
        )

        steps.append({"fstab": "root"})

        for s in (script, user_script):
            if s:
                steps.append(
                    {
                        "shell": (
                            "export AUTOPKGTEST_BUILD_QEMU=1; "
                            + "export RELEASE="
                            + shlex.quote(release)
                            + "; "
                            + "export MIRROR="
                            + shlex.quote(mirror)
                            + "; "
                            + shlex.quote(s)
                            + ' "$ROOT"'
                        ),
                        "root-fs": "root",
                    }
                )

        with open(path, "w") as writer:
            # It's really YAML, but YAML is a superset of JSON (except in
            # pathological cases), so writing it out as JSON avoids a
            # dependency on a non-stdlib YAML library.
            json.dump({"steps": steps}, writer)

    def choose_kernel(
        self,
        mirror: str,
        architecture: str,
    ) -> str:
        if "ubuntu" in mirror:
            return "linux-image-virtual"

        return DEBIAN_KERNELS.get(architecture, "linux-image-" + architecture)


if __name__ == "__main__":
    try:
        BuildQemu().run()
    except UsageError as e:
        logger.error("%s", e)
        sys.exit(2)
    except subprocess.CalledProcessError as e:
        logger.error("%s", e)
        sys.exit(e.returncode or 1)
