Atomic Red Team support in threatest

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/stretchr/testify v1.8.1
 	golang.org/x/crypto v0.7.0
 	gopkg.in/alessio/shellescape.v1 v1.0.0-20170105083845-52074bc9df61
+	gopkg.in/yaml.v3 v3.0.1
 	sigs.k8s.io/yaml v1.3.0
 )
 
@@ -104,7 +105,6 @@ require (
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/api v0.25.4 // indirect
 	k8s.io/apimachinery v0.25.4 // indirect
 	k8s.io/client-go v0.25.4 // indirect

pkg/atomic/fmt.go

@@ -0,0 +1,47 @@
+package atomic
+
+import (
+	"fmt"
+	"strings"
+)
+
+func (t *Test) FormatCommand(arguments map[string]string) (string, error) {
+	if t.Executor.Name != "bash" {
+		return "", fmt.Errorf("invalid executor `%s`, only `bash` is currently supported", t.Executor.Name)
+	}
+
+	if t.Executor.Command == nil {
+		return "", fmt.Errorf("no command was specified for this test")
+	}
+
+	// TODO: handle dependencies install
+	// for _, dependency := range t.Dependencies {
+	// 	t.InterpolateCommand(dependency.PreReqCommand, arguments)
+	// 	t.InterpolateCommand(dependency.GetPreReqCommand, arguments)
+	// }
+
+	// TODO: handle additional files management
+
+	return t.interpolateCommand(*t.Executor.Command, arguments), nil
+
+	// TODO: handle cleanup command
+	// if t.Executor.CleanupCommand != nil {
+	// 	t.InterpolateCommand(*t.Executor.CleanupCommand, arguments)
+	// }
+}
+
+func (t *Test) interpolateCommand(command string, arguments map[string]string) string {
+	for parameterName, parameterDefinition := range t.InputArguments {
+		var parameterValue string
+
+		if _, ok := arguments[parameterName]; ok {
+			parameterValue = arguments[parameterName]
+		} else if parameterDefinition.Default != nil {
+			parameterValue = *parameterDefinition.Default
+		}
+
+		command = strings.ReplaceAll(command, fmt.Sprintf("${%s}", parameterName), parameterValue)
+	}
+
+	return command
+}

pkg/atomic/fmt_test.go

@@ -0,0 +1,63 @@
+package atomic
+
+import (
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInterpolateCommand(t *testing.T) {
+	testCases := map[string]struct {
+		test           *Test
+		arguments      map[string]string
+		expectedOutput string
+	}{
+		"no templating and no arguments": {
+			test: &Test{
+				Executor:       Executor{Command: aws.String("echo 'Hello World!'")},
+				InputArguments: nil,
+			},
+			arguments:      nil,
+			expectedOutput: "echo 'Hello World!'",
+		},
+		"no templating and some arguments": {
+			test: &Test{
+				Executor:       Executor{Command: aws.String("echo 'Hello World!'")},
+				InputArguments: map[string]InputArgument{"first_name": {}, "last_name": {}},
+			},
+			arguments:      map[string]string{"first_name": "John", "last_name": "Doe"},
+			expectedOutput: "echo 'Hello World!'",
+		},
+		"templating and no arguments": {
+			test: &Test{
+				Executor:       Executor{Command: aws.String("echo 'Hello ${first_name} ${last_name}!'")},
+				InputArguments: map[string]InputArgument{"first_name": {}, "last_name": {}},
+			},
+			arguments:      nil,
+			expectedOutput: "echo 'Hello  !'",
+		},
+		"templating and some arguments": {
+			test: &Test{
+				Executor:       Executor{Command: aws.String("echo 'Hello ${first_name} ${last_name}!'")},
+				InputArguments: map[string]InputArgument{"first_name": {}, "last_name": {}},
+			},
+			arguments:      map[string]string{"first_name": "John"},
+			expectedOutput: "echo 'Hello John !'",
+		},
+		"templating and all arguments": {
+			test: &Test{
+				Executor:       Executor{Command: aws.String("echo 'Hello ${first_name} ${last_name}!'")},
+				InputArguments: map[string]InputArgument{"first_name": {}, "last_name": {}},
+			},
+			arguments:      map[string]string{"first_name": "John", "last_name": "Doe"},
+			expectedOutput: "echo 'Hello John Doe!'",
+		},
+	}
+
+	for testCaseName, testCaseData := range testCases {
+		t.Run(testCaseName, func(t *testing.T) {
+			assert.Equal(t, testCaseData.expectedOutput, testCaseData.test.interpolateCommand(*testCaseData.test.Executor.Command, testCaseData.arguments))
+		})
+	}
+}

pkg/atomic/get.go

@@ -0,0 +1,46 @@
+package atomic
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+
+	"gopkg.in/yaml.v3"
+)
+
+func GetTest(technique, name, version string) (*Test, error) {
+	tech, err := GetTechnique(technique, version)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, test := range tech.AtomicTests {
+		if test.Name == name {
+			return &test, nil
+		}
+	}
+
+	return nil, fmt.Errorf("test '%s' not found for technique '%s'", technique, name)
+}
+
+func GetTechnique(id, version string) (*Technique, error) {
+	url := fmt.Sprintf("https://raw.githubusercontent.com/redcanaryco/atomic-red-team/%s/atomics/%s/%s.yaml", version, id, id)
+
+	res, err := http.Get(url)
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+
+	body, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	var technique *Technique
+	if err := yaml.Unmarshal(body, &technique); err != nil {
+		return nil, err
+	}
+
+	return technique, nil
+}

pkg/atomic/spec.go

@@ -0,0 +1,41 @@
+package atomic
+
+// The following structures were designed following the Atomic test specification:
+// https://github.com/redcanaryco/atomic-red-team/blob/master/atomic_red_team/spec.yaml
+
+type Technique struct {
+	AttackTechnique string `yaml:"attack_technique"`
+	DisplayName     string `yaml:"display_name"`
+	AtomicTests     []Test `yaml:"atomic_tests"`
+}
+
+type Test struct {
+	Name                   string                   `yaml:"name"`
+	AutoGeneratedGUID      string                   `yaml:"auto_generated_guid"`
+	Description            string                   `yaml:"description"`
+	SupportedPlatforms     []string                 `yaml:"supported_platforms"`
+	InputArguments         map[string]InputArgument `yaml:"input_arguments,omitempty"`
+	DependencyExecutorName *string                  `yaml:"dependency_executor_name,omitempty"`
+	Dependencies           []Dependency             `yaml:"dependencies,omitempty"`
+	Executor               Executor                 `yaml:"executor"`
+}
+
+type InputArgument struct {
+	Description string  `yaml:"description"`
+	Type        string  `yaml:"type"`
+	Default     *string `yaml:"default,omitempty"`
+}
+
+type Dependency struct {
+	Description      string `yaml:"description"`
+	PreReqCommand    string `yaml:"prereq_command"`
+	GetPreReqCommand string `yaml:"get_prereq_command"`
+}
+
+type Executor struct {
+	Name              string  `yaml:"name"`
+	ElevationRequired *bool   `yaml:"elevation_required,omitempty"`
+	Command           *string `yaml:"command,omitempty"`
+	CleanupCommand    *string `yaml:"cleanup_command,omitempty"`
+	Steps             *string `yaml:"steps,omitempty"`
+}

pkg/threatest/parser/main.go

@@ -2,12 +2,14 @@ package parser
 
 import (
 	"fmt"
+	"strings"
+	"time"
+
+	"github.com/datadog/threatest/pkg/atomic"
 	"github.com/datadog/threatest/pkg/threatest"
 	"github.com/datadog/threatest/pkg/threatest/detonators"
 	"github.com/datadog/threatest/pkg/threatest/matchers/datadog"
 	"sigs.k8s.io/yaml" // we use this library as it provides a handy "YAMLToJSON" function
-	"strings"
-	"time"
 )
 
 // Parse turns a YAML input string into a list of Threatest scenarios
@@ -42,10 +44,50 @@ func buildScenarios(parsed *ThreatestSchemaJson, sshHostname string, sshUsername
 
 		// Detonation
 		if localDetonator := parsedScenario.Detonate.LocalDetonator; localDetonator != nil {
-			commandToRun := strings.Join(parsedScenario.Detonate.LocalDetonator.Commands, "; ")
+			var commandToRun string
+
+			if localDetonator.AtomicReadTeam != nil {
+				version := "master" // default git tree to fetch atomic red tests from
+				if localDetonator.AtomicReadTeam.Version != nil {
+					version = *localDetonator.AtomicReadTeam.Version
+				}
+
+				test, err := atomic.GetTest(localDetonator.AtomicReadTeam.Technique, localDetonator.AtomicReadTeam.Name, version)
+				if err != nil {
+					return nil, fmt.Errorf("failed to retrieve atomic red team test '%s' (%s): %w", localDetonator.AtomicReadTeam.Name, localDetonator.AtomicReadTeam.Technique, err)
+				}
+
+				commandToRun, err = test.FormatCommand(localDetonator.AtomicReadTeam.Inputs)
+				if err != nil {
+					return nil, err
+				}
+			} else {
+				commandToRun = strings.Join(localDetonator.Commands, "; ")
+			}
+
 			scenario.Detonator = detonators.NewCommandDetonator(&detonators.LocalCommandExecutor{}, commandToRun)
 		} else if remoteDetonator := parsedScenario.Detonate.RemoteDetonator; remoteDetonator != nil {
-			commandToRun := strings.Join(remoteDetonator.Commands, "; ")
+			var commandToRun string
+
+			if remoteDetonator.AtomicReadTeam != nil {
+				version := "master" // default git tree to fetch atomic red tests from
+				if remoteDetonator.AtomicReadTeam.Version != nil {
+					version = *remoteDetonator.AtomicReadTeam.Version
+				}
+
+				test, err := atomic.GetTest(remoteDetonator.AtomicReadTeam.Technique, remoteDetonator.AtomicReadTeam.Name, version)
+				if err != nil {
+					return nil, fmt.Errorf("failed to retrieve atomic red team test '%s' (%s): %w", remoteDetonator.AtomicReadTeam.Name, remoteDetonator.AtomicReadTeam.Technique, err)
+				}
+
+				commandToRun, err = test.FormatCommand(remoteDetonator.AtomicReadTeam.Inputs)
+				if err != nil {
+					return nil, err
+				}
+			} else {
+				commandToRun = strings.Join(remoteDetonator.Commands, "; ")
+			}
+
 			//TODO: decouple
 			//TODO: confirm 1 SSH executor per attack makes sense
 			sshExecutor, err := detonators.NewSSHCommandExecutor(sshHostname, sshUsername, sshKey)