Major refactor to make the script more object-oriented.

terraform.py

@@ -1,173 +1,235 @@
 #!/usr/bin/env python
 
-import bisect
 import sys
 import json
 import os
 import re
+import traceback
 from subprocess import Popen, PIPE
 
 TERRAFORM_PATH = os.environ.get('ANSIBLE_TF_BIN', 'terraform')
 TERRAFORM_DIR = os.environ.get('ANSIBLE_TF_DIR', os.getcwd())
 TERRAFORM_WS_NAME = os.environ.get('ANSIBLE_TF_WS_NAME', 'default')
 
+class TerraformState:
+    def __init__(self, state_json):
+        self.ansible_resources = []
 
-def _extract_dict(attrs, key):
-    out = {}
-    for k in attrs.keys():
-        match = re.match(r"^" + key + r"\.(.*)", k)
-        if not match or match.group(1) == "%":
-            continue
+        if "modules" in state_json:
+            # uses pre-0.12
+            self.flat_attrs = True
 
-        out[match.group(1)] = attrs[k]
-    return out
+            for module in state_json["modules"]:
+                self._filter_resources(module["resources"].values())
+        else:
+            # state format for 0.12+
+            self.flat_attrs = False
+            self._filter_resources(state_json["resources"])
 
+    def _filter_resources(self, resources):
+        for resource in resources:
+            if self.flat_attrs:
+                tf_resource = TerraformResource(resource, flat_attrs=True)
+                if tf_resource.is_ansible():
+                    self.ansible_resources.append(tf_resource)
+            else:
+                for instance in resource["instances"]:
+                    tf_resource = TerraformResource(instance, type=resource["type"])
+                    if tf_resource.is_ansible():
+                        self.ansible_resources.append(tf_resource)
 
-def _extract_list(attrs, key):
-    out = []
 
-    length_key = key + ".#"
-    if length_key not in attrs.keys():
-        return []
+class TerraformResource:
+    def __init__(self, source_json, flat_attrs=False, type=None):
+        self.flat_attrs = flat_attrs
+        self._type = type
+        self.source_json = source_json
 
-    length = int(attrs[length_key])
-    if length < 1:
-        return []
 
-    for i in range(0, length):
-        out.append(attrs["{}.{}".format(key, i)])
+    def is_ansible(self):
+        return self.type().startswith("ansible_")
 
-    return out
 
+    def type(self):
+        if self._type:
+            return self._type
+        return self.source_json["type"]
 
-def _init_group(children=None, hosts=None, vars=None):
-    if children is None:
-        children = []
-    if hosts is None:
-        hosts = []
-    if vars is None:
-        vars = {}
 
-    return {
-        "hosts": sorted(hosts),
-        "vars": vars,
-        "children": sorted(children)
-    }
+    def read_dict_attr(self, key):
+        attrs = self._raw_attributes()
 
+        if self.flat_attrs:
+            out = {}
+            for k in attrs.keys():
+                match = re.match(r"^" + key + r"\.(.*)", k)
+                if not match or match.group(1) == "%":
+                    continue
 
-def _add_host(inventory, hostname, groups, host_vars):
-    if host_vars is None:
-        host_vars = {}
-    if groups is None:
-        groups = []
+                out[match.group(1)] = attrs[k]
+            return out
+        return attrs.get(key, {})
 
-    if "all" not in groups:
-        groups.append("all")
+    def read_list_attr(self, key):
+        attrs = self._raw_attributes()
 
-    if not hostname in inventory["_meta"]["hostvars"]:
-        # make a copy because subsequent calls are going to mutate the instance
-        inventory["_meta"]["hostvars"][hostname] = host_vars.copy()
-    else:
-        inventory["_meta"]["hostvars"][hostname].update(host_vars)
+        if self.flat_attrs:
+            out = []
 
+            length_key = key + ".#"
+            if length_key not in attrs.keys():
+                return []
 
-    for group in groups:
-        if group not in inventory.keys():
-            inventory[group] = _init_group(hosts=[hostname])
-        elif hostname not in inventory[group]["hosts"]:
-            bisect.insort(inventory[group]["hosts"], hostname)
+            length = int(attrs[length_key])
+            if length < 1:
+                return []
 
