#!/usr/bin/python3
#
# autopkgtest-virt-lxd is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2015 Canonical Ltd.
#
# Author: Martin Pitt <martin.pitt@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 csv
import os
import random
import re
import string
import subprocess
import sys
import time

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 adtlog
import VirtSubproc
from autopkgtest_deps import Dependency, Executable, check_dependencies

capabilities = [
    "reboot",
    "revert",
    "revert-full-system",
    "root-on-testbed",
    "snapshot",
]

args = None
instance_name = None
normal_user = None


def parse_args() -> None:
    global args

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--lxd",
        dest="command",
        default="",
        action="store_const",
        const="lxc",
        help="Use LXD",
    )
    parser.add_argument(
        "--incus", dest="command", action="store_const", const="incus", help="Use Incus"
    )
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Enable debugging output"
    )
    parser.add_argument(
        "-r",
        "--remote",
        default="",
        help="Run instance on given remote host instead of "
        'the current default; see "lxc remote list" or '
        '"incus remote list"',
    )
    parser.add_argument(
        "-p",
        "--project",
        default="",
        help="Run instance in given project instead of the current one;"
        'see "lxc project list" or "incus project list"',
    )
    parser.add_argument(
        "--vm",
        action="store_true",
        help="Explicitly request a VM instead of a container",
    )
    parser.add_argument("--name", help="Instance name (autogenerated by default)")
    parser.add_argument(
        "--delete-existing",
        help="If there is an instance with the same name, replace it",
        action="store_true",
    )
    parser.add_argument("image", help="LXD/Incus image name")
    parser.add_argument(
        "launchargs",
        nargs=argparse.REMAINDER,
        help='Additional arguments to pass to "lxc init" or "incus init"',
    )
    args = parser.parse_args()
    if args.name is not None and (
        not re.fullmatch(r"[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?", args.name)
        or args.name.isdecimal()
    ):
        parser.error(
            "--name must contain 1-63 ASCII letters, digits or hyphens, "
            "must not start or end with a hyphen, and must not be purely numeric"
        )
    if args.debug:
        adtlog.verbosity = 2

    if not args.command:
        tail = os.path.basename(sys.argv[0])

        if tail == "autopkgtest-virt-lxd":
            args.command = "lxc"
        elif tail == "autopkgtest-virt-incus":
            args.command = "incus"
        else:
            parser.error(
                "Must be invoked as autopkgtest-virt-lxd or "
                "autopkgtest-virt-incus, or with --lxd or --incus option"
            )

    # dnsmasq is in /usr/sbin, make sure we have that in the PATH
    # before checking for required tools
    path = os.environ.get("PATH", "/usr/bin:/bin")

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

    deps: list[Dependency] = [
        Executable("dnsmasq", "dnsmasq-base"),
    ]

    if args.command == "lxc":
        deps.append(Executable("lxc", "lxd (snap)"))
    elif args.command == "incus":
        deps.append(Executable("incus", "incus"))

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

    # Probe for the current remote & project, and store them for later use.
    # This allows a user to start autopkgtest and then use `lxc project switch`
    # or `lxc remote switch` without autopkgtest losing access to the testbed.
    if not args.remote:
        args.remote = VirtSubproc.check_exec(
            client_cmd(["remote", "get-default"]),
            outp=True,
            timeout=30,
            fail_on_stderr=False,
        )

    if args.remote and not args.remote.endswith(":"):
        args.remote += ":"

    if not args.project:
        # No good way to identify the current project; replace this with
        # `{lxc,incus} project get-current` when most LXD versions support it
        # https://github.com/canonical/lxd/pull/17975
        projects = VirtSubproc.check_exec(
            client_cmd(["project", "list", args.remote, "--format=csv"]),
            outp=True,
            timeout=30,
            fail_on_stderr=False,
        ).splitlines()

        for project in projects:
            if " (current)" in project:
                args.project = project.split(" ")[0]
                break


def client_cmd(arguments: list[str]) -> list[str]:
    assert args is not None
    if args.project:
        return [args.command, "--project", args.project] + arguments
    else:
        return [args.command] + arguments


def get_available_instance_name(command):
    """Return an instance name that isn't already taken"""

    if args.command == "lxc":
        prefix = "autopkgtest-lxd-"
    else:
        assert args.command == "incus"
        prefix = "autopkgtest-incus-"

    while True:
        # generate random instance name
        rnd = [random.choice(string.ascii_lowercase) for i in range(6)]
        candidate = prefix + "".join(rnd)

        rc = VirtSubproc.execute_timeout(
            None,
            60,
            client_cmd(["info", candidate]),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
        )[0]
        if rc != 0:
            return candidate


def wait_booted():
    VirtSubproc.wait_booted(
        client_cmd(["exec", "--disable-stdin", instance_name, "--"]), timeout=900
    )


