refactor: use cobra for cmd flags (#18)

* refactor: use cobra for cmd flags

* fix: lint

* fix: README

Author: notdodo

.github/workflows/golangci.yml

@@ -35,7 +35,7 @@ jobs:
       - uses: actions/setup-go@0c52d547c9bc32b1aa3301fd7a9cb496313a4491
         # v5.0.0
         with:
-          go-version-file: "go.mod"
+          go-version: "oldstable"
       - name: golangci-lint
         uses: golangci/golangci-lint-action@82d40c283aeb1f2b6595839195e95c2d6a49081b
         # v5.0.0
@@ -77,6 +77,8 @@ jobs:
         # v5.0.0
         with:
           go-version-file: "go.mod"
+          cache-dependency-path: |
+            go.sum
       - name: Install dependencies
         run: go get .
       - name: Build

.vscode/extensions.json

@@ -3,6 +3,7 @@ include .env
 BUILD_FLAGS=-ldflags="-s -w"
 OUT_DIR=dist
 OUT_PREFIX=nuvola
+NEO4J_PASS ?= neo4j
 
 define ANNOUNCE_BODY
 CALL apoc.export.cypher.all("/var/lib/neo4j/import/all.cypher", {
@@ -13,12 +14,23 @@ YIELD file, batches, source, format, nodes, relationships, properties, time, row
 RETURN file, batches, source, format, nodes, relationships, properties, time, rows, batchSize;
 endef
 
+.PHONY: backup
+export ANNOUNCE_BODY
+backup:
+	@echo "$$ANNOUNCE_BODY" | docker-compose exec -T neo4j cypher-shell -u neo4j -p ${PASSWORD} -d nuvoladb --non-interactive
+	docker-compose exec -T neo4j cp /var/lib/neo4j/import/all.cypher /backup
 
-start:
-	@if [ ! -f ./.env ]; then\
-	  cp .env_example .env;\
+.PHONY: check-dependencies
+check-dependencies:
+	@command -v docker compose >/dev/null 2>&1 || { echo >&2 "docker compose not installed"; exit 1; }
+	@command -v go >/dev/null 2>&1 || { echo >&2 "Go not installed"; exit 1; }
+
+.PHONY: check-neo4j-password
+check-neo4j-password:
+	@if [ -z "${NEO4J_PASS}" ]; then \
+		echo "NEO4J_PASS is not set. Set it in the .env file"; \
+		exit 1; \
 	fi
-	docker-compose up -d
 
 build:
 	go build ${BUILD_FLAGS} -o ${OUT_PREFIX}
@@ -30,23 +42,28 @@ compile: build
 	GOOS=darwin GOARCH=amd64 go build ${BUILD_FLAGS} -o ${OUT_DIR}/${OUT_PREFIX}-darwin-amd64
 	GOOS=darwin GOARCH=arm64 go build ${BUILD_FLAGS} -o ${OUT_DIR}/${OUT_PREFIX}-darwin-arm64
 
-clean:
+.PHONY: clean
+clean: stop-containers
 	@rm -rf ./${OUT_DIR}
 	@rm -f ./${OUT_PREFIX}
-	@docker-compose stop
-	@docker-compose down -v
-	@docker-compose rm -fv
-
-.PHONY: backup
-export ANNOUNCE_BODY
-backup:
-	@echo "$$ANNOUNCE_BODY" | docker-compose exec -T neo4j cypher-shell -u neo4j -p ${PASSWORD} -d nuvoladb --non-interactive
-	docker-compose exec -T neo4j cp /var/lib/neo4j/import/all.cypher /backup
 
 .PHONY: restore
 restore:
 	@cat ./backup/all.cypher | docker-compose exec -T neo4j cypher-shell -u neo4j -p ${PASSWORD} -d nuvoladb --non-interactive
 
+.PHONY: start-containers
+start-containers: check-dependencies
+	@if [ ! -f ./.env ]; then\
+	  cp .env_example .env;\
+	fi
+	@docker compose up -d
+
+.PHONY: stop-containers
+stop-containers:
+	@docker compose stop
+	@docker compose down -v
+	@docker compose rm -fv
+
 .PHONY: tests
-tests: start build
-	$(MAKE) -C tests all
+tests: start-containers build
+	cd ./assets/ && $(MAKE) -C tests all

.vscode/launch.json

@@ -61,7 +61,7 @@ make build
 ./nuvola assess -import ~/DumpDumpFolder/nuvola-default_RO_20220901.zip
 ```
 
-3. To only perform static assessments on the data loaded into the Neo4j database using the [predefined ruleset](https://github.com/primait/nuvola/tree/master/assess/rules):
+3. To only perform static assessments on the data loaded into the Neo4j database using the [predefined ruleset](https://github.com/primait/nuvola/tree/master/assets/rules):
 
 ```bash
 ./nuvola assess

Makefile

@@ -1,4 +1,4 @@
-package assess
+package cmd
 
 import (
 	"archive/zip"
@@ -8,16 +8,38 @@ import (
 	"log"
 	"strings"
 
-	connector "github.com/primait/nuvola/connector"
-
-	clioutput "github.com/primait/nuvola/tools/cli/output"
-	nuvolaerror "github.com/primait/nuvola/tools/error"
+	"github.com/primait/nuvola/connector"
+	"github.com/primait/nuvola/pkg/io/logging"
 	"github.com/primait/nuvola/tools/filesystem/files"
 	unzip "github.com/primait/nuvola/tools/filesystem/zip"
 	"github.com/primait/nuvola/tools/yamler"
+	"github.com/spf13/cobra"
 )
 
-func ImportZipFile(connector *connector.StorageConnector, zipfile string) {
+var assessCmd = &cobra.Command{
+	Use:   "assess",
+	Short: "Execute assessment queries against data loaded in Neo4J",
+	Run: func(cmd *cobra.Command, args []string) {
+		if cmd.Flags().Changed(flagVerbose) {
+			logger.SetVerboseLevel()
+		}
+		if cmd.Flags().Changed(flagDebug) {
+			logger.SetDebugLevel()
+		}
+
+		connector.SetActions()
+		storageConnector := connector.NewStorageConnector()
+		if importFile != "" && !noImport {
+			logger.Info("Flushing database")
+			logger.Info(fmt.Sprintf("Importing %s", importFile))
+			importZipFile(storageConnector, importFile)
+		}
+
+		assess(storageConnector, "./assets/rules/")
+	},
+}
+
+func importZipFile(connector *connector.StorageConnector, zipfile string) {
 	connector.FlushAll()
 	var ordering = []string{
 		"Groups",
@@ -47,7 +69,7 @@ func ImportZipFile(connector *connector.StorageConnector, zipfile string) {
 	for _, f := range orderedFiles {
 		rc, err := f.Open()
 		if err != nil {
-			nuvolaerror.HandleError(err, "Assess", "Opening content of ZIP")
+			logging.HandleError(err, "Assess", "Opening content of ZIP")
 		}
 		defer func() {
 			if err := rc.Close(); err != nil {
@@ -58,27 +80,27 @@ func ImportZipFile(connector *connector.StorageConnector, zipfile string) {
 		buf := new(bytes.Buffer)
 		_, err = io.Copy(buf, rc) // #nosecG110
 		if err != nil {
-			nuvolaerror.HandleError(err, "Assess", "Copying buffer from ZIP")
+			logging.HandleError(err, "Assess", "Copying buffer from ZIP")
 		}
 		connector.ImportResults(f.Name, buf.Bytes())
 	}
 }
 
-func Assess(connector *connector.StorageConnector, rulesPath string) {
+func assess(connector *connector.StorageConnector, rulesPath string) {
 	// perform checks based on pre-defined static rules
 	for _, rule := range files.GetFiles(rulesPath, ".ya?ml") {
 		var c = yamler.GetConf(rule)
 		if c.Enabled {
 			query, args := yamler.PrepareQuery(c)
 			results := connector.Query(query, args)
 
-			clioutput.PrintRed("Running rule: " + rule)
-			clioutput.PrintGreen("Name: " + c.Name)
-			clioutput.PrintGreen("Arguments:")
-			clioutput.PrintDarkGreen(yamler.ArgsToQueryNeo4jBrowser(args))
-			clioutput.PrintGreen("Query:")
-			clioutput.PrintDarkGreen(query)
-			clioutput.PrintGreen("Description: " + c.Description)
+			logging.PrintRed("Running rule: " + rule)
+			logging.PrintGreen("Name: " + c.Name)
+			logging.PrintGreen("Arguments:")
+			logging.PrintDarkGreen(yamler.ArgsToQueryNeo4jBrowser(args))
+			logging.PrintGreen("Query:")
+			logging.PrintDarkGreen(query)
+			logging.PrintGreen("Description: " + c.Description)
 
 			for _, resultMap := range results {
 				for key, value := range resultMap {
@@ -98,3 +120,7 @@ func Assess(connector *connector.StorageConnector, rulesPath string) {
 		}
 	}
 }
+
+func init() {
+	rootCmd.AddCommand(assessCmd)
+}

README.md

@@ -1,4 +1,4 @@
-package dump
+package cmd
 
 import (
 	"encoding/json"
@@ -7,11 +7,10 @@ import (
 	"time"
 
 	"github.com/primait/nuvola/connector"
-
-	clioutput "github.com/primait/nuvola/tools/cli/output"
-	nuvolaerror "github.com/primait/nuvola/tools/error"
+	"github.com/primait/nuvola/pkg/io/logging"
 	"github.com/primait/nuvola/tools/filesystem/files"
 	"github.com/primait/nuvola/tools/filesystem/zip"
+	"github.com/spf13/cobra"
 )
 
 var AWSResults = map[string]interface{}{
@@ -29,7 +28,39 @@ var AWSResults = map[string]interface{}{
 	"RedshiftDBs":      nil,
 }
 
-func DumpData(storageConnector *connector.StorageConnector, cloudConnector *connector.CloudConnector) map[string]interface{} {
+var dumpCmd = &cobra.Command{
+	Use:   "dump",
+	Short: "Dump AWS resources and policies information and store them in Neo4j",
+	Run: func(cmd *cobra.Command, args []string) {
+		startTime := time.Now()
+		markAsRequired("aws-profile")
+		if err := rootCmd.ValidateRequiredFlags(); err != nil {
+			logger.Error("Required flags not provided", "err", err)
+		}
+		if cmd.Flags().Changed(flagVerbose) {
+			logger.SetVerboseLevel()
+		}
+		if cmd.Flags().Changed(flagDebug) {
+			logger.SetDebugLevel()
+		}
+
+		cloudConnector, err := connector.NewCloudConnector(awsProfile, awsEndpointUrl)
+		if err != nil {
+			logger.Error(err.Error())
+		}
+
+		if dumpOnly {
+			dumpData(nil, cloudConnector)
+		} else {
+			storageConnector := connector.NewStorageConnector().FlushAll()
+			dumpData(storageConnector, cloudConnector)
+		}
+		saveResults(awsProfile, outputDirectory, outputFormat)
+		logger.Info("Execution Time", "seconds", time.Since(startTime))
+	},
+}
+
+func dumpData(storageConnector *connector.StorageConnector, cloudConnector *connector.CloudConnector) map[string]interface{} {
 	dataChan := make(chan map[string]interface{})
 	go cloudConnector.DumpAll("aws", dataChan)
 	for {
@@ -41,15 +72,15 @@ func DumpData(storageConnector *connector.StorageConnector, cloudConnector *conn
 		mapKey := v.MapKeys()[0].Interface().(string)
 		obj, err := json.Marshal(a[mapKey])
 		if err != nil {
-			nuvolaerror.HandleError(err, "DumpData", "Marshalling output")
+			logger.Error("DumpData: error marshalling output", err)
 		}
 		storageConnector.ImportResults(mapKey, obj)
 		AWSResults[mapKey] = a[mapKey]
 	}
 	return AWSResults
 }
 
-func SaveResults(awsProfile string, outputDir string, outputFormat string) {
+func saveResults(awsProfile string, outputDir string, outputFormat string) {
 	if awsProfile == "" {
 		awsProfile = "default"
 	}
@@ -59,11 +90,14 @@ func SaveResults(awsProfile string, outputDir string, outputFormat string) {
 
 	today := time.Now().Format("20060102")
 	for key, value := range AWSResults {
-		clioutput.PrintGreen(key + ":")
-		fmt.Printf("%s\n", clioutput.PrettyJSON(value))
+		logger.Info(key, logging.PrettyJSON(value))
 
 		if outputFormat == "json" {
 			files.PrettyJSONToFile(outputDir, fmt.Sprintf("%s_%s.json", key, today), value)
 		}
 	}
 }
+
+func init() {
+	rootCmd.AddCommand(dumpCmd)
+}

controller/assess/rules/CloudFormation-privesc.yaml assets/rules/CloudFormation-privesc.yaml

@@ -0,0 +1,58 @@
+package cmd
+
+import (
+	"github.com/primait/nuvola/pkg/io/logging"
+	"github.com/spf13/cobra"
+)
+
+const (
+	flagVerbose         = "verbose"
+	flagDebug           = "debug"
+	flagAWSProfile      = "aws-profile"
+	flagAWSEndpointUrl  = "aws-endpoint-url"
+	flagOutputDirectory = "output-dir"
+	flagOutputFormat    = "output-format"
+	flagDumpOnly        = "dump-only"
+	flagImportFile      = "import"
+	flagNoImport        = "no-import"
+)
+
+var (
+	logger          logging.LogManager
+	awsProfile      string
+	awsEndpointUrl  string
+	outputDirectory string
+	outputFormat    string
+	dumpOnly        bool
+	importFile      string
+	noImport        bool
+	rootCmd         = &cobra.Command{
+		Use:   "nuvola",
+		Short: "A tool to dump and perform automatic and manual security analysis on AWS",
+	}
+)
+
+func init() {
+	logger = logging.GetLogManager()
+	rootCmd.PersistentFlags().BoolP(flagVerbose, "v", false, "Verbose output")
+	rootCmd.PersistentFlags().BoolP(flagDebug, "d", false, "Debug output")
+	dumpCmd.PersistentFlags().StringVarP(&awsProfile, flagAWSProfile, "p", "default", "AWS Profile to use")
+	dumpCmd.PersistentFlags().StringVarP(&outputDirectory, flagOutputDirectory, "o", "", "Output folder where the files will be saved (default: \".\")")
+	dumpCmd.PersistentFlags().StringVarP(&outputFormat, flagOutputFormat, "f", "zip", "Output format: ZIP or json files")
+	dumpCmd.PersistentFlags().BoolVarP(&dumpOnly, flagDumpOnly, "", false, "Flag to prevent loading data into Neo4j (default: \"false\")")
+
+	assessCmd.PersistentFlags().StringVarP(&importFile, flagImportFile, "i", "", "Input ZIP file to load")
+	assessCmd.PersistentFlags().BoolVarP(&noImport, flagNoImport, "", false, "Use stored data from Neo4j without import (default)")
+}
+
+func Execute() {
+	if err := rootCmd.Execute(); err != nil {
+		logger.Error("Error executing command", "err", err)
+	}
+}
+
+func markAsRequired(flag string) {
+	if err := rootCmd.MarkFlagRequired(flag); err != nil {
+		logger.Error("Required flags not provided", "err", err, "flag", flag)
+	}
+}