-
Notifications
You must be signed in to change notification settings - Fork 197
/
status_cmd.go
60 lines (56 loc) · 1.65 KB
/
status_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
package main
import (
"github.com/datadog/stratus-red-team/v2/internal/state"
"github.com/datadog/stratus-red-team/v2/pkg/stratus"
"github.com/fatih/color"
"github.com/jedib0t/go-pretty/v6/table"
"github.com/spf13/cobra"
)
func buildStatusCmd() *cobra.Command {
statusCmd := &cobra.Command{
Use: "status",
Short: "Display the status of TTPs.",
Args: func(cmd *cobra.Command, args []string) error {
if len(args) == 0 {
return nil // no technique specified == all techniques
}
_, err := resolveTechniques(args)
return err
},
Run: func(cmd *cobra.Command, args []string) {
if len(args) > 0 {
techniques, _ := resolveTechniques(args)
doStatusCmd(techniques)
} else {
doStatusCmd(stratus.GetRegistry().ListAttackTechniques())
}
},
}
return statusCmd
}
func doStatusCmd(techniques []*stratus.AttackTechnique) {
t := GetDisplayTable()
t.AppendHeader(table.Row{"ID", "Name", "Status"})
for i := range techniques {
stateManager := state.NewFileSystemStateManager(techniques[i])
techniqueState := stateManager.GetTechniqueState()
if techniqueState == "" {
techniqueState = stratus.AttackTechniqueStatusCold
}
t.AppendRow(table.Row{techniques[i].ID, techniques[i].FriendlyName, colorState(techniqueState)})
}
t.Render()
}
func colorState(state stratus.AttackTechniqueState) string {
stateString := string(state)
switch state {
case stratus.AttackTechniqueStatusCold:
return color.CyanString(stateString)
case stratus.AttackTechniqueStatusWarm:
return color.YellowString(stateString)
case stratus.AttackTechniqueStatusDetonated:
return color.MagentaString(stateString)
default:
return stateString
}
}