def determine_normal_user():
    """Check for a normal user to run tests as."""

    global normal_user

    # get the first UID in the Debian Policy §9.2.2 "dynamically allocated
    # user account" range
    cmd = client_cmd(
        [
            "exec",
            "--disable-stdin",
            instance_name,
            "--",
            "sh",
            "-c",
            (
                "getent passwd | sort -t: -nk3 | "
                "awk -F: '{if ($3 >= 1000 && $3 <= 59999) { print $1; exit } }'"
            ),
        ]
    )
    out = VirtSubproc.execute_timeout(None, 60, cmd, stdout=subprocess.PIPE)[1].strip()
    if out:
        normal_user = out
        capabilities.append("suggested-normal-user=" + normal_user)
        adtlog.debug(f'determine_normal_user: got user "{normal_user}"')
    else:
        adtlog.debug("determine_normal_user: no uid in [1000,59999] available")


def warn_if_slow_snapshots():
    """Warn if the instance's storage makes snapshots expensive."""

    try:
        pool_name = subprocess.check_output(
            client_cmd(
                [
                    "list",
                    args.remote,
                    f"^{instance_name.removeprefix(args.remote)}$",
                    "--columns=b",
                    "--format=csv",
                ]
            ),
            stdin=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            text=True,
            timeout=30,
        ).strip()
        pools = subprocess.check_output(
            client_cmd(
                ["storage", "list", args.remote, "--columns=n,D", "--format=csv"]
            ),
            stdin=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            text=True,
            timeout=30,
        )
        driver = dict(csv.reader(pools.splitlines()))[pool_name]
    except (OSError, subprocess.SubprocessError, csv.Error, KeyError, ValueError) as e:
        # Storage inspection is advisory and must not prevent opening a testbed.
        adtlog.debug(f"Cannot determine storage pool driver: {e}")
        return

    if driver in ("dir", "lvmcluster"):
        adtlog.warning(
            f"Storage pool {pool_name} uses {driver}; snapshot operations will be slow. "
            "Consider using btrfs or zfs."
        )


def hook_open():
    global instance_name

    extra_args = args.launchargs

    if ":" not in args.image:
        image_name = args.remote + args.image
    else:
        image_name = args.image

    vm = args.vm
    if not vm:
        # For unambiguous images (e.g. local images, or a specific fingerprint)
        # the image type can be detected. This does not work for simplestreams
        # remotes, where a VM image can share its alias with a container image;
        # there --vm is the only way to select the VM.
        image_info = VirtSubproc.check_exec(
            client_cmd(["image", "info", image_name]),
            outp=True,
            timeout=30,
            fail_on_stderr=False,
        ).splitlines()
        vm = "Type: virtual-machine" in image_info

    if vm:
        adtlog.debug("will start a VM (with isolation-machine capability)")
        capabilities.append("isolation-machine")
        extra_args += ["--vm"]
    else:
        adtlog.debug("will start a container (with isolation-container capability)")
        capabilities.append("isolation-container")

    instance_name = args.remote + (
        args.name or get_available_instance_name(args.command)
    )
    adtlog.debug(f"using instance name {instance_name}")
    if args.delete_existing:
        out = VirtSubproc.check_exec(
            client_cmd(
                [
                    "list",
                    args.remote,
                    f"^{instance_name.removeprefix(args.remote)}$",
                    "--format=csv",
                ]
            ),
            outp=True,
            timeout=600,
        )
        if out:
            VirtSubproc.check_exec(
                client_cmd(["delete", "--force", instance_name]), timeout=600
            )
    VirtSubproc.check_exec(
        client_cmd(["init", image_name, instance_name]) + extra_args,
        outp=True,
        timeout=600,
    )
    try:
        warn_if_slow_snapshots()
        # Set outp=True to keep startup messages off the virt server's protocol
        VirtSubproc.check_exec(
            client_cmd(["start", instance_name]), outp=True, timeout=600
        )
        adtlog.debug("waiting for instance start")
        wait_booted()
        adtlog.debug("instance started")
        determine_normal_user()
        # provide a minimal and clean environment in the instance
        # We also want to avoid exiting with 255 as that's auxverb's exit code
        # if the auxverb itself failed; so we translate that to 253.
        # Tests or builds sometimes leak background processes which might still
        # be connected to lxc/incus exec's stdout/err; we need to kill these
        # after the main program (build or test script) finishes, otherwise we
        # get eternal hangs.
        VirtSubproc.auxverb = client_cmd(
            [
                "exec",
                instance_name,
                "--",
                "env",
                "-i",
                "bash",
                "-c",
                (
                    "set -a; "
                    "[ -r /etc/environment ] && . /etc/environment 2>/dev/null || true; "
                    "[ -r /etc/default/locale ] && . /etc/default/locale 2>/dev/null || true; "
                    "[ -r /etc/profile ] && . /etc/profile >/dev/null 2>&1 || true; "
                    "set +a;"
                    '"$@"; RC=$?; [ $RC != 255 ] || RC=253; '
                    "set -e;"
                    "myout=$(readlink /proc/$$/fd/1);"
                    "myerr=$(readlink /proc/$$/fd/2);"
                    'myout="${myout/[/\\\\[}"; myout="${myout/]/\\\\]}";'
                    'myerr="${myerr/[/\\\\[}"; myerr="${myerr/]/\\\\]}";'
                    "PS=$(ls -l /proc/[0-9]*/fd/* 2>/dev/null | sed -nr '\\#('\"$myout\"'|'\"$myerr\"')# { s#^.*/proc/([0-9]+)/.*$#\\1#; p}'|sort -u);"
                    'KILL="";'
                    "for pid in $PS; do"
                    "    [ $pid -ne $$ ] && [ $pid -ne $PPID ] || continue;"
                    '    KILL="$KILL $pid";'
                    "done;"
                    '[ -z "$KILL" ] || kill -9 $KILL >/dev/null 2>&1 || true;'
                    "exit $RC"
                ),
                "--",
            ]
        )
    except Exception:
        # Clean up on failure
        VirtSubproc.execute_timeout(
            None, 300, client_cmd(["delete", "--force", instance_name])
        )
        raise


