diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 25d2b5ad..d57134c7 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,7 +22,7 @@ repos: rev: v3.0.3 hooks: - id: prettier - files: \.(json|yaml|yml)$ + files: \.(json)$ exclude: docs/container.md - repo: https://github.com/jumanjihouse/pre-commit-hooks @@ -84,13 +84,6 @@ repos: entry: .hooks/go-vet.sh files: '\.go$' - - id: go-critic - name: go-critic - language: script - files: '\.go$' - entry: .hooks/go-critic.sh - args: ["check"] - - id: go-licenses name: Run go-licenses language: script diff --git a/README.md b/README.md index e0107507..949c24ee 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,6 @@ targets and mediums. - [Using the TTPForge Dev Container](docs/container.md) - [Code Standards](docs/code-standards.md) - [Creating a new release](docs/release.md) -- [TTPForge Building Blocks](docs/building-blocks.md) --- diff --git a/cmd/run.go b/cmd/run.go index 102d8804..7ac0d83d 100755 --- a/cmd/run.go +++ b/cmd/run.go @@ -51,12 +51,12 @@ func buildRunCommand(cfg *Config) *cobra.Command { // based on the TTPs argument value specifications ttpCfg.Repo = foundRepo - ttp, err := blocks.LoadTTP(ttpAbsPath, foundRepo.GetFs(), &ttpCfg, argsList) + ttp, execCtx, err := blocks.LoadTTP(ttpAbsPath, foundRepo.GetFs(), &ttpCfg, argsList) if err != nil { return fmt.Errorf("could not load TTP at %v:\n\t%v", ttpAbsPath, err) } - if _, err := ttp.RunSteps(ttpCfg); err != nil { + if _, err := ttp.Execute(execCtx); err != nil { return fmt.Errorf("failed to run TTP at %v: %v", ttpAbsPath, err) } return nil diff --git a/cmd/run_test.go b/cmd/run_test.go index c0ca461b..80ca4145 100644 --- a/cmd/run_test.go +++ b/cmd/run_test.go @@ -22,12 +22,9 @@ package cmd_test import ( "bytes" "path/filepath" - "strings" "testing" "github.com/facebookincubator/ttpforge/cmd" - "github.com/facebookincubator/ttpforge/pkg/blocks" - "github.com/spf13/afero" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -70,7 +67,18 @@ func TestRun(t *testing.T) { testConfigFilePath, "another-repo//simple-inline.yaml", }, - expectedStdout: "simple inline was executed\n", + expectedStdout: "simple inline was executed\ncleaning up simple inline\n", + }, + { + name: "subttp-cleanup", + description: "verify that execution of a subTTP with cleanup succeeds", + args: []string{ + "-c", + testConfigFilePath, + "another-repo//sub-ttp-example/ttp.yaml", + }, + expectedStdout: "subttp1_step_1\nsubttp1_step_2\nsubttp2_step_1\nsubttp2_step_1_cleanup\nsubttp1_step_2_cleanup\nsubttp1_step_1_cleanup\n", + wantError: true, }, { name: "dry-run-success", @@ -94,6 +102,27 @@ func TestRun(t *testing.T) { }, wantError: true, }, + { + name: "no-cleanup", + description: "Using the no-cleanup flag should prevent cleanup", + args: []string{ + "-c", + testConfigFilePath, + "--no-cleanup", + "another-repo//simple-inline.yaml", + }, + expectedStdout: "simple inline was executed\n", + }, + { + name: "cleanup-stress-test", + description: "run many different execute+cleanup combinations", + args: []string{ + "-c", + testConfigFilePath, + "another-repo//cleanup-tests/stress-tests.yaml", + }, + expectedStdout: "execute_step_1\nexecute_step_2\nexecute_step_3\nexecute_step_4\ncleanup_step_4\ncleanup_step_3\ncleanup_step_2\ncleanup_step_1\n", + }, } for _, tc := range testCases { @@ -114,89 +143,3 @@ func TestRun(t *testing.T) { }) } } - -func TestNoCleanupFlag(t *testing.T) { - afs := afero.NewOsFs() - testCases := []struct { - name string - content string - execConfig blocks.TTPExecutionConfig - expectedDirExist bool - wantError bool - }{ - { - name: "Test No Cleanup Behavior - Directory Creation", - content: ` ---- -name: test-cleanup -steps: - - name: step_one - inline: mkdir testDir - cleanup: - inline: rm -rf testDir`, - execConfig: blocks.TTPExecutionConfig{ - NoCleanup: true, - }, - expectedDirExist: true, - wantError: false, - }, - { - name: "Test Cleanup Behavior - Directory Deletion", - content: ` ---- -name: test-cleanup-2 -steps: - - name: step_two - inline: mkdir testDir2 - cleanup: - inline: rm -rf testDir2`, - execConfig: blocks.TTPExecutionConfig{ - NoCleanup: false, - }, - expectedDirExist: false, - wantError: false, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - // Create a temp directory to work within - tempDir, err := afero.TempDir(afs, "", "testCleanup") - require.NoError(t, err) - - // Update content to work within the temp directory - tc.content = strings.ReplaceAll(tc.content, "mkdir ", "mkdir "+tempDir+"/") - tc.content = strings.ReplaceAll(tc.content, "rm -rf ", "rm -rf "+tempDir+"/") - - // Render the templated TTP first - ttp, err := blocks.RenderTemplatedTTP(tc.content, &tc.execConfig) - require.NoError(t, err) - - // Handle potential error from RemoveAll within a deferred function - defer func() { - err := afs.RemoveAll(tempDir) // cleanup temp directory - if err != nil { - t.Errorf("failed to remove temp directory: %v", err) - } - }() - - _, err = ttp.RunSteps(tc.execConfig) - if tc.wantError { - require.Error(t, err) - } else { - require.NoError(t, err) - } - - // Determine which directory to check based on the test case content - dirName := tempDir + "/testDir" - if strings.Contains(tc.content, "testDir2") { - dirName = tempDir + "/testDir2" - } - - // Check if the directory exists - dirExists, err := afero.DirExists(afs, dirName) - require.NoError(t, err) - assert.Equal(t, tc.expectedDirExist, dirExists) - }) - } -} diff --git a/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/stress-tests.yaml b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/stress-tests.yaml new file mode 100644 index 00000000..1b19570c --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/stress-tests.yaml @@ -0,0 +1,30 @@ +--- +name: cleanup-stress-test +description: | + Cleanup a whole bunch of different ways +steps: + - name: file-step-with-inline-cleanup + file: test.sh + args: + - execute_step_1 + cleanup: + inline: | + echo "cleanup_step_1" + - name: print-with-file-cleanup + print_str: execute_step_2 + cleanup: + file: test.sh + args: + - cleanup_step_2 + - name: ttp-with-print-cleanup + ttp: cleanup-tests/subttp/ttp.yaml + args: + to_print: execute_step_3 + cleanup: + print_str: cleanup_step_3 + - name: inline-with-ttp-cleanup + inline: echo execute_step_4 + cleanup: + ttp: cleanup-tests/subttp/ttp.yaml + args: + to_print: cleanup_step_4 diff --git a/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/subttp/ttp.yaml b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/subttp/ttp.yaml new file mode 100644 index 00000000..bb40cea3 --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/subttp/ttp.yaml @@ -0,0 +1,7 @@ +--- +name: print-stuff +args: + - name: to_print +steps: + - name: print-arg + print_str: {{.Args.to_print}} diff --git a/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/test.sh b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/test.sh new file mode 100755 index 00000000..ccc95ec1 --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/cleanup-tests/test.sh @@ -0,0 +1,2 @@ +#!/bin/bash +echo "$1" diff --git a/cmd/test-resources/repos/another-repo/some-ttps/simple-inline.yaml b/cmd/test-resources/repos/another-repo/some-ttps/simple-inline.yaml index 0028acf4..0989a699 100644 --- a/cmd/test-resources/repos/another-repo/some-ttps/simple-inline.yaml +++ b/cmd/test-resources/repos/another-repo/some-ttps/simple-inline.yaml @@ -3,3 +3,5 @@ name: basic_inline steps: - name: hello inline: echo "simple inline was executed" + cleanup: + inline: echo "cleaning up simple inline" diff --git a/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp1.yaml b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp1.yaml new file mode 100644 index 00000000..1cc19db9 --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp1.yaml @@ -0,0 +1,11 @@ +--- +name: subttp1 +steps: + - name: subttp1_step_1 + inline: echo subttp1_step_1 + cleanup: + inline: echo subttp1_step_1_cleanup + - name: subttp1_step_2 + inline: echo subttp1_step_2 + cleanup: + inline: echo subttp1_step_2_cleanup diff --git a/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp2.yaml b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp2.yaml new file mode 100644 index 00000000..6b482b74 --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/subttp2.yaml @@ -0,0 +1,11 @@ +--- +name: subttp2 +steps: + - name: subttp2_step_1 + inline: echo subttp2_step_1 + cleanup: + inline: echo subttp2_step_1_cleanup + - name: intentional_failure + inline: foobarbaz + cleanup: + inline: echo subttp2_step_2_cleanup diff --git a/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/ttp.yaml b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/ttp.yaml new file mode 100644 index 00000000..519cb4c6 --- /dev/null +++ b/cmd/test-resources/repos/another-repo/some-ttps/sub-ttp-example/ttp.yaml @@ -0,0 +1,7 @@ +--- +name: subttp_cleanup_test +steps: + - name: first_sub_ttp + ttp: sub-ttp-example/subttp1.yaml + - name: second_sub_ttp + ttp: sub-ttp-example/subttp2.yaml diff --git a/docs/building-blocks.md b/docs/building-blocks.md deleted file mode 100644 index caf5c689..00000000 --- a/docs/building-blocks.md +++ /dev/null @@ -1,48 +0,0 @@ -# Building Blocks - -This document introduces our use of Golang Interfaces and Structs -in `TTPForge` to create the concepts of `Acts` and `Steps`, which -form the building blocks of Tactics, Techniques, and Procedures (TTPs). - -## Act - -`Acts` contain the fields and functionality that are common across -different types of `Steps`. - -This interface facilitates building TTPs with steps that don't need to -explicitly understand what the others are doing to function. Each step -has its own enforcement mechanisms and can function without needing to -understand how the previous step worked, simply providing information -for the current step to use. - -This design enables two different step types to offer precise control of -process lineage and corresponding indicators of compromise. - -## Step - -`Steps` contain a high-level interface that can be used to define -instructions or `Steps` that make up a TTP. The following `Steps` exist today: - -- `BasicStep`: Simulates interactive interpreter abuse, which can - be used to simulate hands-on-keyboard TTPs. -- `FileStep`: Simulates direct process execution without an - intermediate interpreter, which can be used to simulate an attacker - executing logic from a file. -- `SubTTPStep`: A Collection of `TTPs` that are used together to represent a `TTP`. - -## TTP - -`TTPs` are collections of `Steps` that make up the logic to represent -a component of a Tactic, Technique, and Procedure. - -## Cleanup - -`Cleanup` is an embedded structure for all Acts to make use of in order -to clean up anything that a TTP has done. - -## Resources - -- [LogRocket: Exploring Structs and Interfaces in Go](https://blog.logrocket.com/exploring-structs-interfaces-go/) -- [Go by Example: Interfaces](https://gobyexample.com/interfaces) -- [A Tour of Go: Interfaces](https://go.dev/tour/methods/9) -- [Introduction to Tokenizers and Compilers](https://www.cs.man.ac.uk/~pjj/farrell/comp3.html) diff --git a/pkg/blocks/README.md b/pkg/blocks/README.md index 272aa830..58965fb5 100644 --- a/pkg/blocks/README.md +++ b/pkg/blocks/README.md @@ -17,173 +17,16 @@ The `blocks` package is a part of the TTPForge. ## Functions -### Act.CheckCondition() - -```go -CheckCondition() bool, error -``` - -CheckCondition checks the condition specified for an Act and returns true -if it matches the current OS, false otherwise. If the condition is "always", -the function returns true. -If an error occurs while checking the condition, it is returned. - -**Returns:** - -bool: true if the condition matches the current OS or the -condition is "always", false otherwise. - -error: An error if an error occurs while checking the condition. - ---- - -### Act.ExplainInvalid() - -```go -ExplainInvalid() error -``` - -ExplainInvalid returns an error explaining why the Act is invalid. - -**Returns:** - -error: An error explaining why the Act is invalid, or nil -if the Act is valid. - ---- - -### Act.IsNil() - -```go -IsNil() bool -``` - -IsNil checks whether the Act is nil (i.e., it does not have a name). - -**Returns:** - -bool: True if the Act has no name, false otherwise. - ---- - -### Act.MakeCleanupStep(*yaml.Node) - -```go -MakeCleanupStep(*yaml.Node) CleanupAct, error -``` - -MakeCleanupStep creates a CleanupAct based on the given yaml.Node. -If the node is empty or invalid, it returns nil. If the node contains a -BasicStep or FileStep, the corresponding CleanupAct is created and returned. - -**Parameters:** - -node: A pointer to a yaml.Node containing the parameters to -create the CleanupAct. - -**Returns:** - -CleanupAct: The created CleanupAct, or nil if the node is empty or invalid. - -error: An error if the node contains invalid parameters. - ---- - -### Act.SetDir(string) - -```go -SetDir(string) -``` - -SetDir sets the working directory for the Act. - -**Parameters:** - -dir: A string representing the directory path to be set -as the working directory. - ---- - -### Act.StepName() - -```go -StepName() string -``` - -StepName returns the name of the Act. - -**Returns:** - -string: The name of the Act. - ---- - -### Act.Validate() - -```go -Validate() error -``` - -Validate checks the Act for any validation errors, such as the presence of -spaces in the name. - -**Returns:** - -error: An error if any validation errors are found, or nil if -the Act is valid. - ---- - -### BasicStep.Cleanup(TTPExecutionContext) - -```go -Cleanup(TTPExecutionContext) *ActResult, error -``` - -Cleanup is an implementation of the CleanupAct interface's Cleanup method. - ---- - ### BasicStep.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` Execute runs the BasicStep and returns an error if any occur. --- -### BasicStep.ExplainInvalid() - -```go -ExplainInvalid() error -``` - -ExplainInvalid returns an error with an explanation of why a BasicStep is invalid. - ---- - -### BasicStep.GetCleanup() - -```go -GetCleanup() []CleanupAct -``` - -GetCleanup returns the cleanup steps for a BasicStep. - ---- - -### BasicStep.GetType() - -```go -GetType() StepType -``` - -GetType returns the step type for a BasicStep. - ---- - ### BasicStep.IsNil() ```go @@ -194,16 +37,6 @@ IsNil checks if a BasicStep is considered empty or uninitialized. --- -### BasicStep.UnmarshalYAML(*yaml.Node) - -```go -UnmarshalYAML(*yaml.Node) error -``` - -UnmarshalYAML custom unmarshaler for BasicStep to handle decoding from YAML. - ---- - ### BasicStep.Validate(TTPExecutionContext) ```go @@ -214,60 +47,24 @@ Validate validates the BasicStep, checking for the necessary attributes and depe --- -### CreateFileStep.Cleanup(TTPExecutionContext) - -```go -Cleanup(TTPExecutionContext) *ActResult, error -``` - -Cleanup is a method to establish a link with the Cleanup interface. -Assumes that the type is the cleanup step and is invoked by -s.CleanupStep.Cleanup. - ---- - ### CreateFileStep.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` Execute runs the step and returns an error if any occur. --- -### CreateFileStep.ExplainInvalid() - -```go -ExplainInvalid() error -``` - -ExplainInvalid returns an error message explaining why the step -is invalid. - -**Returns:** - -error: An error message explaining why the step is invalid. - ---- - -### CreateFileStep.GetCleanup() +### CreateFileStep.GetDefaultCleanupAction() ```go -GetCleanup() []CleanupAct +GetDefaultCleanupAction() Action ``` -GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. - ---- - -### CreateFileStep.GetType() - -```go -GetType() StepType -``` - -GetType returns the type of the step as StepType. +GetDefaultCleanupAction will instruct the calling code +to remove the path created by this action --- @@ -281,26 +78,6 @@ IsNil checks if the step is nil or empty and returns a boolean value. --- -### CreateFileStep.UnmarshalYAML(*yaml.Node) - -```go -UnmarshalYAML(*yaml.Node) error -``` - -UnmarshalYAML decodes a YAML node into a CreateFileStep instance. It uses -the provided struct as a template for the YAML data, and initializes the -CreateFileStep instance with the decoded values. - -**Parameters:** - -node: A pointer to a yaml.Node representing the YAML data to decode. - -**Returns:** - -error: An error if there is a problem decoding the YAML data. - ---- - ### CreateFileStep.Validate(TTPExecutionContext) ```go @@ -318,35 +95,13 @@ error: An error if any validation checks fail. ### EditStep.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` Execute runs the EditStep and returns an error if any occur. --- -### EditStep.GetCleanup() - -```go -GetCleanup() []CleanupAct -``` - -GetCleanup returns the cleanup steps for a EditStep. -Currently this is always empty because we use backup -files instead for this type of step - ---- - -### EditStep.GetType() - -```go -GetType() StepType -``` - -GetType returns the step type for a EditStep. - ---- - ### EditStep.IsNil() ```go @@ -427,48 +182,13 @@ f.CleanupStep.Cleanup. ### FetchURIStep.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` Execute runs the FetchURIStep and returns an error if any occur. --- -### FetchURIStep.ExplainInvalid() - -```go -ExplainInvalid() error -``` - -ExplainInvalid returns an error message explaining why the FetchURIStep -is invalid. - -**Returns:** - -error: An error message explaining why the FetchURIStep is invalid. - ---- - -### FetchURIStep.GetCleanup() - -```go -GetCleanup() []CleanupAct -``` - -GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. - ---- - -### FetchURIStep.GetType() - -```go -GetType() StepType -``` - -GetType returns the type of the step as StepType. - ---- - ### FetchURIStep.IsNil() ```go @@ -479,26 +199,6 @@ IsNil checks if the FetchURIStep is nil or empty and returns a boolean value. --- -### FetchURIStep.UnmarshalYAML(*yaml.Node) - -```go -UnmarshalYAML(*yaml.Node) error -``` - -UnmarshalYAML decodes a YAML node into a FetchURIStep instance. It uses -the provided struct as a template for the YAML data, and initializes the -FetchURIStep instance with the decoded values. - -**Parameters:** - -node: A pointer to a yaml.Node representing the YAML data to decode. - -**Returns:** - -error: An error if there is a problem decoding the YAML data. - ---- - ### FetchURIStep.Validate(TTPExecutionContext) ```go @@ -533,48 +233,13 @@ f.CleanupStep.Cleanup. ### FileStep.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` Execute runs the FileStep and returns an error if any occur. --- -### FileStep.ExplainInvalid() - -```go -ExplainInvalid() error -``` - -ExplainInvalid returns an error message explaining why the FileStep -is invalid. - -**Returns:** - -error: An error message explaining why the FileStep is invalid. - ---- - -### FileStep.GetCleanup() - -```go -GetCleanup() []CleanupAct -``` - -GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. - ---- - -### FileStep.GetType() - -```go -GetType() StepType -``` - -GetType returns the type of the step as StepType. - ---- - ### FileStep.IsNil() ```go @@ -585,26 +250,6 @@ IsNil checks if the FileStep is nil or empty and returns a boolean value. --- -### FileStep.UnmarshalYAML(*yaml.Node) - -```go -UnmarshalYAML(*yaml.Node) error -``` - -UnmarshalYAML decodes a YAML node into a FileStep instance. It uses -the provided struct as a template for the YAML data, and initializes the -FileStep instance with the decoded values. - -**Parameters:** - -node: A pointer to a yaml.Node representing the YAML data to decode. - -**Returns:** - -error: An error if there is a problem decoding the YAML data. - ---- - ### FileStep.Validate(TTPExecutionContext) ```go @@ -671,7 +316,7 @@ returns it as a string. ### LoadTTP(string, afero.Fs, *TTPExecutionConfig, []string) ```go -LoadTTP(string, afero.Fs, *TTPExecutionConfig, []string) *TTP, error +LoadTTP(string afero.Fs *TTPExecutionConfig []string) *TTP *TTPExecutionContext error ``` LoadTTP reads a TTP file and creates a TTP instance based on its contents. @@ -684,7 +329,8 @@ fsys: an afero.Fs that contains the specified TTP file path **Returns:** -ttp: Pointer to the created TTP instance, or nil if the file is empty or invalid. +*TTP: Pointer to the created TTP instance, or nil if the file is empty or invalid. +TTPExecutionContext: the initialized TTPExecutionContext suitable for passing to TTP.Execute(...) err: An error if the file contains invalid data or cannot be read. --- @@ -759,6 +405,74 @@ NewSubTTPStep creates a new SubTTPStep and returns a pointer to it. --- +### PrintStrAction.Execute(TTPExecutionContext) + +```go +Execute(TTPExecutionContext) *ActResult, error +``` + +Execute runs the step and returns an error if any occur. + +--- + +### PrintStrAction.IsNil() + +```go +IsNil() bool +``` + +IsNil checks if the step is nil or empty and returns a boolean value. + +--- + +### PrintStrAction.Validate(TTPExecutionContext) + +```go +Validate(TTPExecutionContext) error +``` + +Validate validates the step + +**Returns:** + +error: An error if any validation checks fail. + +--- + +### RemovePathAction.Execute(TTPExecutionContext) + +```go +Execute(TTPExecutionContext) *ActResult, error +``` + +Execute runs the step and returns an error if any occur. + +--- + +### RemovePathAction.IsNil() + +```go +IsNil() bool +``` + +IsNil checks if the step is nil or empty and returns a boolean value. + +--- + +### RemovePathAction.Validate(TTPExecutionContext) + +```go +Validate(TTPExecutionContext) error +``` + +Validate validates the step + +**Returns:** + +error: An error if any validation checks fail. + +--- + ### RenderTemplatedTTP(string, *TTPExecutionConfig) ```go @@ -782,77 +496,117 @@ error: An error if the rendering or unmarshaling process fails. --- -### SubTTPStep.Cleanup(TTPExecutionContext) +### ShouldUseImplicitDefaultCleanup(Action) + +```go +ShouldUseImplicitDefaultCleanup(Action) bool +``` + +ShouldUseImplicitDefaultCleanup is a hack +to make subTTPs always run their default +cleanup process even when `cleanup: default` is +not explicitly specified - this is purely for backward +compatibility + +--- + +### Step.Cleanup(TTPExecutionContext) ```go Cleanup(TTPExecutionContext) *ActResult, error ``` -Cleanup runs the cleanup actions associated with all successful sub-steps +Cleanup runs the cleanup action associated with this step --- -### SubTTPStep.Execute(TTPExecutionContext) +### Step.Execute(TTPExecutionContext) ```go -Execute(TTPExecutionContext) *ExecutionResult, error +Execute(TTPExecutionContext) *ActResult, error ``` -Execute runs each step of the TTP file associated with the SubTTPStep -and manages the outputs and cleanup steps. +Execute runs the action associated with this step --- -### SubTTPStep.ExplainInvalid() +### Step.ParseAction(*yaml.Node) ```go -ExplainInvalid() error +ParseAction(*yaml.Node) Action, error ``` -ExplainInvalid checks for invalid data in the SubTTPStep -and returns an error explaining any issues found. -Currently, it checks if the TtpFile field is empty. +ParseAction decodes an action (from step or cleanup) in YAML +format into the appropriate struct --- -### SubTTPStep.GetCleanup() +### Step.ShouldCleanupOnFailure() ```go -GetCleanup() []CleanupAct +ShouldCleanupOnFailure() bool ``` -GetCleanup returns a slice of CleanupAct associated with the SubTTPStep. +ShouldCleanupOnFailure specifies that this step should be cleaned +up even if its Execute(...) failed. +We usually don't want to do this - for example, +you shouldn't try to remove_path a create_file that failed) +However, certain step types (especially SubTTPs) need to run cleanup even if they fail --- -### SubTTPStep.GetType() +### Step.UnmarshalYAML(*yaml.Node) ```go -GetType() StepType +UnmarshalYAML(*yaml.Node) error ``` -GetType returns the type of the step (StepSubTTP for SubTTPStep). +UnmarshalYAML implements custom deserialization +process to ensure that the step action and its +cleanup action are decoded to the correct struct type --- -### SubTTPStep.IsNil() +### Step.Validate(TTPExecutionContext) ```go -IsNil() bool +Validate(TTPExecutionContext) error ``` -IsNil checks if the SubTTPStep is empty or uninitialized. +Validate checks that both the step action and cleanup +action are valid --- -### SubTTPStep.UnmarshalYAML(*yaml.Node) +### SubTTPStep.Execute(TTPExecutionContext) ```go -UnmarshalYAML(*yaml.Node) error +Execute(TTPExecutionContext) *ActResult, error +``` + +Execute runs each step of the TTP file associated with the SubTTPStep +and manages the outputs and cleanup steps. + +--- + +### SubTTPStep.GetDefaultCleanupAction() + +```go +GetDefaultCleanupAction() Action ``` -UnmarshalYAML is a custom unmarshaller for SubTTPStep which decodes -a YAML node into a SubTTPStep instance. +GetDefaultCleanupAction will instruct the calling code +to cleanup all successful steps of this subTTP + +--- + +### SubTTPStep.IsNil() + +```go +IsNil() bool +``` + +IsNil checks if the SubTTPStep is empty or uninitialized. --- @@ -872,6 +626,26 @@ If any of these conditions are not met, an error is returned. --- +### TTP.Execute(*TTPExecutionContext) + +```go +Execute(*TTPExecutionContext) *StepResultsRecord, error +``` + +Execute executes all of the steps in the given TTP, +then runs cleanup if appropriate + +**Parameters:** + +execCfg: The TTPExecutionConfig for the current TTP. + +**Returns:** + +*StepResultsRecord: A StepResultsRecord containing the results of each step. +error: An error if any of the steps fail to execute. + +--- + ### TTP.MarshalYAML() ```go @@ -889,86 +663,109 @@ error: An error if the encoding process fails. --- -### TTP.RunSteps(TTPExecutionConfig) +### TTP.RunSteps(*TTPExecutionContext) ```go -RunSteps(TTPExecutionConfig) *StepResultsRecord, error +RunSteps(*TTPExecutionContext) *StepResultsRecord, int, error ``` RunSteps executes all of the steps in the given TTP. **Parameters:** -execCfg: The TTPExecutionConfig for the current TTP. +execCtx: The current TTPExecutionContext **Returns:** *StepResultsRecord: A StepResultsRecord containing the results of each step. +int: the index of the step where cleanup should start (usually the last successful step) error: An error if any of the steps fail to execute. --- -### TTP.UnmarshalYAML(*yaml.Node) +### TTP.Validate(TTPExecutionContext) ```go -UnmarshalYAML(*yaml.Node) error +Validate(TTPExecutionContext) error ``` -UnmarshalYAML is a custom unmarshalling implementation for the TTP structure. -It decodes a YAML Node into a TTP object, handling the decoding and -validation of the individual steps within the TTP. +Validate ensures that all components of the TTP are valid +It checks key fields, then iterates through each step +and validates them in turn **Parameters:** -node: A pointer to a yaml.Node that represents the TTP structure -to be unmarshalled. +execCtx: The TTPExecutionContext for the current TTP. **Returns:** -error: An error if the decoding process fails or if the TTP structure contains invalid steps. +error: An error if any part of the validation fails, otherwise nil. --- -### TTP.ValidateSteps(TTPExecutionContext) +### TTPExecutionContext.ExpandVariables([]string) ```go -ValidateSteps(TTPExecutionContext) error +ExpandVariables([]string) []string, error ``` -ValidateSteps iterates through each step in the TTP and validates it. -It sets the working directory for each step before calling its Validate -method. If any step fails validation, the method returns an error. -If all steps are successfully validated, the method returns nil. +ExpandVariables takes a string containing the following types of variables +and expands all of them to their appropriate values: + +* Step outputs: ($forge.steps.bar.outputs.baz) **Parameters:** -execCtx: The TTPExecutionContext for the current TTP. +inStrs: the list of strings that have variables expanded **Returns:** -error: An error if any step validation fails, otherwise nil. +[]string: the corresponding strings with variables expanded +error: an error if there is a problem --- -### TTPExecutionContext.ExpandVariables([]string) +### actionDefaults.GetDefaultCleanupAction() ```go -ExpandVariables([]string) []string, error +GetDefaultCleanupAction() Action ``` -ExpandVariables takes a string containing the following types of variables -and expands all of them to their appropriate values: +GetDefaultCleanupAction provides a default implementation +of the GetDefaultCleanupAction method from the Action interface. +This saves us from having to declare this function for every steps +If a specific action needs a default cleanup action (such as a create_file action), +it can override this step -* Step outputs: ($forge.steps.bar.outputs.baz) +--- -**Parameters:** +### subTTPCleanupAction.Execute(TTPExecutionContext) -inStrs: the list of strings that have variables expanded +```go +Execute(TTPExecutionContext) *ActResult, error +``` -**Returns:** +Execute will cleanup the subTTP starting from the last successful step -[]string: the corresponding strings with variables expanded -error: an error if there is a problem +--- + +### subTTPCleanupAction.IsNil() + +```go +IsNil() bool +``` + +IsNil is not needed here, as this is not a user-accessible step type + +--- + +### subTTPCleanupAction.Validate(TTPExecutionContext) + +```go +Validate(TTPExecutionContext) error +``` + +Validate is not needed here, as this is not a user-accessible step type --- diff --git a/pkg/blocks/actions.go b/pkg/blocks/actions.go new file mode 100644 index 00000000..50969a65 --- /dev/null +++ b/pkg/blocks/actions.go @@ -0,0 +1,41 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks + +// Action is an interface that is implemented +// by all action types used in steps/cleanups +// (such as create_file, inline, etc) +type Action interface { + Execute(execCtx TTPExecutionContext) (*ActResult, error) + Validate(execCtx TTPExecutionContext) error + GetDefaultCleanupAction() Action + IsNil() bool +} + +type actionDefaults struct{} + +// GetDefaultCleanupAction provides a default implementation +// of the GetDefaultCleanupAction method from the Action interface. +// This saves us from having to declare this function for every steps +// If a specific action needs a default cleanup action (such as a create_file action), +// it can override this step +func (ad *actionDefaults) GetDefaultCleanupAction() Action { + return nil +} diff --git a/pkg/blocks/basicstep.go b/pkg/blocks/basicstep.go index e6a8e8b9..5c590dbb 100755 --- a/pkg/blocks/basicstep.go +++ b/pkg/blocks/basicstep.go @@ -31,108 +31,37 @@ import ( "github.com/facebookincubator/ttpforge/pkg/logging" "github.com/facebookincubator/ttpforge/pkg/outputs" "go.uber.org/zap" - "gopkg.in/yaml.v3" +) + +// These are all the different executors that could run +// our inline command +const ( + ExecutorPython = "python3" + ExecutorBash = "bash" + ExecutorSh = "sh" + ExecutorPowershell = "powershell" + ExecutorRuby = "ruby" + ExecutorBinary = "binary" + ExecutorCmd = "cmd.exe" ) // BasicStep is a type that represents a basic execution step. type BasicStep struct { - *Act `yaml:",inline"` - Executor string `yaml:"executor,omitempty"` - Inline string `yaml:"inline,flow"` - Args []string `yaml:"args,omitempty,flow"` - CleanupStep CleanupAct `yaml:"cleanup,omitempty"` + actionDefaults `yaml:"-"` + Executor string `yaml:"executor,omitempty"` + Inline string `yaml:"inline,flow"` + Environment map[string]string `yaml:"env,omitempty"` + Outputs map[string]outputs.Spec `yaml:"outputs,omitempty"` } // NewBasicStep creates a new BasicStep instance with an initialized Act struct. func NewBasicStep() *BasicStep { - return &BasicStep{ - Act: &Act{ - Type: StepBasic, - }, - } -} - -// UnmarshalYAML custom unmarshaler for BasicStep to handle decoding from YAML. -func (b *BasicStep) UnmarshalYAML(node *yaml.Node) error { - type BasicStepTmpl struct { - Act `yaml:",inline"` - Executor string `yaml:"executor,omitempty"` - Inline string `yaml:"inline,flow"` - Args []string `yaml:"args,omitempty,flow"` - CleanupStep yaml.Node `yaml:"cleanup,omitempty"` - } - - var tmpl BasicStepTmpl - // there is an issue with strict fields not being managed https://github.com/go-yaml/yaml/issues/460 - if err := node.Decode(&tmpl); err != nil { - return err - } - - b.Act = &tmpl.Act - b.Args = tmpl.Args - b.Executor = tmpl.Executor - b.Inline = tmpl.Inline - - if b.IsNil() { - return b.ExplainInvalid() - } - - // we do it piecemiel to build our struct - if tmpl.CleanupStep.IsZero() || b.Type == StepCleanup { - return nil - } - - logging.L().Debugw("step", "name", tmpl.Name) - cleanup, err := b.MakeCleanupStep(&tmpl.CleanupStep) - logging.L().Debugw("step", "err", err) - if err != nil { - return err - } - - b.CleanupStep = cleanup - - return nil -} - -// Cleanup is an implementation of the CleanupAct interface's Cleanup method. -func (b *BasicStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { - result, err := b.Execute(execCtx) - if err != nil { - return nil, err - } - return &result.ActResult, err -} - -// GetCleanup returns the cleanup steps for a BasicStep. -func (b *BasicStep) GetCleanup() []CleanupAct { - if b.CleanupStep != nil { - return []CleanupAct{b.CleanupStep} - } - return []CleanupAct{} -} - -// GetType returns the step type for a BasicStep. -func (b *BasicStep) GetType() StepType { - return b.Type -} - -// ExplainInvalid returns an error with an explanation of why a BasicStep is invalid. -func (b *BasicStep) ExplainInvalid() error { - var err error - if b.Inline == "" { - err = fmt.Errorf("(inline) empty") - } - if b.Name != "" && err != nil { - return fmt.Errorf("[!] invalid basicstep: [%s] %w", b.Name, err) - } - return err + return &BasicStep{} } // IsNil checks if a BasicStep is considered empty or uninitialized. func (b *BasicStep) IsNil() bool { switch { - case b.Act.IsNil(): - return true case b.Inline == "": return true default: @@ -142,12 +71,6 @@ func (b *BasicStep) IsNil() bool { // Validate validates the BasicStep, checking for the necessary attributes and dependencies. func (b *BasicStep) Validate(execCtx TTPExecutionContext) error { - // Validate Act - if err := b.Act.Validate(); err != nil { - logging.L().Error(zap.Error(err)) - return err - } - // Check if Inline is provided if b.Inline == "" { err := errors.New("inline must be provided") @@ -172,21 +95,13 @@ func (b *BasicStep) Validate(execCtx TTPExecutionContext) error { return err } - // Validate CleanupStep if it is not nil - if b.CleanupStep != nil { - if err := b.CleanupStep.Validate(execCtx); err != nil { - logging.L().Errorw("error validating cleanup step", zap.Error(err)) - return err - } - } - logging.L().Debugw("command found in path", "executor", b.Executor) return nil } // Execute runs the BasicStep and returns an error if any occur. -func (b *BasicStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (b *BasicStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { ctx, cancel := context.WithTimeout(context.Background(), 100*time.Minute) defer cancel() @@ -206,7 +121,7 @@ func (b *BasicStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, erro return result, nil } -func (b *BasicStep) executeBashStdin(ptx context.Context, execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (b *BasicStep) executeBashStdin(ptx context.Context, execCtx TTPExecutionContext) (*ActResult, error) { ctx, cancel := context.WithCancel(ptx) defer cancel() @@ -224,7 +139,7 @@ func (b *BasicStep) executeBashStdin(ptx context.Context, execCtx TTPExecutionCo return nil, err } - cmd := b.prepareCommand(ctx, expandedEnvAsList, expandedStrs[0]) + cmd := b.prepareCommand(ctx, execCtx, expandedEnvAsList, expandedStrs[0]) result, err := streamAndCapture(*cmd, execCtx.Cfg.Stdout, execCtx.Cfg.Stderr) if err != nil { @@ -238,10 +153,10 @@ func (b *BasicStep) executeBashStdin(ptx context.Context, execCtx TTPExecutionCo return result, nil } -func (b *BasicStep) prepareCommand(ctx context.Context, envAsList []string, inline string) *exec.Cmd { +func (b *BasicStep) prepareCommand(ctx context.Context, execCtx TTPExecutionContext, envAsList []string, inline string) *exec.Cmd { cmd := exec.CommandContext(ctx, b.Executor) cmd.Env = envAsList - cmd.Dir = b.WorkDir + cmd.Dir = execCtx.WorkDir cmd.Stdin = strings.NewReader(inline) return cmd diff --git a/pkg/blocks/basicstep_test.go b/pkg/blocks/basicstep_test.go index ff84a75a..f67ed137 100755 --- a/pkg/blocks/basicstep_test.go +++ b/pkg/blocks/basicstep_test.go @@ -61,18 +61,6 @@ steps: `, wantError: false, }, - { - name: "Invalid basic", - content: ` -name: test -description: this is a test -steps: - - noname: testinline - inline: | - ls - `, - wantError: true, - }, } for _, tc := range testCases { diff --git a/pkg/blocks/context.go b/pkg/blocks/context.go index 9f99a76d..9b1b3d3d 100644 --- a/pkg/blocks/context.go +++ b/pkg/blocks/context.go @@ -45,6 +45,7 @@ type TTPExecutionConfig struct { // TTPExecutionContext - holds config and context for the currently executing TTP type TTPExecutionContext struct { Cfg TTPExecutionConfig + WorkDir string StepResults *StepResultsRecord } diff --git a/pkg/blocks/createfile.go b/pkg/blocks/createfile.go index ab5e4214..47841442 100755 --- a/pkg/blocks/createfile.go +++ b/pkg/blocks/createfile.go @@ -20,14 +20,11 @@ THE SOFTWARE. package blocks import ( - "errors" "fmt" "os" "github.com/facebookincubator/ttpforge/pkg/logging" "github.com/spf13/afero" - "go.uber.org/zap" - "gopkg.in/yaml.v3" ) // CreateFileStep creates a new file and populates it @@ -36,122 +33,21 @@ import ( // through an editor program or via a C2, where there is no // corresponding shell history telemetry type CreateFileStep struct { - *Act `yaml:",inline"` - Path string `yaml:"create_file,omitempty"` - Contents string `yaml:"contents,omitempty"` - Overwrite bool `yaml:"overwrite,omitempty"` - Mode int `yaml:"mode,omitempty"` - CleanupStep CleanupAct `yaml:"cleanup,omitempty,flow"` - FileSystem afero.Fs `yaml:"-,omitempty"` + Path string `yaml:"create_file,omitempty"` + Contents string `yaml:"contents,omitempty"` + Overwrite bool `yaml:"overwrite,omitempty"` + Mode int `yaml:"mode,omitempty"` + FileSystem afero.Fs `yaml:"-,omitempty"` } // NewCreateFileStep creates a new CreateFileStep instance and returns a pointer to it. func NewCreateFileStep() *CreateFileStep { - return &CreateFileStep{ - Act: &Act{ - Type: StepCreateFile, - }, - } -} - -// UnmarshalYAML decodes a YAML node into a CreateFileStep instance. It uses -// the provided struct as a template for the YAML data, and initializes the -// CreateFileStep instance with the decoded values. -// -// **Parameters:** -// -// node: A pointer to a yaml.Node representing the YAML data to decode. -// -// **Returns:** -// -// error: An error if there is a problem decoding the YAML data. -func (s *CreateFileStep) UnmarshalYAML(node *yaml.Node) error { - - type createFileStepTmpl struct { - Act `yaml:",inline"` - Path string `yaml:"create_file,omitempty"` - Contents string `yaml:"contents,omitempty"` - Overwrite bool `yaml:"overwrite,omitempty"` - Mode int `yaml:"mode,omitempty"` - CleanupStep yaml.Node `yaml:"cleanup,omitempty,flow"` - } - - // Decode the YAML node into the provided template. - var tmpl createFileStepTmpl - if err := node.Decode(&tmpl); err != nil { - return err - } - - // Initialize the instance with the decoded values. - s.Act = &tmpl.Act - s.Path = tmpl.Path - s.Contents = tmpl.Contents - s.Overwrite = tmpl.Overwrite - s.Mode = tmpl.Mode - - // Check for invalid steps. - if s.IsNil() { - return s.ExplainInvalid() - } - - // If there is no cleanup step or if this step is the cleanup step, exit. - if tmpl.CleanupStep.IsZero() || s.Type == StepCleanup { - return nil - } - - // Create a CleanupStep instance and add it to this step - cleanup, err := s.MakeCleanupStep(&tmpl.CleanupStep) - if err != nil { - return err - } - - s.CleanupStep = cleanup - - return nil -} - -// GetType returns the type of the step as StepType. -func (s *CreateFileStep) GetType() StepType { - return StepCreateFile -} - -// Cleanup is a method to establish a link with the Cleanup interface. -// Assumes that the type is the cleanup step and is invoked by -// s.CleanupStep.Cleanup. -func (s *CreateFileStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { - result, err := s.Execute(execCtx) - if err != nil { - return nil, err - } - return &result.ActResult, err -} - -// GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. -func (s *CreateFileStep) GetCleanup() []CleanupAct { - if s.CleanupStep != nil { - return []CleanupAct{s.CleanupStep} - } - return []CleanupAct{} -} - -// ExplainInvalid returns an error message explaining why the step -// is invalid. -// -// **Returns:** -// -// error: An error message explaining why the step is invalid. -func (s *CreateFileStep) ExplainInvalid() error { - if s.Path == "" { - return errors.New("empty `create_file:` provided") - } - return nil + return &CreateFileStep{} } // IsNil checks if the step is nil or empty and returns a boolean value. func (s *CreateFileStep) IsNil() bool { switch { - case s.Act.IsNil(): - return true case s.Path == "": return true default: @@ -160,7 +56,7 @@ func (s *CreateFileStep) IsNil() bool { } // Execute runs the step and returns an error if any occur. -func (s *CreateFileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (s *CreateFileStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { logging.L().Infof("Creating file %v", s.Path) fsys := s.FileSystem if fsys == nil { @@ -191,7 +87,15 @@ func (s *CreateFileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, return nil, err } - return &ExecutionResult{}, nil + return &ActResult{}, nil +} + +// GetDefaultCleanupAction will instruct the calling code +// to remove the path created by this action +func (s *CreateFileStep) GetDefaultCleanupAction() Action { + return &RemovePathAction{ + Path: s.Path, + } } // Validate validates the step @@ -200,20 +104,8 @@ func (s *CreateFileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, // // error: An error if any validation checks fail. func (s *CreateFileStep) Validate(execCtx TTPExecutionContext) error { - if err := s.Act.Validate(); err != nil { - return err - } - if s.Path == "" { return fmt.Errorf("path field cannot be empty") } - - if s.CleanupStep != nil { - if err := s.CleanupStep.Validate(execCtx); err != nil { - logging.L().Errorw("error validating cleanup step", zap.Error(err)) - return err - } - } - return nil } diff --git a/pkg/blocks/editstep.go b/pkg/blocks/editstep.go index e2c642ef..7741fb58 100755 --- a/pkg/blocks/editstep.go +++ b/pkg/blocks/editstep.go @@ -37,39 +37,21 @@ type Edit struct { // EditStep represents one or more edits to a specific file type EditStep struct { - *Act `yaml:",inline"` - FileToEdit string `yaml:"edit_file,omitempty"` - Edits []*Edit `yaml:"edits,omitempty"` - FileSystem afero.Fs `yaml:"-,omitempty"` - BackupFile string `yaml:"backup_file,omitempty"` + actionDefaults `yaml:"-"` + FileToEdit string `yaml:"edit_file,omitempty"` + Edits []*Edit `yaml:"edits,omitempty"` + FileSystem afero.Fs `yaml:"-,omitempty"` + BackupFile string `yaml:"backup_file,omitempty"` } // NewEditStep creates a new EditStep instance with an initialized Act struct. func NewEditStep() *EditStep { - return &EditStep{ - Act: &Act{ - Type: StepEdit, - }, - } -} - -// GetCleanup returns the cleanup steps for a EditStep. -// Currently this is always empty because we use backup -// files instead for this type of step -func (s *EditStep) GetCleanup() []CleanupAct { - return []CleanupAct{} -} - -// GetType returns the step type for a EditStep. -func (s *EditStep) GetType() StepType { - return s.Type + return &EditStep{} } // IsNil checks if an EditStep is considered empty or uninitialized. func (s *EditStep) IsNil() bool { switch { - case s.Act.IsNil(): - return true case s.FileToEdit == "": return true default: @@ -77,15 +59,8 @@ func (s *EditStep) IsNil() bool { } } -// wrapped by exported Validate to standardize -// the error message prefix -func (s *EditStep) check() error { - // Validate Act - if err := s.Act.Validate(); err != nil { - return err - } - - var err error +// Validate validates the EditStep, checking for the necessary attributes and dependencies. +func (s *EditStep) Validate(execCtx TTPExecutionContext) error { if len(s.Edits) == 0 { return fmt.Errorf("no edits specified") } @@ -98,6 +73,7 @@ func (s *EditStep) check() error { return fmt.Errorf("edit #%d is missing 'new:'", editIdx+1) } + var err error if edit.Regexp { edit.oldRegexp, err = regexp.Compile(edit.Old) if err != nil { @@ -110,23 +86,14 @@ func (s *EditStep) check() error { return nil } -// Validate validates the EditStep, checking for the necessary attributes and dependencies. -func (s *EditStep) Validate(execCtx TTPExecutionContext) error { - err := s.check() - if err != nil { - return fmt.Errorf("[!] invalid editstep: [%s] %w", s.Name, err) - } - return nil -} - // Execute runs the EditStep and returns an error if any occur. -func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { fileSystem := s.FileSystem targetPath := s.FileToEdit if fileSystem == nil { fileSystem = afero.NewOsFs() var err error - targetPath, err = FetchAbs(targetPath, s.WorkDir) + targetPath, err = FetchAbs(targetPath, execCtx.WorkDir) if err != nil { return nil, err } @@ -169,5 +136,5 @@ func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error return nil, err } - return &ExecutionResult{}, nil + return &ActResult{}, nil } diff --git a/pkg/blocks/editstep_test.go b/pkg/blocks/editstep_test.go index a07e8548..eec0baea 100755 --- a/pkg/blocks/editstep_test.go +++ b/pkg/blocks/editstep_test.go @@ -45,7 +45,7 @@ steps: err := yaml.Unmarshal([]byte(content), &ttp) require.NoError(t, err) - err = ttp.ValidateSteps(blocks.TTPExecutionContext{}) + err = ttp.Validate(blocks.TTPExecutionContext{}) require.NoError(t, err) } @@ -66,10 +66,10 @@ steps: err := yaml.Unmarshal([]byte(content), &ttp) require.NoError(t, err) - err = ttp.ValidateSteps(blocks.TTPExecutionContext{}) + err = ttp.Validate(blocks.TTPExecutionContext{}) require.Error(t, err) - assert.Equal(t, "[!] invalid editstep: [missing_new] edit #2 is missing 'new:'", err.Error()) + assert.Equal(t, "edit #2 is missing 'new:'", err.Error()) } func TestUnmarshalEditNoOld(t *testing.T) { @@ -89,10 +89,10 @@ steps: err := yaml.Unmarshal([]byte(content), &ttp) require.NoError(t, err) - err = ttp.ValidateSteps(blocks.TTPExecutionContext{}) + err = ttp.Validate(blocks.TTPExecutionContext{}) require.Error(t, err) - assert.Equal(t, "[!] invalid editstep: [missing_old] edit #1 is missing 'old:'", err.Error()) + assert.Equal(t, "edit #1 is missing 'old:'", err.Error()) } func TestUnmarshalNonListEdits(t *testing.T) { @@ -119,8 +119,8 @@ steps: err := yaml.Unmarshal([]byte(content), &ttp) require.NoError(t, err) - err = ttp.ValidateSteps(blocks.TTPExecutionContext{}) - assert.Equal(t, "[!] invalid editstep: [no_edits] no edits specified", err.Error()) + err = ttp.Validate(blocks.TTPExecutionContext{}) + assert.Equal(t, "no edits specified", err.Error()) } func TestExecuteSimple(t *testing.T) { diff --git a/pkg/blocks/fetchuri.go b/pkg/blocks/fetchuri.go index a056edd8..665e5d11 100755 --- a/pkg/blocks/fetchuri.go +++ b/pkg/blocks/fetchuri.go @@ -31,144 +31,35 @@ import ( "github.com/facebookincubator/ttpforge/pkg/logging" "github.com/spf13/afero" "go.uber.org/zap" - "gopkg.in/yaml.v3" ) // FetchURIStep represents a step in a process that consists of a main action, // a cleanup action, and additional metadata. type FetchURIStep struct { - *Act `yaml:",inline"` - FetchURI string `yaml:"fetch_uri,omitempty"` - Retries string `yaml:"retries,omitempty"` - Location string `yaml:"location,omitempty"` - Proxy string `yaml:"proxy,omitempty"` - Overwrite bool `yaml:"overwrite,omitempty"` - CleanupStep CleanupAct `yaml:"cleanup,omitempty,flow"` - FileSystem afero.Fs `yaml:"-,omitempty"` + actionDefaults `yaml:"-"` + FetchURI string `yaml:"fetch_uri,omitempty"` + Retries string `yaml:"retries,omitempty"` + Location string `yaml:"location,omitempty"` + Proxy string `yaml:"proxy,omitempty"` + Overwrite bool `yaml:"overwrite,omitempty"` + FileSystem afero.Fs `yaml:"-,omitempty"` } // NewFetchURIStep creates a new FetchURIStep instance and returns a pointer to it. func NewFetchURIStep() *FetchURIStep { - return &FetchURIStep{ - Act: &Act{ - Type: StepFetchURI, - }, - } -} - -// UnmarshalYAML decodes a YAML node into a FetchURIStep instance. It uses -// the provided struct as a template for the YAML data, and initializes the -// FetchURIStep instance with the decoded values. -// -// **Parameters:** -// -// node: A pointer to a yaml.Node representing the YAML data to decode. -// -// **Returns:** -// -// error: An error if there is a problem decoding the YAML data. -func (f *FetchURIStep) UnmarshalYAML(node *yaml.Node) error { - - type fileStepTmpl struct { - Act `yaml:",inline"` - FetchURI string `yaml:"fetch_uri,omitempty"` - Retries string `yaml:"retries,omitempty"` - Location string `yaml:"location,omitempty"` - Proxy string `yaml:"proxy,omitempty"` - Overwrite bool `yaml:"overwrite,omitempty"` - CleanupStep yaml.Node `yaml:"cleanup,omitempty,flow"` - } - - // Decode the YAML node into the provided template. - var tmpl fileStepTmpl - if err := node.Decode(&tmpl); err != nil { - return err - } - - // Initialize the FetchURIStep instance with the decoded values. - f.Act = &tmpl.Act - f.FetchURI = tmpl.FetchURI - f.Location = tmpl.Location - f.Retries = tmpl.Retries - f.Proxy = tmpl.Proxy - f.Overwrite = tmpl.Overwrite - - // Check for invalid steps. - if f.IsNil() { - return f.ExplainInvalid() - } - - // If there is no cleanup step or if this step is the cleanup step, exit. - if tmpl.CleanupStep.IsZero() || f.Type == StepCleanup { - return nil - } - - // Create a CleanupStep instance and add it to the FetchURIStep instance. - logging.L().Debugw("step", "name", tmpl.Name) - cleanup, err := f.MakeCleanupStep(&tmpl.CleanupStep) - logging.L().Debugw("step", zap.Error(err)) - if err != nil { - logging.L().Errorw("error creating cleanup step", zap.Error(err)) - return err - } - - f.CleanupStep = cleanup - - return nil -} - -// GetType returns the type of the step as StepType. -func (f *FetchURIStep) GetType() StepType { - return StepFetchURI + return &FetchURIStep{} } // Cleanup is a method to establish a link with the Cleanup interface. // Assumes that the type is the cleanup step and is invoked by // f.CleanupStep.Cleanup. func (f *FetchURIStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { - result, err := f.Execute(execCtx) - if err != nil { - return nil, err - } - return &result.ActResult, err -} - -// GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. -func (f *FetchURIStep) GetCleanup() []CleanupAct { - if f.CleanupStep != nil { - return []CleanupAct{f.CleanupStep} - } - return []CleanupAct{} -} - -// ExplainInvalid returns an error message explaining why the FetchURIStep -// is invalid. -// -// **Returns:** -// -// error: An error message explaining why the FetchURIStep is invalid. -func (f *FetchURIStep) ExplainInvalid() error { - var err error - if f.FetchURI == "" { - err = errors.New("empty FetchURI provided") - } - - if f.Location == "" && err != nil { - err = errors.New("empty Location provided") - } - - if f.Name != "" && err != nil { - err = fmt.Errorf("[!] invalid FetchURIStep: [%s] %v", f.Name, zap.Error(err)) - } - - return err + return f.Execute(execCtx) } // IsNil checks if the FetchURIStep is nil or empty and returns a boolean value. func (f *FetchURIStep) IsNil() bool { switch { - case f.Act.IsNil(): - return true case f.FetchURI == "": return true case f.Location == "": @@ -179,7 +70,7 @@ func (f *FetchURIStep) IsNil() bool { } // Execute runs the FetchURIStep and returns an error if any occur. -func (f *FetchURIStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (f *FetchURIStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { logging.L().Info("========= Executing ==========") if err := f.fetchURI(execCtx); err != nil { @@ -189,7 +80,7 @@ func (f *FetchURIStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, e logging.L().Info("========= Result ==========") - return &ExecutionResult{}, nil + return &ActResult{}, nil } // fetchURI executes the FetchURIStep with the specified Location, Uri, and additional arguments, @@ -201,7 +92,7 @@ func (f *FetchURIStep) fetchURI(execCtx TTPExecutionContext) error { if appFs == nil { var err error appFs = afero.NewOsFs() - absLocal, err = FetchAbs(f.Location, f.WorkDir) + absLocal, err = FetchAbs(f.Location, execCtx.WorkDir) if err != nil { return err } @@ -259,11 +150,6 @@ func (f *FetchURIStep) fetchURI(execCtx TTPExecutionContext) error { // // error: An error if any validation checks fail. func (f *FetchURIStep) Validate(execCtx TTPExecutionContext) error { - if err := f.Act.Validate(); err != nil { - logging.L().Error(zap.Error(err)) - return err - } - if f.FetchURI == "" { err := errors.New("require FetchURI to be set with fetchURI") logging.L().Error(zap.Error(err)) @@ -286,7 +172,7 @@ func (f *FetchURIStep) Validate(execCtx TTPExecutionContext) error { } // Retrieve the absolute path to the file. - absLocal, err := FetchAbs(f.Location, f.WorkDir) + absLocal, err := FetchAbs(f.Location, execCtx.WorkDir) if err != nil { logging.L().Error(zap.Error(err)) return err @@ -297,13 +183,5 @@ func (f *FetchURIStep) Validate(execCtx TTPExecutionContext) error { logging.L().Errorw("FileStep location exists, remove and retry", "location", absLocal) return errors.New("file exists at location specified, remove and retry") } - - if f.CleanupStep != nil { - if err := f.CleanupStep.Validate(execCtx); err != nil { - logging.L().Errorw("error validating cleanup step", zap.Error(err)) - return err - } - } - return nil } diff --git a/pkg/blocks/fetchuri_test.go b/pkg/blocks/fetchuri_test.go index e70966ac..d59a2024 100755 --- a/pkg/blocks/fetchuri_test.go +++ b/pkg/blocks/fetchuri_test.go @@ -189,7 +189,7 @@ steps: }, } - err = ttps.ValidateSteps(execCtx) + err = ttps.Validate(execCtx) if tc.wantError { assert.Error(t, err) } else { diff --git a/pkg/blocks/filestep.go b/pkg/blocks/filestep.go index fef3cb34..63c9d577 100755 --- a/pkg/blocks/filestep.go +++ b/pkg/blocks/filestep.go @@ -21,7 +21,6 @@ package blocks import ( "errors" - "fmt" "os/exec" "path/filepath" "runtime" @@ -29,133 +28,34 @@ import ( "github.com/facebookincubator/ttpforge/pkg/logging" "github.com/facebookincubator/ttpforge/pkg/outputs" "go.uber.org/zap" - "gopkg.in/yaml.v3" ) // FileStep represents a step in a process that consists of a main action, // a cleanup action, and additional metadata. type FileStep struct { - *Act `yaml:",inline"` - FilePath string `yaml:"file,omitempty"` - Executor string `yaml:"executor,omitempty"` - CleanupStep CleanupAct `yaml:"cleanup,omitempty,flow"` - Args []string `yaml:"args,omitempty,flow"` + actionDefaults `yaml:"-"` + FilePath string `yaml:"file,omitempty"` + Executor string `yaml:"executor,omitempty"` + Environment map[string]string `yaml:"env,omitempty"` + Outputs map[string]outputs.Spec `yaml:"outputs,omitempty"` + Args []string `yaml:"args,omitempty,flow"` } // NewFileStep creates a new FileStep instance and returns a pointer to it. func NewFileStep() *FileStep { - return &FileStep{ - Act: &Act{ - Type: StepFile, - }, - } -} - -// UnmarshalYAML decodes a YAML node into a FileStep instance. It uses -// the provided struct as a template for the YAML data, and initializes the -// FileStep instance with the decoded values. -// -// **Parameters:** -// -// node: A pointer to a yaml.Node representing the YAML data to decode. -// -// **Returns:** -// -// error: An error if there is a problem decoding the YAML data. -func (f *FileStep) UnmarshalYAML(node *yaml.Node) error { - - type fileStepTmpl struct { - Act `yaml:",inline"` - FilePath string `yaml:"file,omitempty"` - Executor string `yaml:"executor,omitempty"` - CleanupStep yaml.Node `yaml:"cleanup,omitempty,flow"` - Args []string `yaml:"args,omitempty,flow"` - } - - // Decode the YAML node into the provided template. - var tmpl fileStepTmpl - if err := node.Decode(&tmpl); err != nil { - return err - } - - // Initialize the FileStep instance with the decoded values. - f.Act = &tmpl.Act - f.Args = tmpl.Args - f.FilePath = tmpl.FilePath - f.Executor = tmpl.Executor - - // Check for invalid steps. - if f.IsNil() { - return f.ExplainInvalid() - } - - // If there is no cleanup step or if this step is the cleanup step, exit. - if tmpl.CleanupStep.IsZero() || f.Type == StepCleanup { - return nil - } - - // Create a CleanupStep instance and add it to the FileStep instance. - logging.L().Debugw("step", "name", tmpl.Name) - cleanup, err := f.MakeCleanupStep(&tmpl.CleanupStep) - logging.L().Debugw("step", zap.Error(err)) - if err != nil { - logging.L().Errorw("error creating cleanup step", zap.Error(err)) - return err - } - - f.CleanupStep = cleanup - - return nil -} - -// GetType returns the type of the step as StepType. -func (f *FileStep) GetType() StepType { - return StepFile + return &FileStep{} } // Cleanup is a method to establish a link with the Cleanup interface. // Assumes that the type is the cleanup step and is invoked by // f.CleanupStep.Cleanup. func (f *FileStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { - result, err := f.Execute(execCtx) - if err != nil { - return nil, err - } - return &result.ActResult, err -} - -// GetCleanup returns a slice of CleanupAct if the CleanupStep is not nil. -func (f *FileStep) GetCleanup() []CleanupAct { - if f.CleanupStep != nil { - return []CleanupAct{f.CleanupStep} - } - return []CleanupAct{} -} - -// ExplainInvalid returns an error message explaining why the FileStep -// is invalid. -// -// **Returns:** -// -// error: An error message explaining why the FileStep is invalid. -func (f *FileStep) ExplainInvalid() error { - var err error - if f.FilePath == "" { - err = errors.New("empty FilePath provided") - } - - if f.Name != "" && err != nil { - err = fmt.Errorf("[!] invalid FileStep: [%s] %v", f.Name, zap.Error(err)) - } - - return err + return f.Execute(execCtx) } // IsNil checks if the FileStep is nil or empty and returns a boolean value. func (f *FileStep) IsNil() bool { switch { - case f.Act.IsNil(): - return true case f.FilePath == "": return true default: @@ -164,7 +64,7 @@ func (f *FileStep) IsNil() bool { } // Execute runs the FileStep and returns an error if any occur. -func (f *FileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { +func (f *FileStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { var cmd *exec.Cmd expandedArgs, err := execCtx.ExpandVariables(f.Args) if err != nil { @@ -185,13 +85,13 @@ func (f *FileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error return nil, err } cmd.Env = expandedEnvAsList - cmd.Dir = f.WorkDir + cmd.Dir = execCtx.WorkDir result, err := streamAndCapture(*cmd, execCtx.Cfg.Stdout, execCtx.Cfg.Stderr) if err != nil { return nil, err } result.Outputs, err = outputs.Parse(f.Outputs, result.Stdout) - return nil, err + return result, err } // Validate validates the FileStep. It checks that the @@ -210,11 +110,6 @@ func (f *FileStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error // // error: An error if any validation checks fail. func (f *FileStep) Validate(execCtx TTPExecutionContext) error { - if err := f.Act.Validate(); err != nil { - logging.L().Error(zap.Error(err)) - return err - } - if f.FilePath == "" { err := errors.New("a TTP must include inline logic or path to a file with the logic") logging.L().Error(zap.Error(err)) @@ -222,14 +117,14 @@ func (f *FileStep) Validate(execCtx TTPExecutionContext) error { } // If FilePath is set, ensure that the file exists. - fullpath, err := FindFilePath(f.FilePath, f.WorkDir, nil) + fullpath, err := FindFilePath(f.FilePath, execCtx.WorkDir, nil) if err != nil { logging.L().Error(zap.Error(err)) return err } // Retrieve the absolute path to the file. - f.FilePath, err = FetchAbs(fullpath, f.WorkDir) + f.FilePath, err = FetchAbs(fullpath, execCtx.WorkDir) if err != nil { logging.L().Error(zap.Error(err)) return err @@ -238,7 +133,7 @@ func (f *FileStep) Validate(execCtx TTPExecutionContext) error { // Infer executor if it's not set. if f.Executor == "" { f.Executor = InferExecutor(f.FilePath) - logging.L().Infow("executor set via extension", "exec", f.Executor) + logging.L().Debugw("executor set via extension", "exec", f.Executor) } if f.Executor == ExecutorBinary { @@ -249,13 +144,6 @@ func (f *FileStep) Validate(execCtx TTPExecutionContext) error { logging.L().Error(zap.Error(err)) return err } - - if f.CleanupStep != nil { - if err := f.CleanupStep.Validate(execCtx); err != nil { - logging.L().Errorw("error validating cleanup step", zap.Error(err)) - return err - } - } logging.L().Debugw("command found in path", "executor", f.Executor) return nil diff --git a/pkg/blocks/iocapture.go b/pkg/blocks/iocapture.go index 81355491..f9ee4916 100644 --- a/pkg/blocks/iocapture.go +++ b/pkg/blocks/iocapture.go @@ -26,7 +26,7 @@ import ( "os/exec" ) -func streamAndCapture(cmd exec.Cmd, stdout, stderr io.Writer) (*ExecutionResult, error) { +func streamAndCapture(cmd exec.Cmd, stdout, stderr io.Writer) (*ActResult, error) { if stdout == nil { stdout = os.Stdout } @@ -43,7 +43,7 @@ func streamAndCapture(cmd exec.Cmd, stdout, stderr io.Writer) (*ExecutionResult, return nil, err } outStr, errStr := stdoutBuf.String(), stderrBuf.String() - result := ExecutionResult{} + result := ActResult{} result.Stdout = outStr result.Stderr = errStr if err != nil { diff --git a/pkg/blocks/loader.go b/pkg/blocks/loader.go index e88a22ed..8189b40d 100755 --- a/pkg/blocks/loader.go +++ b/pkg/blocks/loader.go @@ -78,18 +78,19 @@ func RenderTemplatedTTP(ttpStr string, execCfg *TTPExecutionConfig) (*TTP, error // // **Returns:** // -// ttp: Pointer to the created TTP instance, or nil if the file is empty or invalid. +// *TTP: Pointer to the created TTP instance, or nil if the file is empty or invalid. +// TTPExecutionContext: the initialized TTPExecutionContext suitable for passing to TTP.Execute(...) // err: An error if the file contains invalid data or cannot be read. -func LoadTTP(ttpFilePath string, fsys afero.Fs, execCfg *TTPExecutionConfig, argsKvStrs []string) (*TTP, error) { +func LoadTTP(ttpFilePath string, fsys afero.Fs, execCfg *TTPExecutionConfig, argsKvStrs []string) (*TTP, *TTPExecutionContext, error) { ttpBytes, err := readTTPBytes(ttpFilePath, fsys) if err != nil { - return nil, err + return nil, nil, err } result, err := preprocess.Parse(ttpBytes) if err != nil { - return nil, err + return nil, nil, err } // linting above establishes that the TTP yaml will be @@ -100,18 +101,18 @@ func LoadTTP(ttpFilePath string, fsys afero.Fs, execCfg *TTPExecutionConfig, arg var tmpContainer ArgSpecContainer err = yaml.Unmarshal(result.PreambleBytes, &tmpContainer) if err != nil { - return nil, err + return nil, nil, err } argValues, err := args.ParseAndValidate(tmpContainer.ArgSpecs, argsKvStrs) if err != nil { - return nil, fmt.Errorf("failed to parse and validate arguments: %v", err) + return nil, nil, fmt.Errorf("failed to parse and validate arguments: %v", err) } execCfg.Args = argValues ttp, err := RenderTemplatedTTP(string(ttpBytes), execCfg) if err != nil { - return nil, err + return nil, nil, err } // embedded fs has no notion of workdirs @@ -121,29 +122,26 @@ func LoadTTP(ttpFilePath string, fsys afero.Fs, execCfg *TTPExecutionConfig, arg case *afero.OsFs: absPath, err := filepath.Abs(ttpFilePath) if err != nil { - return nil, err + return nil, nil, err } ttp.WorkDir = filepath.Dir(absPath) default: wd, err := os.Getwd() if err != nil { - return nil, err + return nil, nil, err } ttp.WorkDir = wd } - - // TODO: refactor directory handling - this is in-elegant - // but has less bugs than previous way - for _, step := range ttp.Steps { - step.SetDir(ttp.WorkDir) - if cleanups := step.GetCleanup(); cleanups != nil { - for _, c := range cleanups { - c.SetDir(ttp.WorkDir) - } - } + execCtx := &TTPExecutionContext{ + Cfg: *execCfg, + WorkDir: ttp.WorkDir, } - return ttp, nil + err = ttp.Validate(*execCtx) + if err != nil { + return nil, nil, err + } + return ttp, execCtx, nil } func readTTPBytes(ttpFilePath string, system afero.Fs) ([]byte, error) { diff --git a/pkg/blocks/printstr.go b/pkg/blocks/printstr.go new file mode 100755 index 00000000..32a67295 --- /dev/null +++ b/pkg/blocks/printstr.go @@ -0,0 +1,75 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks + +import ( + "bytes" + "fmt" + "io" + "os" +) + +// PrintStrAction is used to print a string to the console +type PrintStrAction struct { + actionDefaults `yaml:"-"` + Message string `yaml:"print_str,omitempty"` +} + +// IsNil checks if the step is nil or empty and returns a boolean value. +func (s *PrintStrAction) IsNil() bool { + switch { + case s.Message == "": + return true + default: + return false + } +} + +// Execute runs the step and returns an error if any occur. +func (s *PrintStrAction) Execute(execCtx TTPExecutionContext) (*ActResult, error) { + // needs to be overwritable to capture output during testing + stdout := execCtx.Cfg.Stdout + if stdout == nil { + stdout = os.Stdout + } + expandedStrs, err := execCtx.ExpandVariables([]string{s.Message}) + if err != nil { + return nil, err + } + var stdoutBuf bytes.Buffer + multi := io.MultiWriter(stdout, &stdoutBuf) + fmt.Fprintln(multi, expandedStrs[0]) + result := &ActResult{ + Stdout: stdoutBuf.String(), + } + return result, nil +} + +// Validate validates the step +// +// **Returns:** +// +// error: An error if any validation checks fail. +func (s *PrintStrAction) Validate(execCtx TTPExecutionContext) error { + if s.Message == "" { + return fmt.Errorf("message field cannot be empty") + } + return nil +} diff --git a/pkg/blocks/printstr_test.go b/pkg/blocks/printstr_test.go new file mode 100755 index 00000000..e9dc6ac9 --- /dev/null +++ b/pkg/blocks/printstr_test.go @@ -0,0 +1,83 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks_test + +import ( + "testing" + + "github.com/facebookincubator/ttpforge/pkg/blocks" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPrintStrExecute(t *testing.T) { + testCases := []struct { + name string + description string + action *blocks.PrintStrAction + expectExecuteError bool + expectedStdout string + }{ + { + name: "Simple Print", + description: "Just Print a String", + action: &blocks.PrintStrAction{ + Message: "hello", + }, + expectedStdout: "hello\n", + }, + { + name: "Print Step Output", + description: "Should be Expanded", + action: &blocks.PrintStrAction{ + Message: "value is $forge.steps.first_step.stdout", + }, + expectedStdout: "value is first-step-output\n", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // for use testing output variables + execCtx := blocks.TTPExecutionContext{ + StepResults: &blocks.StepResultsRecord{ + ByName: map[string]*blocks.ExecutionResult{ + "first_step": { + ActResult: blocks.ActResult{ + Stdout: "first-step-output", + }, + }, + }, + }, + } + + // execute and check error + result, err := tc.action.Execute(execCtx) + if tc.expectExecuteError { + require.Error(t, err) + return + } + require.NoError(t, err) + + // Check stdout + assert.Equal(t, tc.expectedStdout, result.Stdout) + }) + } +} diff --git a/pkg/blocks/removepath.go b/pkg/blocks/removepath.go new file mode 100755 index 00000000..191ff186 --- /dev/null +++ b/pkg/blocks/removepath.go @@ -0,0 +1,99 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks + +import ( + "fmt" + + "github.com/facebookincubator/ttpforge/pkg/logging" + "github.com/spf13/afero" +) + +// RemovePathAction is invoked by +// adding remove_path to a given YAML step. +// It will delete the file at the specified path +// You must pass `recursive: true` to delete directories +type RemovePathAction struct { + actionDefaults `yaml:"-"` + Path string `yaml:"remove_path,omitempty"` + Recursive bool `yaml:"recursive,omitempty"` + FileSystem afero.Fs `yaml:"-,omitempty"` +} + +// IsNil checks if the step is nil or empty and returns a boolean value. +func (s *RemovePathAction) IsNil() bool { + switch { + case s.Path == "": + return true + default: + return false + } +} + +// Execute runs the step and returns an error if any occur. +func (s *RemovePathAction) Execute(execCtx TTPExecutionContext) (*ActResult, error) { + logging.L().Infof("Removing path %v", s.Path) + fsys := s.FileSystem + if fsys == nil { + fsys = afero.NewOsFs() + } + + // cannot remove a non-existent path + exists, err := afero.Exists(fsys, s.Path) + if err != nil { + return nil, err + } + if !exists { + return nil, fmt.Errorf("path %v does not exist", s.Path) + } + + // afero fsys.Remove(...) appears to be buggy + // and will remove a directory even if it is not empty + // so we check manually - we use the semantics + // of the macOS `rm` command and refuse to remove even + // empty directories unless recursive is specified + isDir, err := afero.IsDir(fsys, s.Path) + if err != nil { + return nil, err + } + + if isDir && !s.Recursive { + return nil, fmt.Errorf("path %v is a directory and `recursive: true` was not specified - refusing to remove", s.Path) + } + + // actually remove the file + err = fsys.RemoveAll(s.Path) + if err != nil { + return nil, err + } + return &ActResult{}, nil +} + +// Validate validates the step +// +// **Returns:** +// +// error: An error if any validation checks fail. +func (s *RemovePathAction) Validate(execCtx TTPExecutionContext) error { + if s.Path == "" { + return fmt.Errorf("path field cannot be empty") + } + return nil +} diff --git a/pkg/blocks/removepath_test.go b/pkg/blocks/removepath_test.go new file mode 100755 index 00000000..d0e4b773 --- /dev/null +++ b/pkg/blocks/removepath_test.go @@ -0,0 +1,111 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks_test + +import ( + "testing" + + "github.com/facebookincubator/ttpforge/pkg/blocks" + "github.com/facebookincubator/ttpforge/pkg/testutils" + "github.com/spf13/afero" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRemovePathExecute(t *testing.T) { + testCases := []struct { + name string + description string + step *blocks.RemovePathAction + fsysContents map[string][]byte + expectExecuteError bool + }{ + { + name: "Remove Valid File", + description: "Remove a single unremarkable file", + step: &blocks.RemovePathAction{ + Path: "valid-file.txt", + }, + fsysContents: map[string][]byte{ + "valid-file.txt": []byte("whoops"), + }, + }, + { + name: "Remove Non-Existent File", + description: "Remove a non-existent file - should error", + step: &blocks.RemovePathAction{ + Path: "does-not-exist.txt", + }, + fsysContents: map[string][]byte{ + "valid-file.txt": []byte("whoops"), + }, + expectExecuteError: true, + }, + { + name: "Remove Directory - Success", + description: "Set Recursive to make directory removal succeed", + step: &blocks.RemovePathAction{ + Path: "valid-directory", + Recursive: true, + }, + fsysContents: map[string][]byte{ + "valid-directory/valid-file.txt": []byte("whoops"), + }, + }, + { + name: "Remove Directory - Failure", + description: "Refuse to remove directory because `recursive: true` was not specified", + step: &blocks.RemovePathAction{ + Path: "valid-directory", + }, + fsysContents: map[string][]byte{ + "valid-directory/valid-file.txt": []byte("whoops"), + }, + expectExecuteError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // prep filesystem + if tc.fsysContents != nil { + fsys, err := testutils.MakeAferoTestFs(tc.fsysContents) + require.NoError(t, err) + tc.step.FileSystem = fsys + } else { + tc.step.FileSystem = afero.NewMemMapFs() + } + + // execute and check error + var execCtx blocks.TTPExecutionContext + _, err := tc.step.Execute(execCtx) + if tc.expectExecuteError { + require.Error(t, err) + return + } + require.NoError(t, err) + + // Verify that the file is gone + exists, err := afero.Exists(tc.step.FileSystem, tc.step.Path) + require.NoError(t, err) + assert.False(t, exists) + }) + } +} diff --git a/pkg/blocks/step.go b/pkg/blocks/step.go old mode 100755 new mode 100644 index 1cbab450..f00c2e67 --- a/pkg/blocks/step.go +++ b/pkg/blocks/step.go @@ -22,241 +22,187 @@ package blocks import ( "errors" "fmt" - "runtime" - "strings" "github.com/facebookincubator/ttpforge/pkg/logging" - "github.com/facebookincubator/ttpforge/pkg/outputs" - "go.uber.org/zap" "gopkg.in/yaml.v3" ) -// Constants representing supported executor types. -const ( - ExecutorPython = "python3" - ExecutorBash = "bash" - ExecutorSh = "sh" - ExecutorPowershell = "powershell" - ExecutorRuby = "ruby" - ExecutorBinary = "binary" - ExecutorCmd = "cmd.exe" -) - -// StepType denotes the type of a step in a TTP. -type StepType string - -// Constants for defining the types of steps available. -const ( - StepCreateFile = "createFileStep" - StepUnset = "unsetStep" - StepFile = "fileStep" - StepFetchURI = "fetchURIStep" - StepBasic = "basicStep" - StepSubTTP = "subTTPStep" - StepCleanup = "cleanupStep" - StepEdit = "editStep" -) - -// Act represents a single action within a TTP (Tactics, Techniques, -// and Procedures) step. -// -// Condition: The condition that needs to be satisfied for the Act to execute. -// Environment: Environment variables used during the Act's execution. -// Name: The unique name of the Act. -// WorkDir: The working directory of the Act. -// Type: The type of the Act (e.g., Command, File, or Setup). -// success: Indicates whether the execution of the Act was successful. -// stepRef: Reference to other steps in the sequence. -// output: The output of the Act's execution. -type Act struct { - Condition string `yaml:"if,omitempty"` - Environment map[string]string `yaml:"env,omitempty"` - Name string `yaml:"name"` - WorkDir string `yaml:"-"` - Outputs map[string]outputs.Spec `yaml:"outputs,omitempty"` - Type StepType `yaml:"-"` +// CommonStepFields contains the fields +// common to every type of step (such as Name). +// It centralizes validation to simplify the code +type CommonStepFields struct { + Name string `yaml:"name,omitempty"` + Description string `yaml:"description,omitempty"` + + // CleanupSpec is exported so that UnmarshalYAML + // can see it - however, it should be considered + // to be a private detail of this file + // and not referenced elsewhere in the codebase + CleanupSpec yaml.Node `yaml:"cleanup,omitempty"` } -// CleanupAct interface is implemented by anything that requires a cleanup step. -type CleanupAct interface { - Cleanup(execCtx TTPExecutionContext) (*ActResult, error) - StepName() string - SetDir(dir string) - IsNil() bool - Validate(execCtx TTPExecutionContext) error +// Step contains a TTPForge executable action +// and its associated cleanup action (if specified) +type Step struct { + CommonStepFields + + // These are where the actual executable content + // of the step (and its associated cleanup process) + // live - they are not deserialized directly from YAML + // but rather must be decoded by ParseAction + action Action + cleanup Action } -// Step is an interface that represents a TTP step. Types that implement -// this interface must provide methods for setting up the environment and -// output references, setting the working directory, getting the cleanup -// actions, executing the step, checking if the step is empty, explaining -// validation errors, validating the step, fetching arguments, getting output, -// searching output, setting output success status, checking success status, -// returning the step name, and getting the step type. -type Step interface { - SetDir(dir string) - // Need list in case some steps are encapsulating many cleanup steps - GetCleanup() []CleanupAct - // Execute will need to take care of the condition checks/etc... - Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) - IsNil() bool - ExplainInvalid() error - Validate(execCtx TTPExecutionContext) error - StepName() string - GetType() StepType -} +func isDefaultCleanup(cleanupNode *yaml.Node) (bool, error) { + var testStr string + // is it a string? if not, let the subsequent decoding + // in the calling function deal with it + if err := cleanupNode.Decode(&testStr); err != nil { + return false, nil + } -// SetDir sets the working directory for the Act. -// -// **Parameters:** -// -// dir: A string representing the directory path to be set -// as the working directory. -func (a *Act) SetDir(dir string) { - a.WorkDir = dir + // if it is a string, it must be a valid string + if testStr == "default" { + return true, nil + } + return false, fmt.Errorf("invalid cleanup value specified: %v", testStr) } -// IsNil checks whether the Act is nil (i.e., it does not have a name). -// -// **Returns:** -// -// bool: True if the Act has no name, false otherwise. -func (a *Act) IsNil() bool { - switch { - case a.Name == "": +// ShouldCleanupOnFailure specifies that this step should be cleaned +// up even if its Execute(...) failed. +// We usually don't want to do this - for example, +// you shouldn't try to remove_path a create_file that failed) +// However, certain step types (especially SubTTPs) need to run cleanup even if they fail +func (s *Step) ShouldCleanupOnFailure() bool { + switch s.action.(type) { + case *SubTTPStep: return true default: return false } } -// ExplainInvalid returns an error explaining why the Act is invalid. -// -// **Returns:** -// -// error: An error explaining why the Act is invalid, or nil -// if the Act is valid. -func (a *Act) ExplainInvalid() error { - switch { - case a.Name == "": - return errors.New("no name provided for current step") +// ShouldUseImplicitDefaultCleanup is a hack +// to make subTTPs always run their default +// cleanup process even when `cleanup: default` is +// not explicitly specified - this is purely for backward +// compatibility +func ShouldUseImplicitDefaultCleanup(action Action) bool { + switch action.(type) { + case *SubTTPStep: + return true default: - return nil + return false } } -// StepName returns the name of the Act. -// -// **Returns:** -// -// string: The name of the Act. -func (a *Act) StepName() string { - return a.Name -} +// UnmarshalYAML implements custom deserialization +// process to ensure that the step action and its +// cleanup action are decoded to the correct struct type +func (s *Step) UnmarshalYAML(node *yaml.Node) error { + + // Decode all of the shared fields. + // Use of this auxiliary type prevents infinite recursion + var csf CommonStepFields + err := node.Decode(&csf) + if err != nil { + return err + } + s.CommonStepFields = csf -// Validate checks the Act for any validation errors, such as the presence of -// spaces in the name. -// -// **Returns:** -// -// error: An error if any validation errors are found, or nil if -// the Act is valid. -func (a *Act) Validate() error { - // Make sure name is of format we can index - if strings.Contains(a.Name, " ") { - return errors.New("name must not contain spaces") + if s.Name == "" { + return errors.New("no name specified for step") } - return nil -} + // figure out what kind of action is + // associated with executing this step + s.action, err = s.ParseAction(node) + if err != nil { + return err + } -// CheckCondition checks the condition specified for an Act and returns true -// if it matches the current OS, false otherwise. If the condition is "always", -// the function returns true. -// If an error occurs while checking the condition, it is returned. -// -// **Returns:** -// -// bool: true if the condition matches the current OS or the -// condition is "always", false otherwise. -// -// error: An error if an error occurs while checking the condition. -func (a *Act) CheckCondition() (bool, error) { - switch a.Condition { - case "windows": - if runtime.GOOS == "windows" { - return true, nil + // figure out what kind of action is + // associated with cleaning up this step + if csf.CleanupSpec.IsZero() { + // hack for subTTPs - they should always use their default cleanup + if ShouldUseImplicitDefaultCleanup(s.action) { + s.cleanup = s.action.GetDefaultCleanupAction() } - case "darwin": - if runtime.GOOS == "darwin" { - return true, nil + } else { + useDefaultCleanup, err := isDefaultCleanup(&csf.CleanupSpec) + if err != nil { + return err } - case "linux": - if runtime.GOOS == "linux" { - return true, nil + if useDefaultCleanup { + if dca := s.action.GetDefaultCleanupAction(); dca != nil { + s.cleanup = dca + return nil + } + return fmt.Errorf("`cleanup: default` was specified but step %v is not an action type that has a default cleanup action", s.Name) } - // Run even if a previous step has failed. - case "always": - return true, nil - default: - return false, nil + s.cleanup, err = s.ParseAction(&csf.CleanupSpec) + if err != nil { + return err + } } - return false, nil + return nil } -// MakeCleanupStep creates a CleanupAct based on the given yaml.Node. -// If the node is empty or invalid, it returns nil. If the node contains a -// BasicStep or FileStep, the corresponding CleanupAct is created and returned. -// -// **Parameters:** -// -// node: A pointer to a yaml.Node containing the parameters to -// create the CleanupAct. -// -// **Returns:** -// -// CleanupAct: The created CleanupAct, or nil if the node is empty or invalid. -// -// error: An error if the node contains invalid parameters. -func (a *Act) MakeCleanupStep(node *yaml.Node) (CleanupAct, error) { - if node.IsZero() { - return nil, nil - } - - basic, berr := a.tryDecodeBasicStep(node) - if berr == nil && !basic.IsNil() { - logging.L().Debugw("cleanup step found", "basicstep", basic) - return basic, nil - } +// Execute runs the action associated with this step +func (s *Step) Execute(execCtx TTPExecutionContext) (*ActResult, error) { + return s.action.Execute(execCtx) +} - file, ferr := a.tryDecodeFileStep(node) - if ferr == nil && !file.IsNil() { - logging.L().Debugw("cleanup step found", "filestep", file) - return file, nil +// Cleanup runs the cleanup action associated with this step +func (s *Step) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { + if s.cleanup != nil { + return s.cleanup.Execute(execCtx) } - - err := fmt.Errorf("invalid parameters for cleanup steps with basic [%v], file [%v]", berr, ferr) - logging.L().Errorw(err.Error(), zap.Error(err)) - return nil, err + logging.L().Infof("No Cleanup Action Defined for Step %v", s.Name) + return &ActResult{}, nil } -func (a *Act) tryDecodeBasicStep(node *yaml.Node) (*BasicStep, error) { - basic := NewBasicStep() - err := node.Decode(&basic) - if err == nil && basic.Name == "" { - basic.Name = fmt.Sprintf("cleanup-%s", a.Name) - basic.Type = StepCleanup +// Validate checks that both the step action and cleanup +// action are valid +func (s *Step) Validate(execCtx TTPExecutionContext) error { + if err := s.action.Validate(execCtx); err != nil { + return err + } + if s.cleanup != nil { + if err := s.cleanup.Validate(execCtx); err != nil { + return err + } } - return basic, err + return nil } -func (a *Act) tryDecodeFileStep(node *yaml.Node) (*FileStep, error) { - file := NewFileStep() - err := node.Decode(&file) - if err == nil && file.Name == "" { - file.Name = fmt.Sprintf("cleanup-%s", a.Name) - file.Type = StepCleanup +// ParseAction decodes an action (from step or cleanup) in YAML +// format into the appropriate struct +func (s *Step) ParseAction(node *yaml.Node) (Action, error) { + // actionCandidates := []Action{NewBasicStep(), NewFileStep(), NewEditStep(), NewFetchURIStep(), NewCreateFileStep()} + actionCandidates := []Action{NewBasicStep(), NewFileStep(), NewSubTTPStep(), NewEditStep(), NewFetchURIStep(), NewCreateFileStep(), &PrintStrAction{}} + var action Action + for _, actionType := range actionCandidates { + err := node.Decode(actionType) + if err == nil && !actionType.IsNil() { + if action != nil { + // Must catch bad steps with ambiguous types, such as: + // - name: hello + // file: bar + // ttp: foo + // + // we can't use KnownFields to solve this without a massive + // refactor due to https://github.com/go-yaml/yaml/issues/460 + // note: we check for non-empty name earlier so s.Name will be non-empty + return nil, fmt.Errorf("step %v has ambiguous type", s.Name) + } + action = actionType + } + } + if action == nil { + return nil, fmt.Errorf("step %v did not match any valid step type", s.Name) } - return file, err + return action, nil } diff --git a/pkg/blocks/step_test.go b/pkg/blocks/step_test.go old mode 100755 new mode 100644 index 00708847..ead0237d --- a/pkg/blocks/step_test.go +++ b/pkg/blocks/step_test.go @@ -20,76 +20,186 @@ THE SOFTWARE. package blocks_test import ( + "fmt" + "os" "testing" "github.com/facebookincubator/ttpforge/pkg/blocks" + "github.com/spf13/afero" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" ) -func TestMakeCleanupStep(t *testing.T) { - tests := []struct { - name string - yamlData string - expectedType string - expectedError string +func TestStep(t *testing.T) { + testCases := []struct { + name string + content string + wantUnmarshalError bool + wantExecuteError bool + expectedExecuteStdout string + wantCleanupError bool + expectedCleanupStdout string }{ { - name: "BasicStep", - yamlData: ` -name: "cleanup-test" -command: "echo 'cleanup'" -inline: true -`, - expectedType: "BasicStep", + name: "Run inline command (no error)", + content: `name: inline_step +description: runs a valid inline command +inline: echo inline_step_test`, + expectedExecuteStdout: "inline_step_test\n", + }, + { + name: "Run Cleanup (inline - no error)", + content: `name: inline_step +description: runs an invalid inline command +inline: echo executing +cleanup: + inline: echo cleanup`, + expectedExecuteStdout: "executing\n", + expectedCleanupStdout: "cleanup\n", }, { - name: "FileStep", - yamlData: ` -name: "cleanup-test" -src: "source/file" -dest: "destination/file" -filepath: true -`, - expectedType: "FileStep", - expectedError: "empty FilePath provided", + name: "Run inline command (execution error)", + content: `name: inline_step +description: runs an invalid inline command +inline: this will error`, + wantExecuteError: true, }, { - name: "InvalidStep", - yamlData: ` -invalid_key: "invalid_value" -`, - expectedError: "invalid parameters for cleanup steps with basic [(inline) empty], file [empty FilePath provided]", + name: "Step With Empty Name", + content: `inline: echo should_error_before_execution`, + wantUnmarshalError: true, }, } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - var node yaml.Node - require.NoError(t, yaml.Unmarshal([]byte(test.yamlData), &node)) - - act := &blocks.Act{} - cleanupAct, err := act.MakeCleanupStep(&node) - - if test.expectedError != "" { - assert.Error(t, err) - assert.Contains(t, err.Error(), test.expectedError) - assert.Nil(t, cleanupAct) - } else { - assert.NoError(t, err) - - switch test.expectedType { - case "BasicStep": - _, ok := cleanupAct.(*blocks.BasicStep) - assert.True(t, ok, "Expected BasicStep") - case "FileStep": - _, ok := cleanupAct.(*blocks.FileStep) - assert.True(t, ok, "Expected FileStep") - default: - t.Fatalf("Unknown expected type: %s", test.expectedType) - } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var s blocks.Step + var execCtx blocks.TTPExecutionContext + + // parse the step + err := yaml.Unmarshal([]byte(tc.content), &s) + if tc.wantUnmarshalError { + require.Error(t, err) + return + } + require.NoError(t, err) + + // validate the step + err = s.Validate(execCtx) + require.NoError(t, err) + + // execute the step and check output + result, err := s.Execute(execCtx) + if tc.wantExecuteError { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tc.expectedExecuteStdout, result.Stdout) + + // run cleanup and check output + cleanupResult, err := s.Cleanup(execCtx) + if tc.wantCleanupError { + require.Error(t, err) + return } + require.NoError(t, err) + assert.Equal(t, tc.expectedCleanupStdout, cleanupResult.Stdout) + }) + } +} + +func TestCleanupDefault(t *testing.T) { + testCases := []struct { + name string + contentFmtStr string + wantUnmarshalError bool + wantExecuteError bool + expectedFileContents string + wantCleanupError bool + fileShouldExistAfterCleanup bool + }{ + { + name: "create_file Default Cleanup", + contentFmtStr: `name: create_file_step +description: creates a file and then deletes it +create_file: %v +contents: this is a test +cleanup: default`, + expectedFileContents: "this is a test", + }, + { + name: "create_file with invalid cleanup", + contentFmtStr: `name: create_file_step +description: invalid cleanup value +create_file: %v +contents: this is a test +cleanup: invalid`, + wantUnmarshalError: true, + }, + { + name: "create_file with non-default cleanup", + contentFmtStr: `name: create_file_step +description: non-default cleanup value +create_file: %v +contents: testing non default cleanup +cleanup: + inline: echo "will not delete file"`, + expectedFileContents: "testing non default cleanup", + fileShouldExistAfterCleanup: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var s blocks.Step + var execCtx blocks.TTPExecutionContext + + // hack to get a valid temporary path without creating it + tmpFile, err := os.CreateTemp("", "ttpforge-test-cleanup-default") + require.NoError(t, err) + filePath := tmpFile.Name() + err = os.Remove(filePath) + require.NoError(t, err) + + content := fmt.Sprintf(tc.contentFmtStr, filePath) + err = yaml.Unmarshal([]byte(content), &s) + if tc.wantUnmarshalError { + require.Error(t, err) + return + } + require.NoError(t, err) + + // validate the step + err = s.Validate(execCtx) + require.NoError(t, err) + + // execute the step and check file contents + _, err = s.Execute(execCtx) + if tc.wantExecuteError { + require.Error(t, err) + return + } + require.NoError(t, err) + contentBytes, err := os.ReadFile(filePath) + require.NoError(t, err) + assert.Equal(t, tc.expectedFileContents, string(contentBytes)) + + // run cleanup + _, err = s.Cleanup(execCtx) + if tc.wantCleanupError { + require.Error(t, err) + return + } + require.NoError(t, err) + + // verify that file was deleted + fsys := afero.NewOsFs() + exists, err := afero.Exists(fsys, filePath) + require.NoError(t, err) + assert.Equal(t, tc.fileShouldExistAfterCleanup, exists) }) } } diff --git a/pkg/blocks/subttp.go b/pkg/blocks/subttp.go index 9deae136..a05018ff 100644 --- a/pkg/blocks/subttp.go +++ b/pkg/blocks/subttp.go @@ -21,23 +21,19 @@ package blocks import ( "errors" - "fmt" "strings" "github.com/facebookincubator/ttpforge/pkg/logging" - "gopkg.in/yaml.v3" ) // SubTTPStep represents a step within a parent TTP that references a separate TTP file. type SubTTPStep struct { - *Act `yaml:",inline"` TtpFile string `yaml:"ttp"` Args map[string]string `yaml:"args"` - // Omitting because the sub steps will contain the cleanups. - CleanupSteps []CleanupAct `yaml:"-,omitempty"` - ttp *TTP - subExecCtx TTPExecutionContext + ttp *TTP + subExecCtx TTPExecutionContext + firstStepToCleanupIdx int } // NewSubTTPStep creates a new SubTTPStep and returns a pointer to it. @@ -45,9 +41,12 @@ func NewSubTTPStep() *SubTTPStep { return &SubTTPStep{} } -// GetCleanup returns a slice of CleanupAct associated with the SubTTPStep. -func (s *SubTTPStep) GetCleanup() []CleanupAct { - return []CleanupAct{s} +// GetDefaultCleanupAction will instruct the calling code +// to cleanup all successful steps of this subTTP +func (s *SubTTPStep) GetDefaultCleanupAction() Action { + return &subTTPCleanupAction{ + step: s, + } } func aggregateResults(results []*ActResult) *ActResult { @@ -64,41 +63,6 @@ func aggregateResults(results []*ActResult) *ActResult { } } -// Cleanup runs the cleanup actions associated with all successful sub-steps -func (s *SubTTPStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) { - var results []*ActResult - for _, step := range s.CleanupSteps { - result, err := step.Cleanup(s.subExecCtx) - if err != nil { - return nil, err - } - results = append(results, result) - } - return aggregateResults(results), nil -} - -// UnmarshalYAML is a custom unmarshaller for SubTTPStep which decodes -// a YAML node into a SubTTPStep instance. -func (s *SubTTPStep) UnmarshalYAML(node *yaml.Node) error { - type Subtmpl struct { - Act `yaml:",inline"` - TtpFile string `yaml:"ttp"` - Args map[string]string `yaml:"args"` - } - var substep Subtmpl - - if err := node.Decode(&substep); err != nil { - return err - } - logging.L().Debugw("step found", "substep", substep) - - s.Act = &substep.Act - s.TtpFile = substep.TtpFile - s.Args = substep.Args - - return nil -} - func (s *SubTTPStep) processSubTTPArgs(execCtx TTPExecutionContext) ([]string, error) { var argKvStrs []string for k, v := range s.Args { @@ -114,43 +78,26 @@ func (s *SubTTPStep) processSubTTPArgs(execCtx TTPExecutionContext) ([]string, e // Execute runs each step of the TTP file associated with the SubTTPStep // and manages the outputs and cleanup steps. -func (s *SubTTPStep) Execute(execCtx TTPExecutionContext) (*ExecutionResult, error) { - logging.L().Infof("[*] Executing Sub TTP: %s", s.Name) - availableSteps := make(map[string]Step) - - var results []*ActResult - for _, step := range s.ttp.Steps { - stepCopy := step - logging.L().Infof("[+] Running current step: %s", step.StepName()) - - result, err := stepCopy.Execute(s.subExecCtx) - if err != nil { - return nil, err - } - results = append(results, &result.ActResult) - - availableSteps[stepCopy.StepName()] = stepCopy - - stepClean := stepCopy.GetCleanup() - if stepClean != nil { - logging.L().Debugw("adding cleanup step", "cleanup", stepClean) - s.CleanupSteps = append(stepCopy.GetCleanup(), s.CleanupSteps...) - } - - logging.L().Infof("[+] Finished running step: %s", stepCopy.StepName()) +func (s *SubTTPStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) { + logging.L().Infof("[*] Executing Sub TTP: %s", s.TtpFile) + execResults, firstStepToCleanupIdx, runErr := s.ttp.RunSteps(&execCtx) + s.firstStepToCleanupIdx = firstStepToCleanupIdx + if runErr != nil { + return nil, runErr } + logging.L().Info("[*] Completed SubTTP - No Errors :)") - logging.L().Info("Finished execution of sub ttp file") - - return &ExecutionResult{ - ActResult: *aggregateResults(results), - }, nil + // just a little annoying plumbing due to subtle type differences0 + var actResults []*ActResult + for _, execResult := range execResults.ByIndex { + actResults = append(actResults, &execResult.ActResult) + } + return aggregateResults(actResults), nil } // loadSubTTP loads a TTP file into a SubTTPStep instance // and validates the contained steps. func (s *SubTTPStep) loadSubTTP(execCtx TTPExecutionContext) error { - repo := execCtx.Cfg.Repo subTTPAbsPath, err := execCtx.Cfg.Repo.FindTTP(s.TtpFile) if err != nil { @@ -162,49 +109,17 @@ func (s *SubTTPStep) loadSubTTP(execCtx TTPExecutionContext) error { return err } - ttps, err := LoadTTP(subTTPAbsPath, repo.GetFs(), &s.subExecCtx.Cfg, subArgsKv) + ttps, _, err := LoadTTP(subTTPAbsPath, repo.GetFs(), &s.subExecCtx.Cfg, subArgsKv) if err != nil { return err } s.ttp = ttps - - // run validate to flesh out issues - logging.L().Infof("[*] Validating Sub TTP: %s", s.Name) - for _, step := range s.ttp.Steps { - stepCopy := step - if err := stepCopy.Validate(execCtx); err != nil { - return err - } - } - logging.L().Infof("[*] Finished validating Sub TTP") - - return nil -} - -// GetType returns the type of the step (StepSubTTP for SubTTPStep). -func (s *SubTTPStep) GetType() StepType { - return StepSubTTP -} - -// ExplainInvalid checks for invalid data in the SubTTPStep -// and returns an error explaining any issues found. -// Currently, it checks if the TtpFile field is empty. -func (s *SubTTPStep) ExplainInvalid() error { - if s.TtpFile == "" { - err := fmt.Errorf("error: TtpFile is empty") - if s.Name != "" { - return fmt.Errorf("invalid SubTTPStep [%s]: %w", s.Name, err) - } - return err - } return nil } // IsNil checks if the SubTTPStep is empty or uninitialized. func (s *SubTTPStep) IsNil() bool { switch { - case s.Act.IsNil(): - return true case s.TtpFile == "": return true default: @@ -220,10 +135,6 @@ func (s *SubTTPStep) IsNil() bool { // The steps within the TTP file do not contain any nested SubTTPSteps. // If any of these conditions are not met, an error is returned. func (s *SubTTPStep) Validate(execCtx TTPExecutionContext) error { - if err := s.Act.Validate(); err != nil { - return err - } - if s.TtpFile == "" { return errors.New("a TTP file path is required and must not be empty") } @@ -232,14 +143,5 @@ func (s *SubTTPStep) Validate(execCtx TTPExecutionContext) error { return err } - // Check if steps contain any SubTTPSteps. If they do, return an error. - for _, steps := range s.ttp.Steps { - if steps.GetType() == StepSubTTP { - return errors.New( - "nested SubTTPStep detected within a SubTTPStep, " + - "please remove it for successful execution") - } - } - - return nil + return s.ttp.Validate(execCtx) } diff --git a/pkg/blocks/subttp_test.go b/pkg/blocks/subttp_test.go index 53054275..d7c54c48 100755 --- a/pkg/blocks/subttp_test.go +++ b/pkg/blocks/subttp_test.go @@ -28,7 +28,6 @@ import ( "github.com/spf13/afero" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "gopkg.in/yaml.v3" ) @@ -73,13 +72,12 @@ steps: func TestSubTTPExecution(t *testing.T) { tests := []struct { - name string - spec repos.Spec - fsys afero.Fs - stepYAML string - expectError bool - expectedOutput string - expectedCleanupOutput string + name string + spec repos.Spec + fsys afero.Fs + stepYAML string + expectError bool + expectedOutput string }{ { name: "Simple Sub TTP Execution", @@ -115,8 +113,7 @@ args: fsys: makeTestFsForSubTTPs(t), stepYAML: `name: with-cleanup ttp: with/cleanup.yaml`, - expectedOutput: "sub_step_1_output\nsub_step_2_output\n", - expectedCleanupOutput: "cleanup_sub_step_2\ncleanup_sub_step_1\n", + expectedOutput: "sub_step_1_output\nsub_step_2_output\n", }, } @@ -140,14 +137,6 @@ ttp: with/cleanup.yaml`, result, err := step.Execute(execCtx) require.NoError(t, err) assert.Equal(t, tc.expectedOutput, result.Stdout) - - if tc.expectedCleanupOutput != "" { - cleanups := step.GetCleanup() - require.NotNil(t, cleanups) - cleanupResult, err := cleanups[0].Cleanup(execCtx) - require.NoError(t, err) - assert.Equal(t, tc.expectedCleanupOutput, cleanupResult.Stdout) - } }) } } diff --git a/pkg/blocks/subttpcleanup.go b/pkg/blocks/subttpcleanup.go new file mode 100644 index 00000000..30b57bd9 --- /dev/null +++ b/pkg/blocks/subttpcleanup.go @@ -0,0 +1,46 @@ +/* +Copyright © 2023-present, Meta Platforms, Inc. and affiliates +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +package blocks + +// subTTPCleanupAction ensures that individual +// steps of the subTTP are appropriately cleaned up +type subTTPCleanupAction struct { + actionDefaults + step *SubTTPStep +} + +// Execute will cleanup the subTTP starting from the last successful step +func (a *subTTPCleanupAction) Execute(execCtx TTPExecutionContext) (*ActResult, error) { + cleanupResults, err := a.step.ttp.startCleanupAtStepIdx(a.step.firstStepToCleanupIdx, &execCtx) + if err != nil { + return nil, err + } + return aggregateResults(cleanupResults), nil +} + +// IsNil is not needed here, as this is not a user-accessible step type +func (a *subTTPCleanupAction) IsNil() bool { + return false +} + +// Validate is not needed here, as this is not a user-accessible step type +func (a *subTTPCleanupAction) Validate(execCtx TTPExecutionContext) error { + return nil +} diff --git a/pkg/blocks/ttps.go b/pkg/blocks/ttps.go index d884ed19..7203b930 100755 --- a/pkg/blocks/ttps.go +++ b/pkg/blocks/ttps.go @@ -27,7 +27,6 @@ import ( "github.com/facebookincubator/ttpforge/pkg/args" "github.com/facebookincubator/ttpforge/pkg/logging" - "go.uber.org/zap" "gopkg.in/yaml.v3" ) @@ -46,7 +45,7 @@ import ( type TTP struct { Name string `yaml:"name,omitempty"` Description string `yaml:"description"` - MitreAttackMapping MitreAttack `yaml:"mitre,omitempty"` + MitreAttackMapping *MitreAttack `yaml:"mitre,omitempty"` Environment map[string]string `yaml:"env,flow,omitempty"` Steps []Step `yaml:"steps,omitempty,flow"` ArgSpecs []args.Spec `yaml:"args,omitempty,flow"` @@ -119,228 +118,180 @@ func reduceIndentation(b []byte, n int) []byte { return bytes.Join(lines, []byte("\n")) } -// UnmarshalYAML is a custom unmarshalling implementation for the TTP structure. -// It decodes a YAML Node into a TTP object, handling the decoding and -// validation of the individual steps within the TTP. +// Validate ensures that all components of the TTP are valid +// It checks key fields, then iterates through each step +// and validates them in turn // // **Parameters:** // -// node: A pointer to a yaml.Node that represents the TTP structure -// to be unmarshalled. +// execCtx: The TTPExecutionContext for the current TTP. // // **Returns:** // -// error: An error if the decoding process fails or if the TTP structure contains invalid steps. -func (t *TTP) UnmarshalYAML(node *yaml.Node) error { - type TTPTmpl struct { - Name string `yaml:"name,omitempty"` - Description string `yaml:"description"` - Environment map[string]string `yaml:"env,flow,omitempty"` - Steps []yaml.Node `yaml:"steps,omitempty,flow"` - ArgSpecs []args.Spec `yaml:"args,omitempty,flow"` - } - - var tmpl TTPTmpl - if err := node.Decode(&tmpl); err != nil { - return err - } - - t.Name = tmpl.Name - t.Description = tmpl.Description - t.Environment = tmpl.Environment - t.ArgSpecs = tmpl.ArgSpecs +// error: An error if any part of the validation fails, otherwise nil. +func (t *TTP) Validate(execCtx TTPExecutionContext) error { + logging.L().Info("[*] Validating Steps") - // Check for and handle a mitre node - var mitreNode *yaml.Node - for i := 0; i < len(node.Content)-1; i += 2 { - keyNode := node.Content[i] - if keyNode.Value == "mitre" { - mitreNode = node.Content[i+1] - break - } + // validate MITRE mapping + if t.MitreAttackMapping != nil && len(t.MitreAttackMapping.Tactics) == 0 { + return fmt.Errorf("TTP '%s' has a MitreAttackMapping but no Tactic is defined", t.Name) } - if mitreNode != nil { - if err := mitreNode.Decode(&t.MitreAttackMapping); err != nil { + for _, step := range t.Steps { + stepCopy := step + if err := stepCopy.Validate(execCtx); err != nil { return err } - // if we have a MitreAttackMapping, ensure there's a tactic - if len(t.MitreAttackMapping.Tactics) == 0 { - return fmt.Errorf("TTP '%s' has a MitreAttackMapping but no Tactic is defined", t.Name) - } - } - - return t.decodeSteps(tmpl.Steps) -} - -func (t *TTP) decodeSteps(steps []yaml.Node) error { - for stepIdx, stepNode := range steps { - decoded := false - // these candidate steps are pointers, so this line - // MUST be inside the outer step loop or horrible things will happen - // #justpointerthings - stepTypes := []Step{NewBasicStep(), NewFileStep(), NewSubTTPStep(), NewEditStep(), NewFetchURIStep(), NewCreateFileStep()} - for _, stepType := range stepTypes { - err := stepNode.Decode(stepType) - if err == nil && !stepType.IsNil() { - // Must catch bad steps with ambiguous types, such as: - // - name: hello - // file: bar - // ttp: foo - // - // we can't use KnownFields to solve this without a massive - // refactor due to https://github.com/go-yaml/yaml/issues/460 - if decoded { - return fmt.Errorf("step #%v has ambiguous type", stepIdx+1) - } - logging.L().Debugw("decoded step", "step", stepType) - t.Steps = append(t.Steps, stepType) - decoded = true - } - } - - if !decoded { - return fmt.Errorf("Step #%v does not match any supported step type", stepIdx+1) - } } - + logging.L().Info("[+] Finished validating steps") return nil } -func (t *TTP) setWorkingDirectory() error { - if t.WorkDir != "" { - return nil +func (t *TTP) chdir() (func(), error) { + // note: t.WorkDir may not be set in tests but should + // be set when actually using `ttpforge run` + if t.WorkDir == "" { + return func() {}, nil } - - path, err := os.Getwd() + origDir, err := os.Getwd() if err != nil { - return err + return nil, err } - t.WorkDir = path - return nil + if err := os.Chdir(t.WorkDir); err != nil { + return nil, err + } + return func() { + if err := os.Chdir(origDir); err != nil { + logging.L().Errorf("could not restore original directory %v: %v", origDir, err) + } + }, nil } -// ValidateSteps iterates through each step in the TTP and validates it. -// It sets the working directory for each step before calling its Validate -// method. If any step fails validation, the method returns an error. -// If all steps are successfully validated, the method returns nil. +// Execute executes all of the steps in the given TTP, +// then runs cleanup if appropriate // // **Parameters:** // -// execCtx: The TTPExecutionContext for the current TTP. +// execCfg: The TTPExecutionConfig for the current TTP. // // **Returns:** // -// error: An error if any step validation fails, otherwise nil. -func (t *TTP) ValidateSteps(execCtx TTPExecutionContext) error { - logging.L().Info("[*] Validating Steps") - - for _, step := range t.Steps { - stepCopy := step - // pass in the directory - stepCopy.SetDir(t.WorkDir) - if err := stepCopy.Validate(execCtx); err != nil { - logging.L().Errorw("failed to validate %s step: %v", step, zap.Error(err)) - return err - } +// *StepResultsRecord: A StepResultsRecord containing the results of each step. +// error: An error if any of the steps fail to execute. +func (t *TTP) Execute(execCtx *TTPExecutionContext) (*StepResultsRecord, error) { + stepResults, firstStepToCleanupIdx, runErr := t.RunSteps(execCtx) + if runErr != nil { + // we need to run cleanup so we don't return here + logging.L().Errorf("[*] Error executing TTP: %v", runErr) + } else { + logging.L().Info("[*] Completed TTP - No Errors :)") } - logging.L().Info("[+] Finished validating steps") - return nil -} - -func (t *TTP) executeSteps(execCtx TTPExecutionContext) (*StepResultsRecord, []CleanupAct, error) { - logging.L().Infof("[+] Running current TTP: %s", t.Name) - stepResults := NewStepResultsRecord() - execCtx.StepResults = stepResults - var cleanup []CleanupAct - - for _, step := range t.Steps { - stepCopy := step - logging.L().Infof("[+] Running current step: %s", step.StepName()) - - execResult, err := stepCopy.Execute(execCtx) + if !execCtx.Cfg.NoCleanup { + if execCtx.Cfg.CleanupDelaySeconds > 0 { + logging.L().Infof("[*] Sleeping for Requested Cleanup Delay of %v Seconds", execCtx.Cfg.CleanupDelaySeconds) + time.Sleep(time.Duration(execCtx.Cfg.CleanupDelaySeconds) * time.Second) + } + cleanupResults, err := t.startCleanupAtStepIdx(firstStepToCleanupIdx, execCtx) if err != nil { - return stepResults, cleanup, err + return nil, err + } + // since ByIndex and ByName both contain pointers to + // the same underlying struct, this will update both + for cleanupIdx, cleanupResult := range cleanupResults { + execCtx.StepResults.ByIndex[cleanupIdx].Cleanup = cleanupResult } - stepResults.ByName[step.StepName()] = execResult - stepResults.ByIndex = append(stepResults.ByIndex, execResult) - - // Enters in reverse order - cleanup = append(stepCopy.GetCleanup(), cleanup...) - logging.L().Infof("[+] Finished running step: %s", step.StepName()) } - return stepResults, cleanup, nil + // still pass up the run error after our cleanup + return stepResults, runErr } // RunSteps executes all of the steps in the given TTP. // // **Parameters:** // -// execCfg: The TTPExecutionConfig for the current TTP. +// execCtx: The current TTPExecutionContext // // **Returns:** // // *StepResultsRecord: A StepResultsRecord containing the results of each step. +// int: the index of the step where cleanup should start (usually the last successful step) // error: An error if any of the steps fail to execute. -func (t *TTP) RunSteps(execCfg TTPExecutionConfig) (*StepResultsRecord, error) { - if err := t.setWorkingDirectory(); err != nil { - return nil, err - } - - execCtx := TTPExecutionContext{ - Cfg: execCfg, - } - - if err := t.ValidateSteps(execCtx); err != nil { - return nil, err +func (t *TTP) RunSteps(execCtx *TTPExecutionContext) (*StepResultsRecord, int, error) { + // go to the configuration directory for this TTP + changeBack, err := t.chdir() + if err != nil { + return nil, -1, err } + defer changeBack() - // stop after validation for dry run if execCtx.Cfg.DryRun { logging.L().Info("[*] Dry-Run Requested - Returning Early") - return nil, nil + return nil, -1, nil } - stepResults, cleanup, err := t.executeSteps(execCtx) - if err != nil { - // we need to run cleanup so we don't return here - logging.L().Errorf("[*] Error executing TTP: %v", err) - } + // actually run all the steps + logging.L().Infof("[+] Running current TTP: %s", t.Name) + stepResults := NewStepResultsRecord() + execCtx.StepResults = stepResults + firstStepToCleanupIdx := -1 + for _, step := range t.Steps { + stepCopy := step + logging.L().Infof("[+] Running current step: %s", step.Name) - logging.L().Info("[*] Completed TTP") + // core execution - run the step action + stepResult, err := stepCopy.Execute(*execCtx) - if !execCtx.Cfg.NoCleanup { - if execCtx.Cfg.CleanupDelaySeconds > 0 { - logging.L().Infof("[*] Sleeping for Requested Cleanup Delay of %v Seconds", execCtx.Cfg.CleanupDelaySeconds) - time.Sleep(time.Duration(execCtx.Cfg.CleanupDelaySeconds) * time.Second) - } - if len(cleanup) > 0 { - logging.L().Info("[*] Beginning Cleanup") - if err := t.executeCleanupSteps(execCtx, cleanup, *stepResults); err != nil { - logging.L().Errorw("error encountered in cleanup step: %v", err) - return nil, err + // this part is tricky - SubTTP steps + // must be cleaned up even on failure + // (because substeps may have succeeded) + // so in those cases, we need to save the result + // even if nil + if err != nil { + if step.ShouldCleanupOnFailure() { + logging.L().Infof("[+] Cleaning up failed step %s", step.Name) + logging.L().Infof("[+] Full Cleanup will Run Afterward") + _, cleanupErr := step.Cleanup(*execCtx) + if cleanupErr != nil { + logging.L().Errorf("error cleaning up failed step %v: %v", step.Name, err) + } } - logging.L().Info("[*] Finished Cleanup") - } else { - logging.L().Info("[*] No Cleanup Steps Found") + return nil, firstStepToCleanupIdx, err } - } + firstStepToCleanupIdx++ - return stepResults, err -} + execResult := &ExecutionResult{ + ActResult: *stepResult, + } + stepResults.ByName[step.Name] = execResult + stepResults.ByIndex = append(stepResults.ByIndex, execResult) -func (t *TTP) executeCleanupSteps(execCtx TTPExecutionContext, cleanupSteps []CleanupAct, stepResults StepResultsRecord) error { - for cleanupIdx, step := range cleanupSteps { - stepCopy := step + // Enters in reverse order + logging.L().Infof("[+] Finished running step: %s", step.Name) + } + return stepResults, firstStepToCleanupIdx, nil +} - cleanupResult, err := stepCopy.Cleanup(execCtx) +func (t *TTP) startCleanupAtStepIdx(firstStepToCleanupIdx int, execCtx *TTPExecutionContext) ([]*ActResult, error) { + // go to the configuration directory for this TTP + changeBack, err := t.chdir() + if err != nil { + return nil, err + } + defer changeBack() + + logging.L().Info("[*] Beginning Cleanup") + var cleanupResults []*ActResult + for cleanupIdx := firstStepToCleanupIdx; cleanupIdx >= 0; cleanupIdx-- { + stepToCleanup := t.Steps[cleanupIdx] + cleanupResult, err := stepToCleanup.Cleanup(*execCtx) + // must be careful to put these in step order, not in execution (reverse) order + cleanupResults = append([]*ActResult{cleanupResult}, cleanupResults...) if err != nil { - logging.L().Errorw("error encountered in stepCopy cleanup: %v", err) - return err + logging.L().Errorf("error cleaning up step: %v", err) + logging.L().Errorf("will continue to try to cleanup other steps") + continue } - // since ByIndex and ByName both contain pointers to - // the same underlying struct, this will update both - stepResults.ByIndex[len(cleanupSteps)-cleanupIdx-1].Cleanup = cleanupResult } - return nil + logging.L().Info("[*] Finished Cleanup") + return cleanupResults, nil } diff --git a/pkg/blocks/ttps_test.go b/pkg/blocks/ttps_test.go index 132b7a15..7d7be743 100755 --- a/pkg/blocks/ttps_test.go +++ b/pkg/blocks/ttps_test.go @@ -132,7 +132,7 @@ steps: } } -func TestTTP_ValidateSteps(t *testing.T) { +func TestTTP_Validate(t *testing.T) { testCases := []struct { name string content string @@ -165,7 +165,7 @@ steps: assert.NoError(t, err) } - err = ttp.ValidateSteps(blocks.TTPExecutionContext{}) + err = ttp.Validate(blocks.TTPExecutionContext{}) if tc.wantError { assert.Error(t, err) } else { @@ -317,15 +317,19 @@ steps: return } - stepResults, err := ttp.RunSteps(tc.execConfig) - if tc.wantError && err == nil { - t.Error("expected an error from step execution but got none") - return - } - if !tc.wantError && err != nil { - t.Errorf("didn't expect an error from step execution but got: %s", err) + // validate the TTP + err = ttp.Validate(blocks.TTPExecutionContext{}) + require.NoError(t, err) + + // run it + stepResults, err := ttp.Execute(&blocks.TTPExecutionContext{ + Cfg: tc.execConfig, + }) + if tc.wantError { + require.Error(t, err) return } + require.NoError(t, err) for index, output := range tc.expectedByIndexOut { require.Equal(t, output, stepResults.ByIndex[index].Stdout) @@ -376,6 +380,8 @@ mitre: t.Run(tc.name, func(t *testing.T) { var ttp blocks.TTP err := yaml.Unmarshal([]byte(tc.content), &ttp) + require.NoError(t, err) + err = ttp.Validate(blocks.TTPExecutionContext{}) if tc.wantError { assert.Error(t, err) } else {