#!/usr/bin/python3
#
# autopkgtest-virt-lxc is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2014 Canonical Ltd.
#
# autopkgtest-virt-lxc was derived from autopkgtest-schroot and modified to suit LXC
# by Robie Basak <robie.basak@canonical.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 sys
import os
import string
import random
import subprocess
import tempfile
import time
import shutil
import argparse
from typing import List

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


capabilities = [
    "isolation-container",
    "revert",
    "revert-full-system",
    "root-on-testbed",
]

args = None
lxc_container_name = None
normal_user = None
shared_dir = None
lxcpath = None
lxcpath_opt = []


def parse_args() -> None:
    global args
    global lxcpath
    global lxcpath_opt

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Enable debugging output"
    )
    parser.add_argument(
        "-e",
        "--ephemeral",
        action="store_true",
        help="Use ephemeral overlays instead of cloning (much "
        "faster, but does not support reboot and might cause "
        "errors in some corner cases)",
    )
    parser.add_argument(
        "-P",
        "--lxcpath",
        metavar="PATH",
        help="Create containers under PATH instead of the default location",
    )
    parser.add_argument(
        "-s",
        "--sudo",
        action="store_true",
        help="Run lxc-* commands with sudo; use if you run autopkgtest as normal user",
    )
    parser.add_argument(
        "--name", help="container name (autopkgtest-lxc-XXXXXX by default)"
    )
    parser.add_argument(
        "--disk-limit",
        help="limit amount of disk space that can be used by the container."
        "Causes the container to be backed by a loop device of that size.",
    )
    parser.add_argument(
        "template", help="LXC container name that will be used as a template"
    )
    parser.add_argument(
        "lxcargs",
        nargs=argparse.REMAINDER,
        help="Additional arguments to pass to lxc-start ",
    )
    args = parser.parse_args()
    if args.debug:
        adtlog.verbosity = 2
    if args.lxcpath:
        lxcpath = args.lxcpath
        lxcpath_opt = ["--lxcpath=" + lxcpath]

    deps: List[Dependency] = [
        Executable("lxc-copy", "lxc"),
        Executable("newuidmap", "uidmap"),
    ]

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


def sudoify(command, timeout=None):
    """Prepend sudo to command with the --sudo option"""

    if args.sudo:
        if timeout:
            return ["timeout", str(timeout), "sudo"] + command
        else:
            return ["sudo"] + command
    else:
        return command


def get_available_lxc_container_name():
    """Return an LXC container name that isn't already taken.

    There is a race condition between returning this name and creation of
    the container. Ideally lxc-start-ephemeral would generate a name in a
    race free way and return the name used in machine-readable way, so that
    this function would not be necessary. See LP: #1197754.
    """
    while True:
        # generate random container name
        rnd = [random.choice(string.ascii_lowercase) for i in range(6)]
        candidate = "autopkgtest-lxc-" + "".join(rnd)

        containers = VirtSubproc.check_exec(
            sudoify(["lxc-ls", *lxcpath_opt]), outp=True, timeout=10
        )
        if candidate not in containers:
            return candidate


def wait_booted(lxc_name):
    VirtSubproc.wait_booted(
        sudoify(["lxc-attach", *lxcpath_opt, "--name", lxc_name, "--"]),
        short_command_timeout=20,
    )


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

    global capabilities, normal_user

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


