terraform service dependency graph

#!/usr/bin/env python3

# Folder structure
# <env>/<group>/<template>
# <env>/<group>/stacks/<appstack>

# E.g.
#    staging/services/tf-aws-rds
#    staging/network/tf-aws-dns
#    staging/stacks/staging1/tf-aws-myapp

import re
import os
import sys
import hcl
import argparse
from graphviz import Digraph

parser = argparse.ArgumentParser()
parser.add_argument("-e", dest='envName',
                    help="ENV dir", required=True)
parser.add_argument("-s", dest='stackName',
                    help="Stack Name", required=True)
parser.add_argument("--verbose",
                    help="Up the verbosity, for debugging purposes",
                    action="count")
args = parser.parse_args()

if not os.path.exists(args.envName):
    print("Dir {} does not exist".format(args.envName))
    sys.exit(1)

stackPath = os.path.join(args.envName, 'stacks', args.stackName)
if not os.path.exists(stackPath):
        print("Dir {} does not exist".format(stackPath))
        sys.exit(1)

os.chdir(args.envName)


configs = {}
states = {}


def stripPrefix(f):
    if f.startswith('./'):
        return f[2:]
    else:
        return f


def loadIni(configPath):
    data = {}
    with open(configPath) as cp:
        for line in cp:
                line = line.strip()
                if not line.startswith('#'):
                    lineList = line.split('=', 2)
                    if len(lineList) == 2:
                        data[lineList[0]] = lineList[1].strip('"')
    return data


def loadStateFile(sfPath, baseConfig):
    with open(sfPath) as fp:
        obj = hcl.load(fp)
        sfsRaw = [s[1]['config']['key']
                  for s in obj['data']['terraform_remote_state'].items()]

        stateFiles = []
        for sfRaw in sfsRaw:
            newKey = []
            for c in sfRaw.split('/'):
                m = re.search(r'^\$\{var\.(.*)\}$', c)
                if m and m.group(1) in baseConfig:
                    newKey.append(baseConfig[m.group(1)])
                else:
                    newKey.append(c)

            stateFiles.append('/'.join(newKey))

        pointer = states
        for comp in os.path.dirname(sfPath).split('/'):
            if comp not in pointer:
                pointer[comp] = {}
            pointer = pointer[comp]

        pointer[os.path.basename(sfPath)] = stateFiles


def getConfig(sfPath):
    data = configs["."]
    dirs = os.path.dirname(sfPath).split('/')
    path = []
    for idx, val in enumerate(dirs):
        path.append(val)
        pathKey = '/'.join(path)
        if pathKey in configs:
            data.update(configs[pathKey])

    return data


for root, dirs, files in os.walk('.'):
    if os.path.exists(os.path.join(root, "config.tfvars")) and \
       not os.path.exists(os.path.join(root, "main.tf")):
        configs[stripPrefix(root)] = loadIni(
            os.path.join(root, "config.tfvars"))
        loadIni(os.path.join(root, "config.tfvars"))

for root, dirs, files in os.walk('.'):
    for file in files:
        if file == 'states.tf':
            sfPathOrig = os.path.join(root, file)
            sfPath = stripPrefix(sfPathOrig)

            baseConfig = getConfig(sfPath)
            loadStateFile(sfPathOrig, baseConfig)


def cleanName(name):
    if name.startswith('tf-aws-'):
        return name[7:]
    else:
        return name


def getNameFromState(stateFile):
    parts = stateFile.split('/')
    return "{}-{}".format(
        cleanName(parts[-1][0:-6]),
        parts[-2]
    )


registeredObjects = []


def makeNodes(name, obj, gobj):
    for n, o in obj.items():
        if 'states.tf' in o:
            if name != 'stacks' or n == args.stackName:
                gobj.node("{}-{}".format(
                    cleanName(n),
                    cleanName(name)
                ))
                registeredObjects.append("{}-{}".format(
                    cleanName(n),
                    cleanName(name)
                ))
        else:
            if name != 'stacks' or n == args.stackName:
                with gobj.subgraph(
                    name="cluster_{}_{}".format(
                        cleanName(n),
                        cleanName(name)
                    ),
                ) as gsubobj:
                    gsubobj.attr(label=cleanName(n))
                    gsubobj.attr(color='lightgray')

                    makeNodes(n, o, gsubobj)


def makeLinks(obj, prevName):
    for n, o in obj.items():
        if 'states.tf' in o:
            objName = "{}-{}".format(
                cleanName(n),
                cleanName(prevName)
            )
            if objName in registeredObjects:
                for stateFile in o['states.tf']:
                    linkState = getNameFromState(stateFile)
                    if linkState in registeredObjects:
                        g.edge(objName, linkState)
        else:
            makeLinks(o, n)


g = Digraph('unix', filename="stacks/{}/topology.gv".format(args.stackName))
g.attr(size='6,6')
g.node_attr.update(color='lightblue2', style='filled')

makeNodes(args.envName, states["."], g)
makeLinks(states["."], None)

g.view()