#!/usr/bin/python2

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

ROOT_PASSWORD = "foobar"

VAGRANTFILE = """
Vagrant.configure("2") do |config|
  config.vm.box = "{1}"
  config.vm.box_url = "{0}"
  config.vm.provider "libvirt" do |libvirt|
        libvirt.memory = 1024
        libvirt.cpu_mode = "host-model"
  end
  # disable syncing the default /vagrant share
  config.vm.synced_folder ".", "/vagrant", disabled: true
  config.ssh.username = "root"
  config.vm.provision "shell", inline: <<-SHELL
    # set root password to known value
    echo "{2}" | passwd --stdin root
    # install prerequisite to use ansible
    dnf install -y python2-dnf libselinux-python
  SHELL
end
"""


def main(argv):
    parser = argparse.ArgumentParser(description="Inventory for a Vagrant test box")
    parser.add_argument("--list", action="store_true", help="Verbose output")
    parser.add_argument('--host', help="Get host variables")
    parser.add_argument("subjects", nargs="*",
                        default=shlex.split(os.environ.get("TEST_SUBJECTS", "")))
    opts = parser.parse_args()

    try:
        if opts.host:
            data = host(opts.host)
        else:
            data = list(opts.subjects)
        sys.stdout.write(json.dumps(data, indent=4, separators=(',', ': ')))
    except RuntimeError as ex:
        sys.stderr.write("{0}: {1}\n".format(os.path.basename(sys.argv[0]), str(ex)))
        return 1

    return 0


def list(subjects):
    hosts = []
    variables = {}
    for subject in subjects:
        if subject.endswith(".box"):
            if not subject.endswith((".vagrant-libvirt.box", ".LibVirt.box")):
                sys.stderr.write("WARNING: skipping {0}: only libvirt provider supported for vagrant boxes\n".format(subject))
                continue
            hostname = os.path.basename(subject)
            vars = host(subject)
            if vars:
                hosts.append(hostname)
                variables[hostname] = vars
    return { "localhost": { "hosts": hosts, "vars": {} }, "subjects": { "hosts": hosts, "vars": {} }, "_meta": { "hostvars": variables } }


