From 3d24e730c8a2a87832def9136be7c80b10dde7fe Mon Sep 17 00:00:00 2001 From: aymene01 Date: Sat, 21 Sep 2024 13:19:58 +0200 Subject: [PATCH 1/2] refactor: remove global variables for detonate command --- v2/cmd/stratus/detonate_cmd.go | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/v2/cmd/stratus/detonate_cmd.go b/v2/cmd/stratus/detonate_cmd.go index a7ae93a06..34d04671e 100644 --- a/v2/cmd/stratus/detonate_cmd.go +++ b/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,16 +42,16 @@ 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) @@ -59,12 +59,12 @@ func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool) { // 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 { From 592fc45a6afcf50f77ef95eb2d77543c4be29050 Mon Sep 17 00:00:00 2001 From: aymene01 Date: Mon, 23 Sep 2024 09:59:16 +0200 Subject: [PATCH 2/2] refactor for other commands --- v2/cmd/stratus/cleanup_cmd.go | 55 ++++++++++++++++++----------------- v2/cmd/stratus/list_cmd.go | 17 +++++++---- v2/cmd/stratus/revert_cmd.go | 25 ++++++++-------- v2/cmd/stratus/warmup_cmd.go | 18 +++++++----- 4 files changed, 65 insertions(+), 50 deletions(-) diff --git a/v2/cmd/stratus/cleanup_cmd.go b/v2/cmd/stratus/cleanup_cmd.go index 15a7da9cc..16af4d2aa 100644 --- a/v2/cmd/stratus/cleanup_cmd.go +++ b/v2/cmd/stratus/cleanup_cmd.go @@ -2,17 +2,18 @@ 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"}, @@ -20,17 +21,17 @@ func buildCleanupCmd() *cobra.Command { 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 @@ -38,31 +39,33 @@ func buildCleanupCmd() *cobra.Command { 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) } diff --git a/v2/cmd/stratus/list_cmd.go b/v2/cmd/stratus/list_cmd.go index 64dfc4275..c01792b44 100644 --- a/v2/cmd/stratus/list_cmd.go +++ b/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,13 +27,16 @@ 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 { @@ -40,6 +44,7 @@ func doListCmd(mitreAttackTactic string, platform string) { } filter.Platform = platform } + if mitreAttackTactic != "" { tactic, err := mitreattack.AttackTacticFromString(mitreAttackTactic) if err != nil { @@ -47,7 +52,9 @@ func doListCmd(mitreAttackTactic string, platform string) { } filter.Tactic = tactic } + techniques := stratus.GetRegistry().GetAttackTechniques(&filter) + t := GetDisplayTable() t.AppendHeader(table.Row{"Technique ID", "Technique name", "Platform", "MITRE ATT&CK Tactic"}) diff --git a/v2/cmd/stratus/revert_cmd.go b/v2/cmd/stratus/revert_cmd.go index 1ea46e43e..398c45b4f 100644 --- a/v2/cmd/stratus/revert_cmd.go +++ b/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,14 +37,15 @@ 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) @@ -52,7 +53,7 @@ func doRevertCmd(techniques []*stratus.AttackTechnique) { // 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 } diff --git a/v2/cmd/stratus/warmup_cmd.go b/v2/cmd/stratus/warmup_cmd.go index 86d61fccb..8757ee45d 100644 --- a/v2/cmd/stratus/warmup_cmd.go +++ b/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()