#!/usr/bin/python3
#
# autopkgtest-schroot is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2007 Canonical Ltd.
#
# 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 re
import grp
import pwd
import subprocess
import time
import argparse

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


capabilities = [
    "revert",
]

schroot = None
sessid = None
rootdir = None


def pw_uid(exp_name):
    try:
        return pwd.getpwnam(exp_name).pw_uid
    except KeyError:
        return None


def gr_gid(exp_name):
    try:
        return grp.getgrnam(exp_name).gr_gid
    except KeyError:
        return None


def match(exp_names, ids, extract_id):
    for exp_name in [n for n in exp_names.split(",") if n]:
        if extract_id(exp_name) in ids:
            return True
    return False


def parse_args():
    global schroot, sessid

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Enable debugging output"
    )
    parser.add_argument(
        "-s",
        "--session-id",
        help="custom schroot session ID for easy identification "
        'in "schroot --list --all-sessions"',
    )
    parser.add_argument("schroot", help="name of schroot")

    args = parser.parse_args()
    sessid = args.session_id
    schroot = args.schroot
    if args.debug:
        adtlog.verbosity = 2

    info = VirtSubproc.check_exec(
        ["schroot", "--config", "--chroot", schroot], downp=False, outp=True
    )
    cfg = {}
    ignore_re = re.compile(r"\#|\[|\s*$")
    for entry in info.split("\n"):
        if ignore_re.match(entry):
            continue
        (key, val) = entry.split("=", 1)
        cfg[key] = val

    adtlog.debug("schroot config: %s" % cfg)

    # we don't want to clobber non-ephemeral schroots
    if cfg["type"] == "directory" and cfg.get("union-type", "none") in ("", "none"):
        VirtSubproc.bomb(
            "autopkgtest-schroot requires ephemeral schroot sessions. Set a "
            '"union-type" or use a tarball schroot'
        )

    if (
        match(cfg["root-users"], [os.getuid()], pw_uid)
        or match(cfg["root-groups"], [os.getgid()] + os.getgroups(), gr_gid)
        or os.getuid() == 0
    ):
        adtlog.debug('have "root-on-testbed" capability')
        capabilities.append("root-on-testbed")

        if os.geteuid() != 0:
            username = pwd.getpwuid(os.geteuid()).pw_name
            # Passing this capability promises that it is OK for root on the testbed
            # to completely trust this user.  That is correct, because this user
            # is in schroot's `root-users` configuration - so that trust relationship
            # is already defined by the schroot administrator.
            capabilities.append("suggested-normal-user=" + username)


def hook_open():
    global schroot, sessid
    sessid = VirtSubproc.check_exec(
        ["schroot", "--quiet", "--begin-session", "--chroot", schroot]
        + (sessid and ["--session-name", sessid] or []),
        outp=True,
        fail_on_stderr=False,
    )
    VirtSubproc.auxverb = [
        "schroot",
        "--run-session",
        "--quiet",
        "--directory=/",
        "--chroot",
        sessid,
    ]
    if "root-on-testbed" in capabilities:
        VirtSubproc.auxverb += ["--user=root"]
    VirtSubproc.auxverb += ["--"]


def hook_downtmp(path):
    global capabilities

    d = VirtSubproc.downtmp_mktemp(capabilities, path, None)

    # determine mount location
    location = VirtSubproc.check_exec(
        ["schroot", "--location", "--chroot", "session:" + sessid],
        downp=False,
        outp=True,
    ).strip()
    adtlog.debug("location of schroot session: %s" % location)
    downtmp_host = "%s/%s" % (location, d)
    if os.access(downtmp_host, os.W_OK):
        adtlog.debug("%s is writable, registering as downtmp_host" % downtmp_host)
        capabilities.append("downtmp-host=" + downtmp_host)
    else:
        adtlog.debug("%s is not writable, downtmp_host not supported" % downtmp_host)

    # verify that we have /proc
    (rc, out, err) = VirtSubproc.execute_timeout(
        None, 5, VirtSubproc.auxverb + ["mountpoint", "/proc"], stdout=subprocess.PIPE
    )
    if rc != 0:
        VirtSubproc.bomb("Misconfigured schroot: /proc is not mounted")

    return d


def hook_revert():
    hook_cleanup()
    hook_open()


def hook_cleanup():
    global schroot, sessid, capabilities
    VirtSubproc.downtmp_remove(capabilities)
    # sometimes fails on EBUSY
    retries = 10
    try:
        while retries > 0:
            if (
                VirtSubproc.execute_timeout(
                    None,
                    300,
                    ["schroot", "--quiet", "--end-session", "--chroot", sessid],
                )[0]
                == 0
            ):
                break
            retries -= 1
            adtlog.info("schroot --end-session failed, retrying")
            time.sleep(1)
        else:
            adtlog.warning(
                "schroot --end-session failed repeatedly;please clean up manually"
            )
    except VirtSubproc.TimeoutError as e:
        adtlog.warning(
            "schroot --end-session timed out after %d seconds;please clean up manually",
            e.timeout_secs,
        )

    capabilities = [c for c in capabilities if not c.startswith("downtmp-host=")]


def hook_capabilities():
    return capabilities


parse_args()
VirtSubproc.main()
