#!/usr/bin/python3

import json
import os
import sys
import subprocess
import time


def main(argv):
    """
    Run all standard inventory scripts and return their merged output.

    Note the standard inventory scripts clean up their spawned hosts
    when they detect their parent processes go away. To accomodate
    that behavior, this script forks a child process to run the
    inventory scripts, send back the merged inventory, and then wait
    for the parent of this script (ansible) to die before silently
    exiting. In the mean time, this script outputs the merged
    inventory gathered by the child and then exits.
    """

    # Keep track of our parent
    waitpid = os.getppid()

    tty = err_to_tty()
    pipein, pipeout = os.pipe()

    childpid = os.fork()
    if childpid == 0:
        # this is the child process

        # close the inherited input side of the pipe
        os.close(pipein)

        # run and merge output from standard inventory scripts
        merged_data = merge_standard_inventories(argv[1:])

        # send merged data to parent via output pipe
        os.write(pipeout, merged_data.encode('utf-8'))

        # close the pipe so the parent knows we are done
        os.close(pipeout)

        # wait for the grandparent process to exit
        linger(waitpid, tty)

        # exit cleanly
        sys.exit(0)

    # this is the parent process

    # close the inherited output side of the pipe
    os.close(pipeout)

    # send eveything from the child to stdout
    while True:
        data = os.read(pipein, 999)
        if not data:
            os.close(pipein)
            break
        sys.stdout.write(data.decode('utf-8'))

    return 0


def merge_standard_inventories(args):

    inventory_dir = os.environ.get(
        "TEST_DYNAMIC_INVENTORY_DIRECTORY", "/usr/share/ansible/inventory")

    ignore_ext_string = os.environ.get(
        "ANSIBLE_INVENTORY_IGNORE", "~, .orig, .bak, .ini, .cfg, .retry, .pyc, .pyo")
    ignore_ext_list = []
    for s in ignore_ext_string.split(','):
        if s.strip():
            ignore_ext_list.append(s.strip())
    inventory_ignore_extensions = tuple(ignore_ext_list)

    merged = Inventory()

    for i in os.path.exists(inventory_dir) and os.listdir(inventory_dir) or []:
        ipath = os.path.join(inventory_dir, i)
        if not i.startswith("standard-inventory-"):
            continue
        if i.endswith(inventory_ignore_extensions):
            continue
        if not os.access(ipath, os.X_OK):
            continue

        cmd = [ipath] + args

        try:
            inv_out = subprocess.check_output(cmd, stdin=None, close_fds=True)
        except subprocess.CalledProcessError:
            raise RuntimeError("Could not run: {0}".format(str(cmd)))

        merged.merge(inv_out.decode('utf-8'))

    return merged.dumps()


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

    return tty


def linger(waitpid, tty=None):
    # Go into daemon mode and watch the process

    null = open(os.devnull, 'w')

    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 watched process to go away, then return
    while True:
        time.sleep(3)

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

    return


class Inventory:
    """
    Merge JSON data from standard test dynamic inventory scripts.

    Note: This class is very specific to the JSON data written by the
    ansible dynamic inventory scripts that are provided by the
    standard-test-roles package. In particular, it insists on finding
    and generating "subjects" and "localhost" members.
    """

    def __init__(self):
        self.hosts = []
        self.variables = {}

    def merge_files(self, files):
        for f in files:
            with open(f) as ifile:
                s = ifile.read()
                self.merge(s)

    def merge(self, s):
        # parse provided string as JSON
        inventory = json.loads(s)

        if not isinstance(inventory, dict):
            return

        if "subjects" not in inventory or not isinstance(inventory["subjects"], dict):
            return
        if "hosts" not in inventory["subjects"] or not isinstance(inventory["subjects"]["hosts"], list):
            return

        for h in inventory["subjects"]["hosts"]:
            self.hosts.append(h)

        if "_meta" not in inventory or not isinstance(inventory["_meta"], dict):
            raise ValueError(
                "inventory JSON does not contain the expected [_meta] dictionary")
        if "hostvars" not in inventory["_meta"] or not isinstance(inventory["_meta"]["hostvars"], dict):
            raise ValueError(
                "inventory JSON does not contain the expected [_meta][hostvars] dict")

        for h in inventory["_meta"]["hostvars"]:
            if not isinstance(inventory["_meta"]["hostvars"][h], dict):
                raise ValueError(
                    "inventory JSON does not contain the expected [_meta][hostvars][{0}] dict".format(h))

            self.variables[h] = inventory["_meta"]["hostvars"][h]

        if "localhost" not in inventory or not isinstance(inventory["localhost"], dict):
            raise ValueError(
                "inventory JSON does not contain the expected [localhost] dictionary")
        if "hosts" not in inventory["localhost"] or not isinstance(inventory["localhost"]["hosts"], list):
            raise ValueError(
                "inventory JSON does not contain the expected [localhost][hosts] list")

    def dumps(self):
        if not self.hosts:
            empty_inventory = {}
            return json.dumps(empty_inventory, indent=4, separators=(',', ': '))
        data = {"subjects": {"hosts": self.hosts, "vars": {}}, "localhost": {
            "hosts": self.hosts, "vars": {}}, "_meta": {"hostvars": self.variables}}
        return json.dumps(data, indent=4, separators=(',', ': '))


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