-
Notifications
You must be signed in to change notification settings - Fork 197
/
detonate_cmd.go
87 lines (78 loc) · 2.84 KB
/
detonate_cmd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
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/spf13/cobra"
)
var detonateForce bool
var detonateCleanup bool
func buildDetonateCmd() *cobra.Command {
detonateCmd := &cobra.Command{
Use: "detonate attack-technique-id [attack-technique-id]...",
Short: "Detonate one or multiple attack techniques",
Example: strings.Join([]string{
"stratus detonate aws.defense-evasion.cloudtrail-stop",
"stratus detonate aws.defense-evasion.cloudtrail-stop --cleanup",
}, "\n"),
DisableFlagsInUseLine: true,
PreRunE: func(cmd *cobra.Command, args []string) error {
if len(args) == 0 {
cmd.Help()
os.Exit(0)
}
return nil
},
Args: func(cmd *cobra.Command, args []string) error {
if len(args) == 0 {
return errors.New("you must specify at least one attack technique")
}
_, err := resolveTechniques(args)
return err
},
ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
return getTechniquesCompletion(toComplete), cobra.ShellCompDirectiveNoFileComp
},
Run: func(cmd *cobra.Command, args []string) {
techniques, _ := resolveTechniques(args)
doDetonateCmd(techniques, detonateCleanup)
},
}
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) {
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)
}
// Send attack techniques to detonate
for i := range techniques {
techniquesChan <- techniques[i]
}
close(techniquesChan)
if hadError := handleErrorsChannel(errorsChan, workerCount); hadError {
os.Exit(1)
}
}
func detonateCmdWorker(techniques <-chan *stratus.AttackTechnique, errors chan<- error) {
for technique := range techniques {
stratusRunner := runner.NewRunner(technique, detonateForce)
detonateErr := stratusRunner.Detonate()
if detonateCleanup {
cleanupErr := stratusRunner.CleanUp()
errors <- utils.CoalesceErr(detonateErr, cleanupErr)
} else {
errors <- detonateErr
}
}
}