def host(box):
    # directory-separators in hostnames could be bad
    hostname = os.path.basename(box)
    null = open(os.devnull, 'w')

    try:
        lsmod = subprocess.check_output(["lsmod"], stderr=null)
    except subprocess.CalledProcessError as ex:
        raise RuntimeError("failed to run lsmod\n")

    if "kvm" not in lsmod.decode('utf-8'):
        raise RuntimeError("CPU must support KVM hardware virtualization to run vagrant\n")

    needed_pkgs = [ "vagrant", "vagrant-libvirt" ]
    pkg_error = None
    try:
        subprocess.check_call(["rpm", "--quiet", "-q"] + needed_pkgs,
                              stdout=sys.stderr.fileno())
    except subprocess.CalledProcessError as ex:
        sys.stderr.write("INFO: installing packages needed to run vagrant: {0}\n".format(" ".join(needed_pkgs)))
        if os.path.isfile("/usr/bin/dnf"):
            pkgmgr = "/usr/bin/dnf"
        else:
            pkgmgr = "/usr/bin/yum"

        try:
            subprocess.check_call([pkgmgr, "install", "-y", "-q"] + needed_pkgs,
                                  stdout=sys.stderr.fileno())
        except subprocess.CalledProcessError as ex:
            pkg_error = "Unable to install packages needed to run vagrant: {0}\n".format(" ".join(needed_pkgs))
            raise RuntimeError(pkg_error)

    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

    # Open debug log. Must be binary mode to be unbuffered.
    debuglog = open(os.path.join(artifacts, "inventory-vagrant.log"), "wb", buffering=0)

    try:
        tty = os.open("/dev/tty", os.O_WRONLY)
        os.dup2(tty, 2)
    except OSError:
        tty = None
        pass

    # A directory for temporary stuff
    directory = tempfile.mkdtemp(prefix="inventory-vagrant")

    boxname = re.sub('[^A-Za-z0-9.-_]+', '-', os.path.basename(box))

    vagrantfile = os.path.join(directory, "Vagrantfile")
    with open(vagrantfile, 'w') as f:
        f.write(VAGRANTFILE.format(box, boxname, ROOT_PASSWORD))

    # Determine if vagrant box should be kept available for diagnosis after completion
    try:
        diagnose = distutils.util.strtobool(os.getenv("TEST_DEBUG", "0"))
    except ValueError:
        diagnose = 0

    sys.stderr.write("Launching vagrant box for {0}\n".format(box))

    # Make sure the libvirtd service is running
    try:
        subprocess.check_call(["/usr/sbin/service", "libvirtd", "start"],
                              stdout=sys.stderr.fileno())
    except subprocess.CalledProcessError as ex:
        raise RuntimeError("Could not start libvirtd service")

    # set vagrant's working directory to our temporary directory
    vagrant_env = os.environ.copy()
    vagrant_env['VAGRANT_CWD'] = directory
    debuglog.write("### VAGRANT_CWD={0}\n".format(directory).encode('utf-8'))

    # And launch the actual box
    debuglog.write("### vagrant up\n".encode('utf-8'))
    proc = subprocess.Popen(["/usr/bin/vagrant", "up"],
                            env=vagrant_env, stdout=debuglog, stderr=subprocess.STDOUT)
    rc = proc.wait()

    if rc != 0:
        raise RuntimeError("vagrant failed to start with exit code {0}\n".format(rc))

    debuglog.write("### vagrant ssh-config\n".encode('utf-8'))
    try:
        ssh_config = subprocess.check_output(["/usr/bin/vagrant", "ssh-config"],
                                             env=vagrant_env, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as ex:
        raise RuntimeError("failed to retrieve vagrant SSH configuration\n")
    debuglog.write("{0}".format(ssh_config).encode('utf-8'))

    # The variables
    variables = {
        "ansible_ssh_user": "root",
        "ansible_ssh_pass": ROOT_PASSWORD,
        "ansible_ssh_common_args": "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"
    }

    for line in ssh_config.decode('utf-8').splitlines():
        l = line.split(None, 2)
        if len(l) != 2:
            continue
        key, val = l

        if key == "HostName":
            variables["ansible_ssh_host"] = val
        elif key == "Port":
            variables["ansible_ssh_port"] = val
        elif key == "IdentityFile":
            variables["ansible_ssh_private_key_file"] = val

    # verify we have all critical SSH configuration values
    for key in [ "ansible_ssh_host", "ansible_ssh_port", "ansible_ssh_private_key_file" ]:
        if not key in variables:
            raise RuntimeError("failed to retrieve vagrant SSH configuration value {0}\n".format(key))

    # wait for ssh to come up
    args = " ".join([ "{0}='{1}'".format(*item) for item in variables.items() ])
    inventory = os.path.join(directory, "inventory")
    with open(inventory, "w") as f:
        f.write("{0} {1}\n".format(hostname, args))

    debuglog.write("### inventory: {0} {1}\n".format(hostname, args).encode('utf-8'))

    ping = [
        "/usr/bin/ansible",
        "--inventory",
        inventory,
        hostname,
        "--module-name",
        "ping"
    ]

    for tries in range(0, 10):
        try:
            subprocess.check_call(ping, stdout=null, stderr=null)
            break
        except subprocess.CalledProcessError:
            time.sleep(3)
    else:
        raise RuntimeError("could not access launched vagrant box: {0}".format(box))

    # Process of our parent
    ppid = os.getppid()

    child = os.fork()
    if child:
        return 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 destroy the box
    debuglog.write("### waiting for ppid {0} to go away\n".format(ppid).encode('utf-8'))
    while True:
        time.sleep(3)

        try:
            os.kill(ppid, 0)
        except OSError:
            break # the process no longer exists

    debuglog.write("### {0} no longer exists\n".format(ppid).encode('utf-8'))

    if diagnose:
        sys.stderr.write("\n")
        sys.stderr.write("DIAGNOSE: ssh -p {0} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@{1} # password: {2}\n".format(variables["ansible_ssh_port"], variables["ansible_ssh_host"], ROOT_PASSWORD))
        sys.stderr.write("DIAGNOSE: kill {0} # when finished\n".format(os.getpid()))

        def _signal_handler(*args):
            sys.stderr.write("\nDIAGNOSE ending...\n")

        debuglog.write("### DIAGNOSE: waiting for signal before cleaning up\n".format(ppid).encode('utf-8'))

        signal.signal(signal.SIGTERM, _signal_handler)
        signal.pause()

    # Destroy the vagrant box
    debuglog.write("### vagrant destroy\n".encode('utf-8'))
    subprocess.call(["/usr/bin/vagrant", "destroy"],
                    env=vagrant_env, stdout=debuglog, stderr=subprocess.STDOUT)

    shutil.rmtree(directory)
    sys.exit(0)


if __name__ == '__main__':
    sys.exit(main(sys.argv))
