#!/usr/bin/python3.6

# SPDX Licence identifier MIT
# Copyright (c) 2017-2018 Red Hat Inc.
# Author: Merlin Mathesius <merlinm@redhat.com>
#          Andrei Stepanov <astepano@redhat.com>
#          Bruno Goncalves <bgoncalv@redhat.com>

import os
import sys
import time
import json
import errno
import shlex
import shutil
import signal
import logging
import tempfile
import argparse
import subprocess
import distutils.util


EMPTY_INVENTORY = {}
LOG_FILE = "default_provisioners.log"


def print_bad_inventory():
    """Print bad inventory on any uncatched exception. This will prevent
    running playbook on localhost.
    """
    fake_host = "fake_host"
    fake_hostname = "standard-inventory-qcow2_failed_check_logs"
    hosts = [fake_host]
    bad_inv = {"localhost": {"hosts": hosts, "vars": {}},
               "subjects": {"hosts": hosts, "vars": {}},
               "_meta": {"hostvars": {fake_host: {"ansible_host": fake_hostname}}}}
    sys.stdout.write(json.dumps(bad_inv, indent=4, separators=(',', ': ')))


def get_artifact_path(path=""):
    """Return path to an artifact file in artifacts directory. If path == ""
    than return path artifacts dir.  Create artifacts dir if necessary.
    """
    artifacts = os.environ.get("TEST_ARTIFACTS", os.path.join(os.getcwd(), "artifacts"))
    try:
        os.makedirs(artifacts)
    except OSError as exc:
        if exc.errno != errno.EEXIST or not os.path.isdir(artifacts):
            raise
    return os.path.join(artifacts, path)


def inv_list(subjects, docker_extra_args):
    hosts = []
    variables = {}
    for subject in subjects:
        name, host_vars = inv_host(subject, docker_extra_args)
        if host_vars:
            hosts.append(name)
            variables[name] = host_vars
    if not hosts:
        return EMPTY_INVENTORY
    return {"localhost": {"hosts": hosts, "vars": {}},
            "subjects": {"hosts": hosts, "vars": {}},
            "_meta": {"hostvars": variables}}


def inv_host(subject, docker_extra_args):
    if not subject.startswith("docker:"):
        return None, EMPTY_INVENTORY
    image = subject[7:]
    null = open(os.devnull, 'w')
    try:
        tty = os.open("/dev/tty", os.O_WRONLY)
        os.dup2(tty, 2)
    except OSError:
        tty = None
        pass
    directory = tempfile.mkdtemp(prefix="inventory-docker")
    cidfile = os.path.join(directory, "cid")
    # Check for any additional arguments to include when starting docker container
    try:
        extra_arg_list = shlex.split(docker_extra_args)
    except ValueError:
        raise RuntimeError("Could not parse DOCKER_EXTRA_ARGS")
    logger.info("Launching Docker container for {0}".format(image))
    # Make sure the docker service is running
    try:
        subprocess.check_call(["/usr/bin/systemctl", "is-active", "--quiet", "docker"],
                              stdout=sys.stderr.fileno())
    except subprocess.CalledProcessError:
        try:
            cmd = [
                "/usr/bin/systemctl", "start", "docker"
            ]
            subprocess.check_call(cmd, stdout=sys.stderr.fileno())
        except subprocess.CalledProcessError:
            raise RuntimeError("Could not start docker service")

    # And launch the actual container
    cmd = [
        "/usr/bin/docker", "run", "--detach", "--cidfile={0}".format(cidfile),
    ] + extra_arg_list + [
        "--entrypoint=/bin/sh", image, "-c", "sleep 1000000"
    ]
    try:
        subprocess.check_call(cmd, stdout=sys.stderr.fileno())
    except subprocess.CalledProcessError:
        raise RuntimeError("Could not start container image: {0}".format(image))
    # Read out the container environment variable
    for _ in range(1, 90):
        if os.path.exists(cidfile):
            break
        time.sleep(1)
    else:
        raise RuntimeError("Could not find container file for launched container")
    with open(cidfile, "r") as f:
        name = f.read().strip()
    # Need to figure out what python interpreter to use
    interpreters = ["/usr/bin/python3", "/usr/bin/python2"]
    for interpreter in interpreters:
        check_file = ["/usr/bin/docker", "exec", "--user=root", name, "/usr/bin/ls", interpreter]
        try:
            subprocess.check_call(check_file, stdout=null, stderr=null)
            ansible_python_interpreter = interpreter
            break
        except subprocess.CalledProcessError:
            pass
    else:
        logger.error("Could not set ansible_python_interpreter.")
        return None
    # Directory to place artifacts
    artifacts = os.environ.get("TEST_ARTIFACTS", os.path.join(os.getcwd(), "artifacts"))
    # The variables
    variables = {
        "ansible_connection": "docker",
        "ansible_python_interpreter": ansible_python_interpreter
    }
    # Process of our parent
    ppid = os.getppid()
    child = os.fork()
    if child:
        return name, variables
    # Daemonize and watch the processes
    os.chdir("/")
    os.setsid()
    os.umask(0)
    if tty is None:
        tty = null.fileno()
    # Duplicate standard input to standard output and standard error.
    os.dup2(null.fileno(), 0)
    os.dup2(tty, 1)
    os.dup2(tty, 2)
    # Now wait for the parent process to go away, then kill the VM
    logger.info("docker exec -it {0} /bin/bash".format(name))
    while True:
        time.sleep(3)
        try:
            os.kill(ppid, 0)
        except OSError:
            break  # Either of the processes no longer exist
    if diagnose:
        def _signal_handler(*args):
            logger.info("Diagnose ending.")
        logger.info("kill {0} # when finished".format(os.getpid()))
        signal.signal(signal.SIGTERM, _signal_handler)
        signal.pause()
    # Dump the container logs
    try:
        os.makedirs(artifacts)
    except OSError as exc:
        if exc.errno != errno.EEXIST or not os.path.isdir(artifacts):
            raise
    log = os.path.join(artifacts, "{0}.log".format(os.path.basename(image)))
    # Kill the container
    with open(log, "w") as f:
        subprocess.call(["/usr/bin/docker", "logs", name], stdout=f.fileno())
    subprocess.call(["/usr/bin/docker", "rm", "-f", name], stdout=null)
    shutil.rmtree(directory)
    sys.exit(0)