+            for i in range(0, length):
+                out.append(attrs["{}.{}".format(key, i)])
 
-def _add_group(inventory, group_name, children, group_vars):
-    if children is None:
-        children = []
-    if group_name not in inventory.keys():
-        inventory[group_name] = _init_group(children=children, vars=group_vars)
-    else:
-        # Start out with support for only one "group" with a given name
-        # If there's a second group by the name, last in wins
-        inventory[group_name]["children"] = sorted(children)
-        inventory[group_name]["vars"] = group_vars
+            return out
+        return attrs.get(key, None)
+
+    def read_attr(self, key):
+        return self._raw_attributes().get(key, None)
+
+    def _raw_attributes(self):
+        if self.flat_attrs:
+            return self.source_json["primary"]["attributes"]
+        return self.source_json["attributes"]
+
+
+class AnsibleInventory:
+    def __init__(self):
+        self.groups = {}
+        self.hosts = {}
+        self.inner_json = {}
+
+    def update_hosts(self, hostname, groups=None, host_vars=None):
+        if hostname in self.hosts:
+            host = self.hosts[hostname]
+            host.update(groups=groups, host_vars=host_vars)
+        else:
+            host = AnsibleHost(hostname, groups=groups, host_vars=host_vars)
+            self.hosts[hostname] = host
+
+        if host.groups:
+            for groupname in host.groups:
+                self.update_groups(groupname, hosts=[hostname])
+
+    def update_groups(self, groupname, children=None, hosts=None, group_vars=None):
+        if groupname in self.groups:
+            self.groups[groupname].update(children=children, hosts=hosts, group_vars=group_vars)
+        else:
+            self.groups[groupname] = AnsibleGroup(groupname, children=children, hosts=hosts, group_vars=group_vars)
+
+    def add_host_resource(self, resource):
+        hostname = resource.read_attr("inventory_hostname")
+        groups = resource.read_list_attr("groups")
+        host_vars = resource.read_dict_attr("vars")
 
+        self.update_hosts(hostname, groups=groups, host_vars=host_vars)
 
