fix: bug on missing init logger; fix goroutines (#34)

primait · Jun 9, 2024 · 8d80327 · 8d80327
1 parent ad4f1db
commit 8d80327
Show file tree

Hide file tree

Showing 8 changed files with 62 additions and 127 deletions.
diff --git a/cmd/dump.go b/cmd/dump.go
@@ -3,6 +3,7 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"sync"
 	"time"
 
@@ -64,30 +65,40 @@ func runDumpCmd(cmd *cobra.Command, args []string) {
 func dumpData(storageConnector *connector.StorageConnector, cloudConnector *connector.CloudConnector) {
 	dataChan := make(chan map[string]interface{})
 	var wg sync.WaitGroup
-	wg.Add(1)
+
 	go func() {
-		cloudConnector.DumpAll("aws", dataChan, &wg)
 		defer close(dataChan)
+		cloudConnector.DumpAll("aws", dataChan)
 	}()
 
-	for data := range dataChan {
-		processData(data, storageConnector)
-	}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for data := range dataChan {
+			processData(storageConnector, data)
+		}
+	}()
 	wg.Wait()
 }
 
-func processData(data map[string]interface{}, storageConnector *connector.StorageConnector) {
-	for key, value := range data {
-		obj, err := json.Marshal(value)
-		if err != nil {
-			logger.Error("Error marshalling output", "err", err)
-			continue
-		}
-		if storageConnector != nil {
-			storageConnector.ImportResults(key, obj)
-		}
-		AWSResults[key] = value
+func processData(storageConnector *connector.StorageConnector, data map[string]interface{}) {
+	if len(data) == 0 {
+		return
+	}
+
+	v := reflect.ValueOf(data)
+	if v.Kind() != reflect.Map || v.Len() == 0 {
+		logger.Error("processData: unexpected data format")
 	}
+
+	mapKey := v.MapKeys()[0].Interface().(string)
+	obj, err := json.Marshal(data[mapKey])
+	if err != nil {
+		logger.Error("processData: error marshalling output", "err", err)
+	}
+
+	storageConnector.ImportResults(mapKey, obj)
+	AWSResults[mapKey] = data[mapKey]
 }
 
 func saveResults(awsProfile, outputDir, outputFormat string) {

diff --git a/go.mod b/go.mod
@@ -19,7 +19,6 @@ require (
 	github.com/gocarina/gocsv v0.0.0-20240520201108-78e41c74b4b1
 	github.com/imroc/req/v3 v3.43.7
 	github.com/itchyny/gojq v0.12.16
-	github.com/joho/godotenv v1.5.1
 	github.com/neo4j/neo4j-go-driver/v5 v5.20.0
 	github.com/notdodo/arner v0.0.0-20230222134658-4fe417a6cc9c
 	github.com/notdodo/goflat/v2 v2.1.1
@@ -69,7 +68,6 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
-	github.com/muesli/reflow v0.3.0 // indirect
 	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/onsi/ginkgo/v2 v2.19.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect

diff --git a/go.sum b/go.sum
diff --git a/pkg/connector/cloud_connector.go b/pkg/connector/cloud_connector.go
@@ -3,7 +3,6 @@ package connector
 import (
 	"errors"
 	"strings"
-	"sync"
 
 	awsconfig "github.com/primait/nuvola/pkg/connector/services/aws"
 	"github.com/primait/nuvola/pkg/io/logging"
@@ -24,34 +23,37 @@ func SetActions() {
 	awsconfig.SetActions()
 }
 
-func (cc *CloudConnector) DumpAll(cloudprovider string, c chan map[string]interface{}, wg *sync.WaitGroup) {
+func (cc *CloudConnector) DumpAll(cloudprovider string, c chan map[string]interface{}) {
 	switch strings.ToLower(cloudprovider) {
 	case "aws":
-		cc.dumpAWSData(c, wg)
-	default:
-		cc.logger.Error("Unsupported cloud provider", "cloudprovider", cloudprovider)
-	}
-}
+		dumpFunctions := []struct {
+			name string
+			dump func() interface{}
+		}{
+			{"Whoami", cc.AWSConfig.DumpWhoami},
+			{"CredentialReport", cc.AWSConfig.DumpCredentialReport},
+			{"Groups", cc.AWSConfig.DumpIAMGroups},
+			{"Users", cc.AWSConfig.DumpIAMUsers},
+			{"Roles", cc.AWSConfig.DumpIAMRoles},
+			{"Buckets", cc.AWSConfig.DumpBuckets},
+			{"EC2s", cc.AWSConfig.DumpEC2Instances},
+			{"VPCs", cc.AWSConfig.DumpVpcs},
+			{"Lambdas", cc.AWSConfig.DumpLambdas},
+			{"RDS", cc.AWSConfig.DumpRDS},
+			{"DynamoDBs", cc.AWSConfig.DumpDynamoDBs},
+			{"RedshiftDBs", cc.AWSConfig.DumpRedshiftDBs},
+		}
 
-func (cc *CloudConnector) dumpAWSData(c chan map[string]interface{}, wg *sync.WaitGroup) {
-	defer wg.Done()
-	data := map[string]interface{}{
-		"Whoami":           cc.AWSConfig.DumpWhoami(),
-		"CredentialReport": cc.AWSConfig.DumpCredentialReport(),
-		"Groups":           cc.AWSConfig.DumpIAMGroups(),
-		"Users":            cc.AWSConfig.DumpIAMUsers(),
-		"Roles":            cc.AWSConfig.DumpIAMRoles(),
-		"Buckets":          cc.AWSConfig.DumpBuckets(),
-		"EC2s":             cc.AWSConfig.DumpEC2Instances(),
-		"VPCs":             cc.AWSConfig.DumpVpcs(),
-		"Lambdas":          cc.AWSConfig.DumpLambdas(),
-		"RDS":              cc.AWSConfig.DumpRDS(),
-		"DynamoDBs":        cc.AWSConfig.DumpDynamoDBs(),
-		"RedshiftDBs":      cc.AWSConfig.DumpRedshiftDBs(),
-	}
-
-	for key, value := range data {
-		c <- map[string]interface{}{key: value}
+		for _, df := range dumpFunctions {
+			data := df.dump()
+			if data != nil {
+				c <- map[string]interface{}{
+					df.name: data,
+				}
+			}
+		}
+	default:
+		cc.logger.Error("unsupported cloud provider", "cloudprovider", cloudprovider)
 	}
 }
 

diff --git a/pkg/connector/services/aws/database/rds.go b/pkg/connector/services/aws/database/rds.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 
 	"github.com/primait/nuvola/pkg/connector/services/aws/ec2"
+	"github.com/primait/nuvola/pkg/io/logging"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
@@ -14,7 +15,7 @@ import (
 
 // aws iam list-users
 func ListRDS(cfg aws.Config) (rdsRet *RDS, re *awshttp.ResponseError) {
-	var rdsClient = RDSClient{Config: cfg}
+	var rdsClient = RDSClient{Config: cfg, logger: logging.GetLogManager()}
 
 	rdsRet = &RDS{}
 

diff --git a/pkg/connector/services/aws/database/redshift.go b/pkg/connector/services/aws/database/redshift.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 
 	"github.com/primait/nuvola/pkg/connector/services/aws/ec2"
+	"github.com/primait/nuvola/pkg/io/logging"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/redshift"
@@ -13,7 +14,7 @@ import (
 
 // aws iam list-users
 func ListRedshiftDBs(cfg aws.Config) (redshiftDBs []*RedshiftDB) {
-	var redshiftClient = RedshiftClient{Config: cfg}
+	var redshiftClient = RedshiftClient{Config: cfg, logger: logging.GetLogManager()}
 
 	for i := range ec2.Regions {
 		cfg.Region = ec2.Regions[i]

diff --git a/pkg/connector/services/aws/ec2/vpcs.go b/pkg/connector/services/aws/ec2/vpcs.go
@@ -6,10 +6,11 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/primait/nuvola/pkg/io/logging"
 )
 
 func ListVpcs(cfg aws.Config) (vpcs *VPC) {
-	ec2Client := EC2Client{Config: cfg, client: ec2.NewFromConfig(cfg)}
+	ec2Client := EC2Client{Config: cfg, client: ec2.NewFromConfig(cfg), logger: logging.GetLogManager()}
 
 	vpcs = &VPC{}
 	for _, region := range Regions {

diff --git a/pkg/connector/services/aws/iam/policies.go b/pkg/connector/services/aws/iam/policies.go
@@ -202,7 +202,6 @@ func (ic *IAMClient) listAttachedPolicies(identity string, object string) (attac
 		ic.logger.Warn("no user/role/group defined", "object", object)
 	}
 
-	fmt.Println("Asdfasd")
 	for _, policy := range output {
 		policyVersions := ic.listPolicyVersions(policy.PolicyArn)
 		policyDocument, errj := json.Marshal(policyVersions[0].Document)