refactor: remove global variables for all commands

v2/cmd/stratus/cleanup_cmd.go

@@ -2,67 +2,70 @@ package main
 
 import (
 	"errors"
+	"log"
+	"os"
+
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"github.com/spf13/cobra"
-	"log"
-	"os"
 )
 
-var flagForceCleanup bool
-var flagCleanupAll bool
-
 func buildCleanupCmd() *cobra.Command {
+	var flagForceCleanup bool
+	var flagCleanupAll bool
+
 	cleanupCmd := &cobra.Command{
 		Use:                   "cleanup [attack-technique-id]... | --all",
 		Aliases:               []string{"clean"},
 		Short:                 "Cleans up any leftover infrastructure or configuration from a TTP.",
 		Example:               "stratus cleanup aws.defense-evasion.cloudtrail-stop\nstratus cleanup --all",
 		DisableFlagsInUseLine: true,
 		Args: func(cmd *cobra.Command, args []string) error {
-			if len(args) == 0 && flagCleanupAll {
-				if !flagCleanupAll {
-					return errors.New("pass the ID of the technique to clean up, or --all")
-				}
-				return nil
+			if len(args) == 0 && !flagCleanupAll {
+				return errors.New("pass the ID of the technique to clean up, or --all")
 			}
 
-			// Ensure the technique IDs are valid
-			_, err := resolveTechniques(args)
+			if len(args) > 0 {
+				// Ensure the technique IDs are valid
+				_, err := resolveTechniques(args)
+				return err
+			}
 
-			return err
+			return nil
 		},
 		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 			return getTechniquesCompletion(toComplete), cobra.ShellCompDirectiveNoFileComp
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if len(args) > 0 {
 				techniques, _ := resolveTechniques(args)
-				doCleanupCmd(techniques)
+				doCleanupCmd(techniques, flagForceCleanup)
 				return nil
 			} else if flagCleanupAll {
-				// clean up all techniques that are not in the COLD state
-				doCleanupAllCmd()
+				doCleanupAllCmd(flagForceCleanup)
 				return nil
-			} else {
-				return errors.New("pass the ID of the technique to clean up, or --all")
 			}
+			return errors.New("pass the ID of the technique to clean up, or --all")
 		},
 	}
+
 	cleanupCmd.Flags().BoolVarP(&flagForceCleanup, "force", "f", false, "Force cleanup even if the technique is already COLD")
 	cleanupCmd.Flags().BoolVarP(&flagCleanupAll, "all", "", false, "Clean up all techniques that are not in COLD state")
+
 	return cleanupCmd
 }
 
