simple-node-classifier.py

#!/usr/bin/env python
#
# Author: Hector Rivas <keymon@gmail.com>
#
# This scripts implements a simple plain external node classifier for puppet.
# It reads a plain .yaml file with all definitions returning the gived node name (or all if none)
# 
# It supports:
#  * inheretance, by setting an extra property in the node "include" and referencing
#     other hosts, or by defining classes called "external:<name>"
#  * Reference to other parameters in actual or included nodes, using the syntax %{parameter_name}
#
#  Example
#
#- name: server
#  parameters: 
#   guest_group: "guests"
#  classes: 
#   - generic_server
#
#- name: linux 
#  parameters: 
#   connect_allowed_groups: [ "group1", "group2" ]
#  classes: 
#   - external:server
#
#- name: node1.mydomain.com
#  parameters:
#   connect_allowed_groups: [ "%{connect_allowed_groups}", "group3"]
#  classes:
#   - dbserver
#   - external:linux
#  
#- name: node2.mydomain.com
#  parameters:
#   connect_allowed_groups:  ["group1","group3", "%{guest_group}"]
#  include:
#   - linux
#  classes:
#   - webserver

import yaml
import sys
import re

dbfile='test.yaml'

def load_map(dbfile):
    f = open(dbfile)
    dataArray = yaml.load(f)
    f.close()

    dataDict = {}
    for i in dataArray:
        dataDict[i['name']] = i

    return dataDict
    
def resolve_node(map, nodename):
    if not map.has_key(nodename): return None
    node = map[nodename]
    
    if node.has_key('classes'): 
        classes = set(node['classes'])
    else: 
        classes = set()

    if node.has_key('parameters'): 
        parameters = node['parameters'] 
    else: 
        parameters = {}
    
    if node.has_key('include'): 
        includes = set(node['include'])
    else: 
        includes = set()
        
    # Allow includes in classes with format external:<name>
    for c in list(classes):
        if c.strip().startswith("external:"):
            classes.remove(c)
            includes.add(c.strip().split(":",1)[1])
    
    # recurse by included classes
    merged_parameters={}
    for include in includes:
        # WARNING: Can become a infinite loop
        included=resolve_node(map, include)
        if included: 
            merged_parameters.update(included['parameters'])
            classes = classes.union(included['classes'])
        else:
            sys.stderr.write("Include definition '%s' not found" % include)
    
    # Resolve parameter inclusion... Search for %{}
    re_pattern=re.compile("\w*%\{(.*)\}\w*")
    def resolve_param(param_value, actual_param_key):
        r=re_pattern.match(param_value)
        if r:
            reference_key = r.groups()[0]
            # Search for key in actual node
            if reference_key != actual_param_key and parameters.has_key(reference_key):
                # WARNING: Can become a infinite loop
                return resolve_param(parameters[reference_key], reference_key) 
            # Or in includes 
            elif merged_parameters.has_key(reference_key):
                return merged_parameters[reference_key]
            else:
                return param_value
        else:
            return param_value
        
    for param_key, param_value in parameters.items():
        if isinstance(param_value, list):
            new_param_value=[]
            for i in param_value:
                t = resolve_param(i, param_key)
                if isinstance(t, list):
                    new_param_value.extend(t)
                else:
                    new_param_value.append(t)
        else:
            new_param_value = resolve_param(param_value, param_key)
        merged_parameters[param_key] = new_param_value
        
    node['include'] = list(includes)
    node['classes'] = list(classes)
    node['parameters'] = merged_parameters

    return node
        
map = load_map(dbfile)

if len(sys.argv) > 1: 
    nodename = sys.argv[1]
    node = resolve_node(map, nodename)
    print '---'
    print yaml.dump(node)
else:
    result_map=[]
    for name in map.keys():
        result_map.append(resolve_node(map, name))
    print yaml.dump(result_map)