def main(argv):
    global logger
    global diagnose
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)
    # stderr output
    conhandler = logging.StreamHandler()
    # Print to strerr by default messages with level >= warning, can be changed
    # with setting TEST_DEBUG=1.
    try:
        diagnose = distutils.util.strtobool(os.getenv("TEST_DEBUG", "0"))
    except ValueError:
        diagnose = 0
    conhandler.setLevel(logging.WARNING)
    if diagnose:
        # Collect all messages with any log level to stderr.
        conhandler.setLevel(logging.NOTSET)
    # Log format for stderr.
    log_format = "[%(levelname)-5.5s] {}: %(message)s".format(os.path.basename(__file__))
    formatter = logging.Formatter(log_format)
    conhandler.setFormatter(formatter)
    logger.addHandler(conhandler)
    parser = argparse.ArgumentParser(description="Inventory for a container image in a registry")
    parser.add_argument("--list", action="store_true", help="Verbose output")
    parser.add_argument('--host', help="Get host variables")
    parser.add_argument('--docker-extra-args', help="Extra docker arguments for launching container",
                        default=os.environ.get("TEST_DOCKER_EXTRA_ARGS", ""))
    parser.add_argument("subjects", nargs="*", default=shlex.split(os.environ.get("TEST_SUBJECTS", "")))
    opts = parser.parse_args()
    # Send logs to common logfile for all default provisioners.
    log_file = get_artifact_path(LOG_FILE)
    fhandler = logging.FileHandler(log_file)
    # Collect all messages with any log level to log file.
    fhandler.setLevel(logging.NOTSET)
    log_format = ("%(asctime)s [{}/%(threadName)-12.12s] [%(levelname)-5.5s]:"
                  "%(message)s").format(os.path.basename(__file__))
    logFormatter = logging.Formatter(log_format)
    fhandler.setFormatter(logFormatter)
    logger.addHandler(fhandler)
    logger.info("Start provisioner.")
    if opts.host:
        _, data = inv_host(opts.host, opts.docker_extra_args)
    else:
        data = inv_list(opts.subjects, opts.docker_extra_args)
    sys.stdout.write(json.dumps(data, indent=4, separators=(',', ': ')))


if __name__ == '__main__':
    ret = -1
    try:
        main(sys.argv)
        ret = 0
    except Exception:
        print_bad_inventory()
        # Backtrace stack goes to log file. If TEST_DEBUG == 1, it goes to stderr too.
        logger.info("Fatal error in provision script.", exc_info=True)
    sys.exit(ret)
