Ansible inventory plugin that preserves the order of hosts in a host group

yaml_ordered.py

# Copyright (c) 2017 Ansible Project
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function

import os
from collections import MutableMapping

from yaml.nodes import MappingNode

from ansible import constants as C
from ansible.errors import AnsibleParserError
from ansible.module_utils._text import to_native
from ansible.module_utils.six import string_types
from ansible.parsing.dataloader import DataLoader
from ansible.parsing.utils.addresses import parse_address
from ansible.parsing.yaml.constructor import AnsibleConstructor
from ansible.parsing.yaml.loader import AnsibleLoader
from ansible.parsing.yaml.objects import AnsibleBaseYAMLObject
from ansible.plugins.inventory import (BaseFileInventoryPlugin, detect_range,
                                       expand_hostname_range)
from ansible.template import Templar

__metaclass__ = type

DOCUMENTATION = '''
    inventory: yaml_ordered
    short_description: Uses a specifically YAML file as inventory source.
    description:
        - YAML based inventory that preserves the order of hosts in a host group.
        - Based on the YAML inventory plugin shipped with Ansible 2.4.2.
        - "Starts with the 'all' group and has hosts/vars/children entries."
        - Host entries can have sub-entries defined, which will be treated as variables.
        - Vars entries are normal group vars.
        - "Children are 'child groups', which can also have their own vars/hosts/children and so on."
        - File MUST have a valid extension, defined in configuration
    notes:
        - >-
          It takes the place of the core plugin. Create the directory
          `ansible_plugins/inventory`, copy this file to there. Then add two
          items to `ansible.cfg`:


          [default]
          inventory_plugins = ansible_plugins/inventory


          [inventory]
          enable_plugins = yaml_ordered

        - To function it requires being whitelisted in configuration.
    options:
        yaml_extensions:
            description: list of 'valid' extensions for files containing YAML
            type: list
            default: ['.yaml', '.yml', '.json']
'''
EXAMPLES = '''
all: # keys must be unique, i.e. only one 'hosts' per group
    hosts:
        test1:
        test2:
            var1: value1
    vars:
        group_var1: value2
    children:   # key order does not matter, indentation does
        other_group:
            children:
                group_x:
                    hosts:
                        test5
            vars:
                g2_var2: value3
            hosts:
                test4:
                    ansible_host: 127.0.0.1
        last_group:
            hosts:
                test1 # same host as above, additional group membership
            vars:
                last_var: MYVALUE
'''


try:
    # Part of stdlib from Python 2.7
    from collections import OrderedDict
except ImportError:
    # Backport available from PyPI for older Python releases
    from ordereddict import OrderedDict



try:
    from __main__ import display
except ImportError:
    from ansible.utils.display import Display
    display = Display()


class AnsibleOrderedConstructor(AnsibleConstructor):
    def construct_yaml_map(self, node):
        data = AnsibleOrderedDict()
        yield data
        value = self.construct_mapping(node)
        data.update(value)
        data.ansible_pos = self._node_position_info(node)

    def construct_mapping(self, node, deep=False):
        # Most of this is from yaml.constructor.SafeConstructor.  We replicate
        # it here so that we can warn users when they have duplicate dict keys
        # (pyyaml silently allows overwriting keys)
        if not isinstance(node, MappingNode):
            raise ConstructorError(None, None,
                                   "expected a mapping node, but found %s" % node.id,
                                   node.start_mark)
        self.flatten_mapping(node)
        mapping = AnsibleOrderedDict()

        # Add our extra information to the returned value
        mapping.ansible_pos = self._node_position_info(node)

        for key_node, value_node in node.value:
            key = self.construct_object(key_node, deep=deep)
            try:
                hash(key)
            except TypeError as exc:
                raise ConstructorError("while constructing a mapping", node.start_mark,
                                       "found unacceptable key (%s)" % exc, key_node.start_mark)

            if key in mapping:
                display.warning(u'While constructing a mapping from {1}, line {2}, column {3}, found a duplicate dict key ({0}).'
                                u' Using last defined value only.'.format(key, *mapping.ansible_pos))

            value = self.construct_object(value_node, deep=deep)
            mapping[key] = value

        return mapping

AnsibleOrderedConstructor.add_constructor(
    u'tag:yaml.org,2002:map',
    AnsibleOrderedConstructor.construct_yaml_map)

AnsibleOrderedConstructor.add_constructor(
    u'tag:yaml.org,2002:python/dict',
    AnsibleOrderedConstructor.construct_yaml_map)

AnsibleOrderedConstructor.add_constructor(
    u'tag:yaml.org,2002:str',
    AnsibleOrderedConstructor.construct_yaml_str)

AnsibleOrderedConstructor.add_constructor(
    u'tag:yaml.org,2002:python/unicode',
    AnsibleOrderedConstructor.construct_yaml_str)

AnsibleOrderedConstructor.add_constructor(
    u'tag:yaml.org,2002:seq',
    AnsibleOrderedConstructor.construct_yaml_seq)