-func doCleanupCmd(techniques []*stratus.AttackTechnique) {
+func doCleanupCmd(techniques []*stratus.AttackTechnique, forceCleanup bool) {
 	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
+
 	for i := 0; i < workerCount; i++ {
-		go cleanupCmdWorker(techniquesChan, errorsChan)
+		go cleanupCmdWorker(techniquesChan, errorsChan, forceCleanup)
 	}
-	for i := range techniques {
-		techniquesChan <- techniques[i]
+
+	for _, technique := range techniques {
+		techniquesChan <- technique
 	}
 	close(techniquesChan)
 
@@ -73,16 +76,16 @@ func doCleanupCmd(techniques []*stratus.AttackTechnique) {
 	}
 }
 
-func cleanupCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error) {
+func cleanupCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error, forceCleanup bool) {
 	for technique := range techniques {
-		stratusRunner := runner.NewRunner(technique, flagForceCleanup)
+		stratusRunner := runner.NewRunner(technique, forceCleanup)
 		err := stratusRunner.CleanUp()
 		errors <- err
 	}
 }
 
-func doCleanupAllCmd() {
+func doCleanupAllCmd(forceCleanup bool) {
 	log.Println("Cleaning up all techniques that have been warmed-up or detonated")
 	availableTechniques := stratus.GetRegistry().ListAttackTechniques()
-	doCleanupCmd(availableTechniques)
+	doCleanupCmd(availableTechniques, forceCleanup)
 }

v2/cmd/stratus/detonate_cmd.go

@@ -2,19 +2,19 @@ package main
 
 import (
 	"errors"
-	"github.com/datadog/stratus-red-team/v2/internal/utils"
-	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
-	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"os"
 	"strings"
 
+	"github.com/datadog/stratus-red-team/v2/internal/utils"
+	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
+	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"github.com/spf13/cobra"
 )
 
-var detonateForce bool
-var detonateCleanup bool
-
 func buildDetonateCmd() *cobra.Command {
+	var detonateForce bool
+	var detonateCleanup bool
+
 	detonateCmd := &cobra.Command{
 		Use:   "detonate attack-technique-id [attack-technique-id]...",
 		Short: "Detonate one or multiple attack techniques",
@@ -42,29 +42,29 @@ func buildDetonateCmd() *cobra.Command {
 		},
 		Run: func(cmd *cobra.Command, args []string) {
 			techniques, _ := resolveTechniques(args)
-			doDetonateCmd(techniques, detonateCleanup)
+			doDetonateCmd(techniques, detonateCleanup, detonateForce)
 		},
 	}
 	detonateCmd.Flags().BoolVarP(&detonateCleanup, "cleanup", "", false, "Clean up the infrastructure that was spun up as part of the technique prerequisites")
-	//detonateCmd.Flags().BoolVarP(&detonateNoWarmup, "no-warmup", "", false, "Do not spin up prerequisite infrastructure or configuration. Requires that 'warmup' was used before.")
 	detonateCmd.Flags().BoolVarP(&detonateForce, "force", "f", false, "Force detonation in cases where the technique is not idempotent and has already been detonated")
 
 	return detonateCmd
 }
-func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool) {
+
+func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool, force bool) {
 	VerifyPlatformRequirements(techniques)
 	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 
 	// Create workers
 	for i := 0; i < workerCount; i++ {
-		go detonateCmdWorker(techniquesChan, errorsChan)
+		go detonateCmdWorker(techniquesChan, errorsChan, cleanup, force)
 	}
 
 	// Send attack techniques to detonate
-	for i := range techniques {
-		techniquesChan <- techniques[i]
+	for _, technique := range techniques {
+		techniquesChan <- technique
 	}
 	close(techniquesChan)
 
@@ -73,11 +73,11 @@ func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool) {
 	}
 }
 
-func detonateCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error) {
+func detonateCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error, cleanup bool, force bool) {
 	for technique := range techniques {
-		stratusRunner := runner.NewRunner(technique, detonateForce)
+		stratusRunner := runner.NewRunner(technique, force)
 		detonateErr := stratusRunner.Detonate()
-		if detonateCleanup {
+		if cleanup {
 			cleanupErr := stratusRunner.CleanUp()
 			errors <- utils.CoalesceErr(detonateErr, cleanupErr)
 		} else {

v2/cmd/stratus/list_cmd.go

@@ -2,19 +2,20 @@ package main
 
 import (
 	"fmt"
+	"log"
+	"strings"
+
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus/mitreattack"
 	"github.com/fatih/color"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/spf13/cobra"
-	"log"
-	"strings"
 )
 
-var listPlatform string
-var listMitreAttackTactic string
-
 func buildListCmd() *cobra.Command {
+	var listPlatform string
+	var listMitreAttackTactic string
+
 	listCmd := &cobra.Command{
 		Use:   "list",
 		Short: "List attack techniques",
@@ -26,28 +27,34 @@ func buildListCmd() *cobra.Command {
 			doListCmd(listMitreAttackTactic, listPlatform)
 		},
 	}
+
 	listCmd.Flags().StringVarP(&listPlatform, "platform", "", "", "Filter on specific platform")
 	listCmd.Flags().StringVarP(&listMitreAttackTactic, "mitre-attack-tactic", "", "", "Filter on a specific MITRE ATT&CK tactic.")
+
 	return listCmd
 }
 
 func doListCmd(mitreAttackTactic string, platform string) {
 	filter := stratus.AttackTechniqueFilter{}
+
 	if platform != "" {
 		platform, err := stratus.PlatformFromString(platform)
 		if err != nil {
 			log.Fatal(err)
 		}
 		filter.Platform = platform
 	}
+
 	if mitreAttackTactic != "" {
 		tactic, err := mitreattack.AttackTacticFromString(mitreAttackTactic)
 		if err != nil {
 			log.Fatal(err)
 		}
 		filter.Tactic = tactic
 	}
+
 	techniques := stratus.GetRegistry().GetAttackTechniques(&filter)
+
 	t := GetDisplayTable()
 	t.AppendHeader(table.Row{"Technique ID", "Technique name", "Platform", "MITRE ATT&CK Tactic"})
 

v2/cmd/stratus/revert_cmd.go

@@ -2,18 +2,18 @@ package main
 
 import (
 	"errors"
-	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
-	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"log"
 	"os"
 
+	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
+	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"github.com/spf13/cobra"
 )
 
-var revertForce bool
-
 func buildRevertCmd() *cobra.Command {
-	detonateCmd := &cobra.Command{
+	var revertForce bool
+
+	revertCmd := &cobra.Command{
 		Use:                   "revert attack-technique-id [attack-technique-id]...",
 		Short:                 "Revert the detonation of an attack technique",
 		Example:               "stratus revert aws.defense-evasion.cloudtrail-stop",
@@ -37,22 +37,23 @@ func buildRevertCmd() *cobra.Command {
 		},
 		Run: func(cmd *cobra.Command, args []string) {
 			techniques, _ := resolveTechniques(args)
-			doRevertCmd(techniques)
+			doRevertCmd(techniques, revertForce)
 		},
 	}
-	detonateCmd.Flags().BoolVarP(&revertForce, "force", "f", false, "Force attempt to reverting even if the technique is not in the DETONATED state")
-	return detonateCmd
+
+	revertCmd.Flags().BoolVarP(&revertForce, "force", "f", false, "Force attempt to revert even if the technique is not in the DETONATED state")
+	return revertCmd
 }
 
-func doRevertCmd(techniques []*stratus.AttackTechnique) {
+func doRevertCmd(techniques []*stratus.AttackTechnique, force bool) {
 	VerifyPlatformRequirements(techniques)
 	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 
 	// Create workers
 	for i := 0; i < workerCount; i++ {
-		go revertCmdWorker(techniquesChan, errorsChan)
+		go revertCmdWorker(techniquesChan, errorsChan, force)
 	}
 
 	// Send attack techniques to revert
@@ -68,14 +69,14 @@ func doRevertCmd(techniques []*stratus.AttackTechnique) {
 	}
 }
 
-func revertCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error) {
+func revertCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error, force bool) {
 	for technique := range techniques {
 		if technique.Revert == nil {
 			log.Println("Warning: " + technique.ID + " has no revert function and cannot be reverted.")
 			errors <- nil
 			continue
 		}
-		stratusRunner := runner.NewRunner(technique, revertForce)
+		stratusRunner := runner.NewRunner(technique, force)
 		err := stratusRunner.Revert()
 		errors <- err
 	}

v2/cmd/stratus/warmup_cmd.go

@@ -2,15 +2,16 @@ package main
 
 import (
 	"errors"
+	"os"
+
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner"
 	"github.com/spf13/cobra"
-	"os"
 )
 
-var forceWarmup bool
-
 func buildWarmupCmd() *cobra.Command {
+	var forceWarmup bool
+
 	warmupCmd := &cobra.Command{
 		Use:                   "warmup attack-technique-id [attack-technique-id]...",
 		Short:                 "\"Warm up\" an attack technique by spinning up the prerequisite infrastructure or configuration, without detonating it",
@@ -35,21 +36,24 @@ func buildWarmupCmd() *cobra.Command {
 		},
 		Run: func(cmd *cobra.Command, args []string) {
 			techniques, _ := resolveTechniques(args)
-			doWarmupCmd(techniques)
+			doWarmupCmd(techniques, forceWarmup)
 		},
 	}
+
 	warmupCmd.Flags().BoolVarP(&forceWarmup, "force", "f", false, "Force re-ensuring the prerequisite infrastructure or configuration is up to date")
 	return warmupCmd
 }
 
-func doWarmupCmd(techniques []*stratus.AttackTechnique) {
+func doWarmupCmd(techniques []*stratus.AttackTechnique, forceWarmup bool) {
 	VerifyPlatformRequirements(techniques)
 	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
+
 	for i := 0; i < workerCount; i++ {
-		go warmupCmdWorker(techniquesChan, errorsChan)
+		go warmupCmdWorker(techniquesChan, errorsChan, forceWarmup)
 	}
+
 	for i := range techniques {
 		techniquesChan <- techniques[i]
 	}
@@ -60,7 +64,7 @@ func doWarmupCmd(techniques []*stratus.AttackTechnique) {
 	}
 }
 
-func warmupCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error) {
+func warmupCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error, forceWarmup bool) {
 	for technique := range techniques {
 		stratusRunner := runner.NewRunner(technique, forceWarmup)
 		_, err := stratusRunner.WarmUp()