def start_lxc_copy():
    argv = ["lxc-copy", "--name", args.template, "--newname", lxc_container_name]
    if lxcpath:
        argv += ["--newpath=" + lxcpath]
    if args.disk_limit:
        argv += ["--backingstorage", "loop", "--fssize", args.disk_limit]
    if args.ephemeral:
        argv += ["--ephemeral"]
        if shared_dir:
            argv += ["--mount", "bind=%s:%s" % (shared_dir, shared_dir)]
        argv += args.lxcargs
        rc = VirtSubproc.execute_timeout(
            None, 310, sudoify(argv, 300), stdout=subprocess.DEVNULL
        )[0]
        if rc != 0:
            VirtSubproc.bomb("lxc-copy with exit status %i" % rc)
    else:
        # For some reason once in a while ci.d.n fails on lxc-copy or
        # lxc-start, so lets retry at least once
        max_retries = 1
        for retry in range(max_retries + 1):
            rc = VirtSubproc.execute_timeout(
                None, 310, sudoify(argv, 300), stdout=subprocess.DEVNULL
            )[0]
            if rc == 0:
                break
            if retry == max_retries:
                VirtSubproc.bomb("lxc-copy with exit status %i" % rc)
            adtlog.warning(
                "retrying lxc-copy because it failed with exit status %i" % rc
            )
            # Absolutely not sure if waiting here is helpful and if
            # so, how long is smart.
            time.sleep(5)
        argv = ["lxc-start", *lxcpath_opt, "--name", lxc_container_name, "--daemon"]
        if shared_dir:
            argv += [
                "--define",
                "lxc.mount.entry=%s %s none bind,create=dir 0 0"
                % (shared_dir, shared_dir[1:]),
            ]
        argv += args.lxcargs
        for retry in range(max_retries + 1):
            rc = VirtSubproc.execute_timeout(
                None, 310, sudoify(argv, 300), stdout=subprocess.DEVNULL
            )[0]
            if rc == 0:
                break
            if retry == max_retries:
                # remove already copied container, tmpdir, etc.
                hook_cleanup()
                VirtSubproc.bomb("lxc-start failed with exit status %d" % rc)
            adtlog.warning(
                "retrying lxc-start because it failed with exit status %i" % rc
            )
            # Absolutely not sure if waiting here is helpful and if
            # so, how long is smart, but at least give the container
            # time to recover from ABORTING.
            time.sleep(5)

    # lxc-copy ephemeral containers support reboot too
    capabilities.append("reboot")


def hook_open():
    global args, lxc_container_name, shared_dir

    lxc_container_name = args.name or get_available_lxc_container_name()
    adtlog.debug("using container name %s" % lxc_container_name)
    if shared_dir is None:
        # shared bind mount works poorly for unprivileged containers due to
        # mapped UIDs
        if args.sudo or os.geteuid() == 0:
            shared_dir = tempfile.mkdtemp(prefix="autopkgtest-lxc.")
    else:
        # don't change the name between resets, to provide stable downtmp paths
        os.makedirs(shared_dir)
    if shared_dir:
        os.chmod(shared_dir, 0o755)
    start_lxc_copy()
    try:
        adtlog.debug("waiting for lxc guest start")
        wait_booted(lxc_container_name)
        adtlog.debug("lxc guest started")
        determine_normal_user(lxc_container_name)
        # provide a minimal and clean environment in the container
        # 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-attach's stdout/err; we need to kill these after the
        # main program (build or test script) finishes, otherwise we get
        # eternal hangs.
        VirtSubproc.auxverb = [
            "lxc-attach",
            *lxcpath_opt,
            "--name",
            lxc_container_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 2>/dev/null || 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",
            "--",
        ]
        if args.sudo:
            VirtSubproc.auxverb = ["sudo", "--preserve-env"] + VirtSubproc.auxverb
    except Exception:
        # Clean up on failure
        hook_cleanup()
        raise


def hook_downtmp(path):
    global capabilities, shared_dir

    if shared_dir:
        d = os.path.join(shared_dir, "downtmp")
        # these permissions are ugly, but otherwise we can't clean up files
        # written by the testbed when running as user
        VirtSubproc.check_exec(["mkdir", "-m", "777", d], downp=True, timeout=30)
        capabilities.append("downtmp-host=" + d)
    else:
        d = VirtSubproc.downtmp_mktemp(path, capabilities, None)
    return d


def hook_revert():
    hook_cleanup()
    hook_open()


def hook_wait_reboot(*func_args, **kwargs):
    adtlog.debug("hook_wait_reboot: waiting for container to shut down...")
    VirtSubproc.execute_timeout(
        None,
        65,
        sudoify(
            [
                "lxc-wait",
                *lxcpath_opt,
                "-n",
                lxc_container_name,
                "-s",
                "STOPPED",
                "-t",
                "60",
            ]
        ),
    )
    adtlog.debug("hook_wait_reboot: container shut down, waiting for reboot")
    wait_booted(lxc_container_name)


def hook_cleanup():
    global capabilities, shared_dir, lxc_container_name

    VirtSubproc.downtmp_remove(capabilities)

    if lxc_container_name:
        VirtSubproc.check_exec(
            sudoify(
                [
                    "lxc-destroy",
                    *lxcpath_opt,
                    "--quiet",
                    "--force",
                    "--name",
                    lxc_container_name,
                ],
                300,
            ),
            timeout=310,
        )
        lxc_container_name = None

    if shared_dir:
        shutil.rmtree(shared_dir, ignore_errors=True)


def hook_capabilities():
    return capabilities


parse_args()
VirtSubproc.main()