-def _init_inventory():
-    return {
-        "all": _init_group(),
-        "_meta": {
-            "hostvars": {}
+    def add_host_var_resource(self, resource):
+        hostname = resource.read_attr("inventory_hostname")
+        key = resource.read_attr("key")
+        value = resource.read_attr("value")
+
+        host_vars = {key: value}
+
+        self.update_hosts(hostname, host_vars=host_vars)
+
+    def add_group_resource(self, resource):
+        groupname = resource.read_attr("inventory_group_name")
+        children = resource.read_list_attr("children")
+        group_vars = resource.read_dict_attr("vars")
+
+        self.update_groups(groupname, children=children, group_vars=group_vars)
+
+    def add_group_var_resource(self, resource):
+        groupname = resource.read_attr("inventory_group_name")
+        key = resource.read_attr("key")
+        value = resource.read_attr("value")
+
+        group_vars = {key: value}
+
+        self.update_groups(groupname, group_vars=group_vars)
+
+    def add_resource(self, resource):
+        if resource.type() == "ansible_host":
+            self.add_host_resource(resource)
+        elif resource.type() == "ansible_host_var":
+            self.add_host_var_resource(resource)
+        elif resource.type() == "ansible_group":
+            self.add_group_resource(resource)
+        elif resource.type() == "ansible_group_var":
+            self.add_group_var_resource(resource)
+
+    def to_dict(self):
+        out = {
+            "_meta": {
+                "hostvars": {}
+            }
         }
-    }
 
+        for hostname, host in self.hosts.items():
+            host.tidy()
+            out["_meta"]["hostvars"][hostname] = host.get_vars()
 
-def _handle_host(attrs, inventory):
-    host_vars = _extract_dict(attrs, "vars")
-    groups = _extract_list(attrs, "groups")
-    hostname = attrs["inventory_hostname"]
+        for groupname, group in self.groups.items():
+            group.tidy()
+            out[groupname] = group.to_dict()
 
-    _add_host(inventory, hostname, groups, host_vars)
+        return out
 
 
-def _handle_host_var(attrs, inventory):
-    host_vars = {attrs["key"]: attrs["value"]}
-    hostname = attrs["inventory_hostname"]
+class AnsibleHost:
+    def __init__(self, hostname, groups=None, host_vars=None):
+        self.hostname = hostname
+        self.groups = set(["all"])
+        self.host_vars = {}
 
-    _add_host(inventory, hostname, None, host_vars)
+        self.update(groups=groups, host_vars=host_vars)
 
-def _handle_group_var(attrs, inventory):
-    group_vars = {attrs["key"]: attrs["value"]}
-    group_name = attrs["inventory_group_name"]
+    def update(self, groups=None, host_vars=None):
+        if host_vars:
+            self.host_vars.update(host_vars)
+        if groups:
+            self.groups.update(groups)
 
-    _add_group(inventory, group_name, None, group_vars)
+    def tidy(self):
+        self.groups = sorted(self.groups)
 
-def _handle_group(attrs, inventory):
-    group_vars = _extract_dict(attrs, "vars")
-    children = _extract_list(attrs, "children")
-    group_name = attrs["inventory_group_name"]
+    def get_vars(self):
+        return dict(self.host_vars)
 
-    _add_group(inventory, group_name, children, group_vars)
+class AnsibleGroup:
+    def __init__(self, groupname, children=None, hosts=None, group_vars=None):
+        self.groupname = groupname
+        self.hosts = set()
+        self.children = set()
+        self.group_vars = {}
 
+        self.update(children=children, hosts=hosts, group_vars=group_vars)
 
-def _walk_state(tfstate, inventory):
-    if "modules" in tfstate:
-        # handle Terraform < 0.12
-        for module in tfstate["modules"]:
-            for resource in module["resources"].values():
-                if not resource["type"].startswith("ansible_"):
-                    continue
+    def update(self, children=None, hosts=None, group_vars=None):
+        if hosts:
+            self.hosts.update(hosts)
+        if children:
+            self.children.update(children)
+        if group_vars:
+            self.group_vars.update(group_vars)
 
-                attrs = resource["primary"]["attributes"]
+    def tidy(self):
+        self.hosts = sorted(self.hosts)
+        self.children = sorted(self.children)
 
-                if resource["type"] == "ansible_host":
-                    _handle_host(attrs, inventory)
-                elif resource["type"] == "ansible_group":
-                    _handle_group(attrs, inventory)
-                elif resource["type"] == "ansible_group_var":
-                    _handle_group_var(attrs, inventory)
-                elif resource["type"] == "ansible_host_var":
-                    _handle_host_var(attrs, inventory)
-    else:
-        # handle Terraform >= 0.12
-        for resource in tfstate["resources"]:
-            if resource["provider"] != "provider.ansible":
-                continue
-
-            for instance in resource["instances"]:
-                attrs = instance["attributes"]
-
-                if resource["type"] == "ansible_group":
-                    _add_group(
-                        inventory, attrs["inventory_group_name"], attrs["children"], attrs["vars"])
-                elif resource["type"] == "ansible_group_var":
-                    _add_group(
-                        inventory, attrs["inventory_group_name"], None, {attrs["key"]: attrs["value"]})
-                elif resource["type"] == "ansible_host":
-                    _add_host(
-                        inventory, attrs["inventory_hostname"], attrs["groups"], attrs["vars"])
-                elif resource["type"] == "ansible_host_var":
-                    _add_host(
-                        inventory, attrs["inventory_hostname"], None, {attrs["key"]: attrs["value"]})
-
-    return inventory
+    def to_dict(self):
+        return {
+            "children": list(self.children),
+            "hosts": list(self.hosts),
+            "vars": dict(self.group_vars)
+        }
 
 
 def _execute_shell():
@@ -193,11 +255,15 @@ def _execute_shell():
 
 def _main():
     try:
-        tfstate = _execute_shell()
-        inventory = _walk_state(tfstate, _init_inventory())
-        sys.stdout.write(json.dumps(inventory, indent=2))
+        tfstate = TerraformState(_execute_shell())
+        inventory = AnsibleInventory()
+
+        for resource in tfstate.ansible_resources:
+            inventory.add_resource(resource)
+
+        sys.stdout.write(json.dumps(inventory.to_dict(), indent=2))
     except Exception as e:
-        sys.stderr.write(str(e)+'\n')
+        traceback.print_exc(file=sys.stderr)
         sys.exit(1)