def hook_downtmp(path):
    if path is None:
        return VirtSubproc.downtmp_mktemp()
    VirtSubproc.downtmp_mkdir(path)
    return path


def hook_revert():
    hook_cleanup()
    hook_open()


def hook_snapshot_create(name):
    if "isolation-machine" in capabilities:
        # Flush guest-cached writes before taking a disk-only snapshot.
        VirtSubproc.check_exec(["sync"], downp=True, timeout=60)
    if args.command == "lxc":
        command = ["snapshot", instance_name, name, "--no-expiry"]
    else:
        command = ["snapshot", "create", instance_name, name, "--no-expiry"]
    VirtSubproc.check_exec(client_cmd(command), timeout=900)


def hook_snapshot_restore(name):
    if args.command == "lxc":
        command = ["restore", instance_name, name]
    else:
        command = ["snapshot", "restore", instance_name, name]
    # The server stops and restarts a running instance as part of restoration.
    VirtSubproc.check_exec(client_cmd(command), timeout=900)
    state = VirtSubproc.check_exec(
        client_cmd(
            [
                "list",
                args.remote,
                f"^{instance_name.removeprefix(args.remote)}$",
                "--format=csv",
                "--columns=s",
            ]
        ),
        outp=True,
        timeout=30,
    )
    if state == "STOPPED":
        VirtSubproc.check_exec(client_cmd(["start", instance_name]), timeout=600)
    wait_booted()


def hook_snapshot_delete(name):
    if args.command == "lxc":
        command = ["delete", f"{instance_name}/{name}"]
    else:
        command = ["snapshot", "delete", instance_name, name]
    VirtSubproc.check_exec(client_cmd(command), timeout=600)


def get_boot_id():
    try:
        (rc, out, _) = VirtSubproc.execute_timeout(
            None,
            60,
            client_cmd(
                [
                    "exec",
                    "--disable-stdin",
                    instance_name,
                    "--",
                    "cat",
                    "/proc/sys/kernel/random/boot_id",
                ]
            ),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        if rc != 0:
            return

        return out.strip()
    except IndexError:
        return


def hook_prepare_reboot():
    initial_boot_id = get_boot_id()
    adtlog.debug(
        f"hook_prepare_reboot: fetching boot_id before reboot: {initial_boot_id}"
    )

    return {"initial_boot_id": initial_boot_id}


def hook_wait_reboot(*func_args, **kwargs):
    adtlog.debug("hook_wait_reboot: waiting for instance to shut down...")
    # "lxc/incus exec" exits with 0 when the instance stops, so just wait
    # longer than our timeout
    initial_boot_id = kwargs["initial_boot_id"]

    adtlog.debug(
        f"hook_wait_reboot: instance had boot_id {initial_boot_id}, waiting for reboot"
    )

    for retry in range(20):
        time.sleep(5)

        current_boot_id = get_boot_id()

        # instance is probably in the very late stages of shutting down, just
        # keep trying, if this persists we'll bomb out later on
        if not current_boot_id:
            continue

        if current_boot_id != initial_boot_id:
            adtlog.debug(
                f"hook_wait_reboot: instance boot_id is {current_boot_id} - has rebooted (initial boot_id {initial_boot_id})"
            )
            break
        else:
            adtlog.debug(
                f"hook_wait_reboot: instance boot_id is {current_boot_id} - has not rebooted (initial boot_id {initial_boot_id})"
            )
    else:
        VirtSubproc.bomb(f"timed out waiting for instance {instance_name} to restart")

    adtlog.debug("hook_wait_reboot: instance restarted, waiting for boot to finish")
    wait_booted()


def hook_cleanup():
    VirtSubproc.check_exec(
        client_cmd(["delete", "--force", instance_name]), timeout=600
    )


def hook_capabilities():
    return capabilities


parse_args()
VirtSubproc.main()
