From f312583248e48ca658a679fcb656d9668561fd6e Mon Sep 17 00:00:00 2001
From: Andre Ziviani <andrepziviani@gmail.com>
Date: Tue, 18 May 2021 16:40:53 -0300
Subject: [PATCH] feat/tgroutes: group by destination

---
 internal/chart/tgroutes.go     | 42 ++++++++++++++++++++++++++--------
 internal/vpc/transitgateway.go |  2 +-
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/internal/chart/tgroutes.go b/internal/chart/tgroutes.go
index 03a2f5f..56f3bc9 100644
--- a/internal/chart/tgroutes.go
+++ b/internal/chart/tgroutes.go
@@ -1,10 +1,13 @@
 package chart
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
+	"net"
 	"os"
+	"sort"
 	"sync"
 
 	"github.com/AndreZiviani/aws-fuzzy/internal/common"
@@ -40,11 +43,18 @@ func getTGAttachmentsRoutine(ctx context.Context, ec2client *ec2.Client, i int,
 	}
 }
 
+type TreeDataIp struct {
+	Name     string
+	Children []*net.IPNet
+}
+
 func processTables(ctx context.Context, ec2client *ec2.Client, tables []*vpc.DescribeTransitGatewayRouteTablesOutput) ([]*opts.TreeData, error) {
 	tableNode := make([]*opts.TreeData, 0)
 
 	for _, table := range tables {
-		ipRanges := make([]*opts.TreeData, len(table.Routes))
+		var ipRanges []*opts.TreeData
+		ipRanges = make([]*opts.TreeData, 0) // cant know beforehand how many destinations
+		nodeMapIp := make(map[string]*TreeDataIp)
 
 		c := make(chan *tgAttachmentRoutine)
 		var wg sync.WaitGroup
@@ -79,19 +89,31 @@ func processTables(ctx context.Context, ec2client *ec2.Client, tables []*vpc.Des
 				aws.ToString(attachments[0].ResourceId),
 			)
 
-			destination := []*opts.TreeData{
-				{
-					Name: name,
-				},
+			resourceId := aws.ToString(attachments[0].ResourceId)
+			node, ok := nodeMapIp[resourceId]
+			if !ok {
+				nodeMapIp[resourceId] = &TreeDataIp{
+					Name:     name,
+					Children: make([]*net.IPNet, 0),
+				}
+				node = nodeMapIp[resourceId]
 			}
+			_, netsubnet, _ := net.ParseCIDR(subnet)
+			node.Children = append(node.Children, netsubnet)
 
-			ipNode := &opts.TreeData{
-				Name:     subnet,
-				Children: destination,
-			}
+		}
 
-			ipRanges[attachmentsRoutine.Index] = ipNode
+		for _, v := range nodeMapIp {
+			sort.Slice(v.Children, func(i, j int) bool {
+				return bytes.Compare(v.Children[i].IP, v.Children[j].IP) < 0
+			})
+			tmp := make([]*opts.TreeData, 0)
+			for _, i := range v.Children {
+				tmp = append(tmp, &opts.TreeData{Name: i.String()})
+			}
+			ipRanges = append(ipRanges, &opts.TreeData{Name: v.Name, Children: tmp})
 		}
+
 		tmp := &opts.TreeData{
 			Name:     table.Name,
 			Children: ipRanges,
diff --git a/internal/vpc/transitgateway.go b/internal/vpc/transitgateway.go
index 393c285..fe7f7bb 100644
--- a/internal/vpc/transitgateway.go
+++ b/internal/vpc/transitgateway.go
@@ -175,7 +175,7 @@ func DescribeTransitGatewayRegistrationsFromARN(ctx context.Context, transitGate
 }
 
 func DescribeTransitGatewayRouteTables(ctx context.Context, ec2client *ec2.Client) ([]*DescribeTransitGatewayRouteTablesOutput, error) {
-	span, ctx := opentracing.StartSpanFromContext(ctx, "describetransitgateway")
+	span, ctx := opentracing.StartSpanFromContext(ctx, "describe transitgateway route tables")
 	defer span.Finish()
 
 	tmp, err := ec2client.DescribeTransitGatewayRouteTables(ctx, &ec2.DescribeTransitGatewayRouteTablesInput{})