Pull infrastructure from terraform remote state

[desertfoxmb]

A little something the community might find useful.

This is a python script that will pull remote state of terraform resources (this is known to work with Azure and Terraform 1.3.0) from a storage blob in Azure. It will use that to generically map those resources to diagram node types. It's rudimentary at this point and only covers a few node types in the resource mappings. But, it will automagically create a dependency diagram from remote state.

If your resources have tags or other metadata that indicate any information for data or workflows that would be good to diagram, you could alter the script to look to include those in the edges as well/in place of a simple dependency graph which is what this currently generates.

Assuming you have remote state stored in an Azure storage blob, just fill in the related fields at the bottom with your details, update the resource type mappings as needed and give it a run.

import os
import json
import tempfile
from diagrams import Diagram, Edge
from diagrams.generic.blank import Blank
from azure.storage.blob import BlobServiceClient

# Mapping for AWS and Azure resource types to their corresponding diagrams classes
RESOURCE_MAPPING = {
    # AWS Compute
    'aws_instance': 'diagrams.aws.compute.EC2',
    'aws_lambda_function': 'diagrams.aws.compute.Lambda',

    # AWS Data
    'aws_dynamodb_table': 'diagrams.aws.database.DynamoDB',
    'aws_rds_instance': 'diagrams.aws.database.RDS',
    'aws_s3_bucket': 'diagrams.aws.storage.S3',

    # Azure Compute
    'azurerm_virtual_machine': 'diagrams.azure.compute.VM',
    'azurerm_function_app': 'diagrams.azure.compute.FunctionApp',

    # Azure Data
    'azurerm_cosmosdb_account': 'diagrams.azure.database.CosmosDb',
    'azurerm_sql_database': 'diagrams.azure.database.SQLDatabase',
    'azurerm_storage_account': 'diagrams.azure.storage.StorageAccount',
    'azurerm_cosmosdb_sql_database': 'diagrams.azure.database.SQL',
}

def load_resource_class(resource_name):
    parts = resource_name.split('.')
    module_name = '.'.join(parts[:-1])
    class_name = parts[-1]

    module = __import__(module_name, fromlist=[class_name])
    return getattr(module, class_name)

def get_remote_state_from_azure_blob_storage(account_name, account_key, container_name, blob_name):
    blob_service_client = BlobServiceClient(account_url=f"https://{account_name}.blob.core.windows.net", credential=account_key)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(blob_name)
    return json.loads(blob_client.download_blob().content_as_text())

def get_instances(resource):
    if "instances" in resource:
        return resource["instances"]
    return None

def get_dependencies(instance):
     if "dependencies" in instance:
          return instance["dependencies"]
     return None

def process_instances(instances, nodes, source):
    for instance in instances:
        dependencies = get_dependencies(instance)
        if dependencies is not None:
            generate_nodes(dependencies, nodes, source)

def generate_nodes(dependencies, nodes, source):
     for dependency in dependencies:
         if "data.terraform_remote_state" not in dependency:
            target = dependency.split('.')[1]
            if target in nodes:
                nodes[source] >> Edge(label="depends on") >> nodes[target]

def generate_diagram(state_data):
   # print(state_data)
    resources = state_data["resources"]

    with Diagram("Terraform State", show=False) as diagram:
        nodes = {}
        for resource in resources:
            resource_type = resource['type']
            name = f"{resource['name']}"

            if resource_type in RESOURCE_MAPPING:
                resource_class = load_resource_class(RESOURCE_MAPPING[resource_type])
                nodes[name] = resource_class(name)
            else:
                nodes[name] = Blank(name)
        for node in nodes:
            print(node)

        for resource in resources:
            source = f"{resource['name']}"
            instances = get_instances(resource)
            if instances is not None:
                process_instances(instances, nodes, source)
                
                

if __name__ == "__main__":
    # Replace these with your Azure Blob Storage details
    account_name = ""
    account_key = ""
    container_name = ""
    blob_name = ""

    state_data = get_remote_state_from_azure_blob_storage(account_name, account_key, container_name, blob_name)
    generate_diagram(state_data)