Merge pull request #38 from RoseSecurity/add-support-for-higher-count…

…s-of-resources

Update Tests, Provider Name Cleanup, and Flowchart Functions

.github/workflows/build.yaml

@@ -1,4 +1,4 @@
-name: "Build CLI and attach to GitHub release"
+name: "Build CLI and Attach to GitHub Release"
 
 on:
   release:
@@ -60,6 +60,10 @@ jobs:
           push: true
           tags: rosesecurity/terramaid:latest
 
+      - name: "Verify Image"
+        run: |
+            docker pull ${{ steps.build.outputs.image }}:${{ steps.build.outputs.tag}}
+
   homebrew:
     name: "Bump Homebrew Formula"
     runs-on: ubuntu-latest
@@ -72,4 +76,3 @@ jobs:
           formula-path: Formula/t/terramaid.rb
         env:
           COMMITTER_TOKEN: ${{ secrets.COMMITTER_TOKEN }}
-

.github/workflows/test.yml

@@ -1,4 +1,4 @@
-name: Test
+name: "Test Terramaid Outputs"
 
 on:
   pull_request:
@@ -7,16 +7,42 @@ on:
       - main
 
 jobs:
-  test:
+  test_aws:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
       - uses: hashicorp/setup-terraform@v3
       - uses: actions/setup-go@v5
         with:
-          go-version: '1.22.2'
+          go-version: "1.22.2"
           cache: false
       - run: |
           make build
-          build/terramaid -w test/
+          build/terramaid -w test/aws
+          cat Terramaid.md
+  test_gcp:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: hashicorp/setup-terraform@v3
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.22.2"
+          cache: false
+      - run: |
+          make build
+          build/terramaid -w test/gcp
+          cat Terramaid.md
+  test_az:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: hashicorp/setup-terraform@v3
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.22.2"
+          cache: false
+      - run: |
+          make build
+          build/terramaid -w test/az
           cat Terramaid.md

.goreleaser.yaml

@@ -9,7 +9,8 @@ before:
 builds:
   - env:
       - CGO_ENABLED=0
-    ldflags: -s -w -X main.version={{ .Version }}
+    ldflags:
+      - '-s -w -X "github.com/rosesecurity/Terramaid/cmd.Version={{.Version}}"'
     goos:
       - linux
       - windows

Makefile

@@ -13,7 +13,7 @@ fmt: ## Format Go files
 	gofumpt -w .
 
 build: ## Build Terramaid
-	$(GO) build -ldflags="-s -w" -o build/$(BINARY_NAME) main.go
+	env $(if $(GOOS),GOOS=$(GOOS)) $(if $(GOARCH),GOARCH=$(GOARCH)) $(GO) build -o build/$(BINARY_NAME) -ldflags "-X 'github.com/rosesecurity/Terramaid/cmd.Version=${VERSION}'" main.go
 
 install: ## Install dependencies
 	$(GO) install ./...@latest

README.md

@@ -22,18 +22,39 @@ Terramaid transforms your Terraform resources and plans into visually appealing
 
 ```mermaid
 flowchart TD
- subgraph Hashicorp
- subgraph Terraform
-  aws_db_instance.example_db["aws_db_instance.example_db"]
-  aws_instance.example_instance["aws_instance.example_instance"]
-  aws_s3_bucket.logs["aws_s3_bucket.logs"]
-  aws_s3_bucket.test["aws_s3_bucket.test"]
-  aws_s3_bucket_policy.logs_policy["aws_s3_bucket_policy.logs_policy"]
-  aws_s3_bucket_policy.test_policy["aws_s3_bucket_policy.test_policy"]
-  aws_s3_bucket_policy.logs_policy --> aws_s3_bucket.logs
-  aws_s3_bucket_policy.test_policy --> aws_s3_bucket.test
- end
- end
+	subgraph Terraform
+		subgraph Aws
+			aws_db_instance.main_db["aws_db_instance.main_db"]
+			aws_instance.app_server["aws_instance.app_server"]
+			aws_instance.web_server["aws_instance.web_server"]
+			aws_lb.web["aws_lb.web"]
+			aws_lb_listener.web["aws_lb_listener.web"]
+			aws_lb_target_group.web["aws_lb_target_group.web"]
+			aws_lb_target_group_attachment.web["aws_lb_target_group_attachment.web"]
+			aws_s3_bucket.logs["aws_s3_bucket.logs"]
+			aws_s3_bucket.test["aws_s3_bucket.test"]
+			aws_s3_bucket_policy.logs_policy["aws_s3_bucket_policy.logs_policy"]
+			aws_s3_bucket_policy.test_policy["aws_s3_bucket_policy.test_policy"]
+			aws_security_group.db["aws_security_group.db"]
+			aws_security_group.web["aws_security_group.web"]
+			aws_subnet.private["aws_subnet.private"]
+			aws_subnet.public["aws_subnet.public"]
+			aws_vpc.main["aws_vpc.main"]
+		end
+		aws_lb.web --> aws_security_group.web
+		aws_lb.web --> aws_subnet.public
+		aws_lb_listener.web --> aws_lb.web
+		aws_lb_listener.web --> aws_lb_target_group.web
+		aws_lb_target_group.web --> aws_vpc.main
+		aws_lb_target_group_attachment.web --> aws_instance.web_server
+		aws_lb_target_group_attachment.web --> aws_lb_target_group.web
+		aws_s3_bucket_policy.logs_policy --> aws_s3_bucket.logs
+		aws_s3_bucket_policy.test_policy --> aws_s3_bucket.test
+		aws_security_group.db --> aws_security_group.web
+		aws_security_group.web --> aws_vpc.main
+		aws_subnet.private --> aws_vpc.main
+		aws_subnet.public --> aws_vpc.main
+	end
 ```
 
 > [!TIP]

