Ansible filter for modifying kops cluster spec

kops-filters.py

import yaml
import time
import re

def is_valid_yaml(yaml_content):
    try:
        list(yaml.load_all(yaml_content))
        return True
    except Exception as err:
        print(err)
        return False

# Transforms all yaml documents in a yaml file using the given function.
# The yaml to transform will be added to the result, if the function returns
# `None`. If not, the result will be interpreted as an iterable producing
# yaml documents to add to the result.
def transform_yamls(yaml_content, func):
    yamls = list(yaml.load_all(yaml_content))

    result = []
    for y in yamls:
        if y is None:
            continue

        r = func(y)
        if r is None:
            r = [y]

        result.extend(r)

    return yaml.dump_all(result, explicit_start=True)


# Decorator to convert a function into a yaml transformer.
# It is basically a little bit of functional python magic - it basically
# replaces 'func' with the 'wrapper' function and calls it instead.
def transformer(func):
    def wrapper(yaml_content, *args, **kwargs):
        def transform(yaml):
            return func(yaml, *args, **kwargs)

        return transform_yamls(yaml_content, transform)

    return wrapper


# Removes all secrets from the yaml.
@transformer
def kops_remove_secrets(y):
    if y["kind"] == "Secret" or y["kind"] == "SSHCredential":
        return []

# Remove the bastion ig from the yaml.
@transformer
def kops_remove_bastion(y):
    if y["kind"] == "InstanceGroup" and y["spec"]["role"] == "Bastion":
        return []


# Get the YAML for bastion host.
@transformer
def kops_get_bastion(y):
    if y["kind"] == "InstanceGroup":
        if not y["spec"]["role"] == "Bastion":
            return []
    else:
        return []

class FilterModule(object):
    def filters(self):
        return {
            "kops_remove_secrets": kops_remove_secrets,
            "kops_remove_bastion": kops_remove_bastion,
            "kops_get_bastion": kops_get_bastion,
        }