AnsibleOrderedConstructor.add_constructor(
    u'!unsafe',
    AnsibleOrderedConstructor.construct_yaml_unsafe)

AnsibleOrderedConstructor.add_constructor(
    u'!vault',
    AnsibleOrderedConstructor.construct_vault_encrypted_unicode)

AnsibleOrderedConstructor.add_constructor(
    u'!vault-encrypted',
    AnsibleOrderedConstructor.construct_vault_encrypted_unicode)


class AnsibleOrderedDict(AnsibleBaseYAMLObject, OrderedDict):
    pass


class AnsibleOrderedLoader(AnsibleLoader, AnsibleOrderedConstructor):
    pass


class OrderedDataLoader(DataLoader):
    # Override DataLoader.load instead of parser.utils.yaml._safe_load
    def load(self, data, file_name='<string>', show_content=True):
        loader = AnsibleOrderedLoader(data, file_name, self._vault.secrets)
        try:
            return loader.get_single_data()
        finally:
            try:
                loader.dispose()
            except AttributeError:
                pass  # older versions of yaml don't have dispose function, ignore


class InventoryModule(BaseFileInventoryPlugin):

    NAME = 'yaml_ordered'

    def __init__(self):

        super(InventoryModule, self).__init__()

    def verify_file(self, path):

        valid = False
        if super(InventoryModule, self).verify_file(path):
            file_name, ext = os.path.splitext(path)
            if not ext or ext in C.YAML_FILENAME_EXTENSIONS:
                valid = True
        return valid

    def parse(self, inventory, loader, path, cache=True):
        ''' parses the inventory file '''
        if not isinstance(loader, AnsibleOrderedLoader):
            self.loader = OrderedDataLoader()
            self.loader.set_vault_secrets(loader._vault.secrets)
        else:
            self.loader = loader

        self.inventory = inventory
        self.templar = Templar(loader=loader)

        try:
            data = self.loader.load_from_file(path)
        except Exception as e:
            raise AnsibleParserError(e)

        if not data:
            raise AnsibleParserError('Parsed empty YAML file')
        elif not isinstance(data, MutableMapping):
            raise AnsibleParserError('YAML inventory has invalid structure, it should be a dictionary, got: %s' % type(data))
        elif data.get('plugin'):
            raise AnsibleParserError('Plugin configuration YAML file, not YAML inventory')

        # We expect top level keys to correspond to groups, iterate over them
        # to get host, vars and subgroups (which we iterate over recursivelly)
        if isinstance(data, MutableMapping):
            for group_name in data:
                self._parse_group(group_name, data[group_name])
        else:
            raise AnsibleParserError("Invalid data from file, expected dictionary and got:\n\n%s" % to_native(data))

    def _parse_group(self, group, group_data):

        self.inventory.add_group(group)

        if isinstance(group_data, MutableMapping):
            # make sure they are dicts
            for section in ['vars', 'children', 'hosts']:
                if section in group_data:
                    # convert strings to dicts as these are allowed
                    if isinstance(group_data[section], string_types):
                        group_data[section] = {group_data[section]: None}

                    if not isinstance(group_data[section], MutableMapping):
                        raise AnsibleParserError('Invalid %s entry for %s group, requires a dictionary, found %s instead.' %
                                                 (section, group, type(group_data[section])))

            for key in group_data:
                if key == 'vars':
                    for var in group_data['vars']:
                        self.inventory.set_variable(group, var, group_data['vars'][var])

                elif key == 'children':
                    for subgroup in group_data['children']:
                        self._parse_group(subgroup, group_data['children'][subgroup])
                        self.inventory.add_child(group, subgroup)

                elif key == 'hosts':
                    for host_pattern in group_data['hosts']:
                        hosts, port = self._parse_host(host_pattern)
                        self._populate_host_vars(hosts, group_data['hosts'][host_pattern] or {}, group, port)
                else:
                    self.display.warning('Skipping unexpected key (%s) in group (%s), only "vars", "children" and "hosts" are valid' % (key, group))

    def _parse_host(self, host_pattern):
        '''
        Each host key can be a pattern, try to process it and add variables as needed
        '''
        (hostnames, port) = self._expand_hostpattern(host_pattern)

        return hostnames, port

    def _expand_hostpattern(self, hostpattern):
        '''
        Takes a single host pattern and returns a list of hostnames and an
        optional port number that applies to all of them.
        '''
        # Can the given hostpattern be parsed as a host with an optional port
        # specification?

        try:
            (pattern, port) = parse_address(hostpattern, allow_ranges=True)
        except:
            # not a recognizable host pattern
            pattern = hostpattern
            port = None

        # Once we have separated the pattern, we expand it into list of one or
        # more hostnames, depending on whether it contains any [x:y] ranges.

        if detect_range(pattern):
            hostnames = expand_hostname_range(pattern)
        else:
            hostnames = [pattern]

        return (hostnames, port)