internal/flowchart.go

@@ -11,8 +11,10 @@ import (
 )
 
 type Node struct {
-	ID    string
-	Label string
+	ID       string
+	Label    string
+	Count    int
+	Provider string
 }
 
 type Edge struct {
@@ -21,114 +23,118 @@ type Edge struct {
 }
 
 type Graph struct {
-	Nodes []Node
-	Edges []Edge
+	Nodes   []Node
+	Edges   []Edge
+	NodeMap map[string]int
 }
 
-// Removes unnecessary parts from the label
+var labelCleaner = regexp.MustCompile(`\s*\(expand\)|\s*\(close\)|\[root\]\s*|"`)
+
+// CleanLabel removes unnecessary parts from the label
 func CleanLabel(label string) string {
-	re := regexp.MustCompile(`\s*\(expand\)|\s*\(close\)|\[root\]\s*|"`)
-	return re.ReplaceAllString(label, "")
+	return labelCleaner.ReplaceAllString(label, "")
 }
 
-// Removes unnecessary parts from the ID
+// CleanID removes unnecessary parts from the ID
 func CleanID(id string) string {
-	re := regexp.MustCompile(`\s*\(expand\)|\s*\(close\)|\[root\]\s*|"`)
-	return re.ReplaceAllString(id, "")
+	return labelCleaner.ReplaceAllString(id, "")
 }
 
-// Extracts the provider for separate subgraph
+// ExtractProvider extracts the provider for separate subgraph
 func ExtractProvider(label string) string {
-	if strings.Contains(label, "provider") {
-		parts := strings.Split(label, "/")
-		if len(parts) > 2 {
-			return parts[len(parts)-2]
-		}
+	parts := strings.Split(label, "_")
+	if len(parts) > 0 {
+		// Remove quotes from the provider name
+		return strings.ReplaceAll(parts[0], "\"", "")
 	}
 	return ""
 }
 
-// Transforms the parsed graph into cleaned nodes and edges
+// TransformGraph transforms the parsed graph into cleaned nodes and edges
 func TransformGraph(graph *gographviz.Graph) Graph {
 	nodes := []Node{}
 	edges := []Edge{}
+	nodeMap := make(map[string]int)
 
 	for _, node := range graph.Nodes.Nodes {
 		cleanedID := CleanID(node.Name)
 		cleanedLabel := CleanLabel(node.Attrs["label"])
-		if cleanedLabel != "" && !strings.Contains(cleanedLabel, "provider") {
-			nodes = append(nodes, Node{ID: cleanedID, Label: cleanedLabel})
+		provider := ExtractProvider(cleanedLabel)
+		if cleanedLabel != "" {
+			nodeMap[cleanedLabel]++
+			nodes = append(nodes, Node{ID: cleanedID, Label: cleanedLabel, Count: nodeMap[cleanedLabel], Provider: provider})
 		}
 	}
 
 	for _, edge := range graph.Edges.Edges {
 		fromLabel := CleanLabel(graph.Nodes.Lookup[edge.Src].Attrs["label"])
 		toLabel := CleanLabel(graph.Nodes.Lookup[edge.Dst].Attrs["label"])
-		if fromLabel != "" && toLabel != "" && !strings.Contains(fromLabel, "provider") && !strings.Contains(toLabel, "provider") {
+		if fromLabel != "" && toLabel != "" {
 			edges = append(edges, Edge{From: CleanID(edge.Src), To: CleanID(edge.Dst)})
 		}
 	}
 
-	return Graph{Nodes: nodes, Edges: edges}
+	return Graph{Nodes: nodes, Edges: edges, NodeMap: nodeMap}
 }
 
-// Converts a gographviz graph to a Mermaid.js compatible string.
+// ConvertToMermaidFlowchart converts a gographviz graph to a Mermaid.js compatible string.
 // It accepts a graph, direction, and an optional subgraph name.
 func ConvertToMermaidFlowchart(graph *gographviz.Graph, direction string, subgraphName string) (string, error) {
 	var sb strings.Builder
 
-	// Capitalize the provider name
 	caser := cases.Title(language.English)
-	// Validate the direction of the flowchart. Valid options are: TB, TD, BT, RL, LR
 	validDirections := map[string]bool{
 		"TB": true, "TD": true, "BT": true, "RL": true, "LR": true,
 	}
 	if !validDirections[direction] {
 		return "", fmt.Errorf("invalid direction %s: valid options are: TB, TD, BT, RL, LR", direction)
 	}
-	// Start Mermaid graph definition
+
 	sb.WriteString("```mermaid\n")
 	sb.WriteString("flowchart " + direction + "\n")
 
-	// Add subgraph for providers
+	if subgraphName != "" {
+		sb.WriteString(fmt.Sprintf("\tsubgraph %s\n", subgraphName))
+	}
+
 	providerSubgraphs := make(map[string]bool)
 	for _, n := range graph.Nodes.Nodes {
-		label := n.Attrs["label"]
-		provider := ExtractProvider(label)
+		provider := ExtractProvider(n.Attrs["label"])
 		if provider != "" && !providerSubgraphs[provider] {
-			sb.WriteString(fmt.Sprintf("\tsubgraph %s\n", caser.String(provider)))
+			sb.WriteString(fmt.Sprintf("\t\tsubgraph %s\n", caser.String(provider)))
 			providerSubgraphs[provider] = true
 		}
 	}
 
-	if subgraphName != "" {
-		sb.WriteString(fmt.Sprintf("\tsubgraph %s\n", subgraphName))
-	}
-
-	// Iterate over nodes to add them to the Mermaid graph
+	nodeMap := make(map[string]int)
 	for _, n := range graph.Nodes.Nodes {
 		label := CleanLabel(n.Attrs["label"])
 		nodeName := CleanID(n.Name)
-		if label != "" && nodeName != "" && !strings.Contains(label, "provider") {
-			sb.WriteString(fmt.Sprintf("\t\t%s[\"%s\"]\n", nodeName, label))
+		if label != "" && nodeName != "" {
+			nodeMap[label]++
+			count := nodeMap[label]
+			if count > 1 {
+				sb.WriteString(fmt.Sprintf("\t\t\t%s[\"%s\\nCount: %d\"]\n", nodeName, label, count))
+			} else {
+				sb.WriteString(fmt.Sprintf("\t\t\t%s[\"%s\"]\n", nodeName, label))
+			}
 		}
 	}
 
-	// Iterate over edges to add them to the Mermaid graph
+	for range providerSubgraphs {
+		sb.WriteString("\t\tend\n")
+	}
+
 	for _, edge := range graph.Edges.Edges {
 		srcLabel := CleanLabel(graph.Nodes.Lookup[edge.Src].Attrs["label"])
 		dstLabel := CleanLabel(graph.Nodes.Lookup[edge.Dst].Attrs["label"])
 		srcName := CleanID(edge.Src)
 		dstName := CleanID(edge.Dst)
-		if srcLabel != "" && dstLabel != "" && !strings.Contains(srcLabel, "provider") && !strings.Contains(dstLabel, "provider") {
+		if srcLabel != "" && dstLabel != "" {
 			sb.WriteString(fmt.Sprintf("\t\t%s --> %s\n", srcName, dstName))
 		}
 	}
 
-	// Close all open subgraphs
-	for range providerSubgraphs {
-		sb.WriteString("\tend\n")
-	}
 	if subgraphName != "" {
 		sb.WriteString("\tend\n")
 	}

main.go

@@ -5,11 +5,7 @@ import (
 	u "github.com/RoseSecurity/terramaid/pkg/utils"
 )
 
-var version string
-
 func main() {
-	cmd.Version = version
-
 	if err := cmd.Execute(); err != nil {
 		u.LogErrorAndExit(err)
 	}