Skip to content

Commit

Permalink
Drop non-unique Cleanup logic from expect step (facebookincubator#516)
Browse files Browse the repository at this point in the history 
Summary:
Pull Request resolved: facebookincubator#516

To enable actual clean ups and prevent errors like

```
ERROR   failed to run command:
        could not load TTP at /home/nesusvet/security-ttpcode/ttps/infra/tupperware/ssh-to-container-as-root.yaml:
        could not parse action for step "start-tupperware-container": action fields did not match any valid action type
```

Changes pretty radical but I don't see why we need a custom implementation of cleanup logic for expect step

Consider small change to the [expect.yaml](https://www.internalfb.com/code/fbsource/fbcode/security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml) example TTP
```
hg d
 diff --git a/fbcode/security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml b/fbcode/security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml
 --- a/fbcode/security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml
+++ b/fbcode/security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml
@@ -26,3 +26,4 @@
           response: "John"
         - prompt: "Enter your age:"
           response: "30"
+      cleanup: echo "Done"
```

Try to run it on master and see NO CLEANUP instructions executed:
```
 buck run security/redteam/purple_team/ttpforge:ttpforge -- run security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml
File changed: fbcode//security/redteam/purple_team/ttpforge/example-ttps/actions/expect/expect.yaml
Buck UI: https://www.internalfb.com/buck2/a008be5d-79d8-4591-960a-64fafbc10a44
Network: Up: 0B  Down: 0B
Jobs completed: 4. Time elapsed: 0.0s.
BUILD SUCCEEDED
INFO    RUNNING TTP: Complex Expect Step with Python Script
INFO    ----------------------------------------
INFO    Executing Step #1: "create_python_script"
INFO    ----------------------------------------
INFO    Executing Step facebookincubator#2: "run_expect_script"
Enter your name:
John
Enter your age:
30
Hello John, you are 30 years old!
INFO    ----------------------------------------
INFO    All TTP steps completed successfully!
INFO    ========================================
INFO    CLEANING UP 2 steps of TTP: "Complex Expect Step with Python Script"
INFO    ----------------------------------------
INFO    Cleaning Up Step facebookincubator#2: "run_expect_script"
INFO    No Cleanup Action Defined for Step run_expect_script
INFO    ----------------------------------------
INFO    Cleaning Up Step #1: "create_python_script"
INFO    No Cleanup Action Defined for Step create_python_script
INFO    ----------------------------------------
INFO    Finished Cleanup Successfully
```

Differential Revision: D64108097

fbshipit-source-id: 2ab395b393b88ebfd0f5dd647cd75f34fce52b97
  • Loading branch information
@inesusvet @facebook-github-bot
inesusvet authored and facebook-github-bot committed Oct 10, 2024
1 parent 50f6817 commit c789c78
Show file tree
Hide file tree
Showing 3 changed files with 2 additions and 182 deletions.
2 changes: 2 additions & 0 deletions example-ttps/actions/expect/expect.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,3 +26,5 @@ steps:
response: "John"
- prompt: "Enter your age:"
response: "30"
cleanup:
inline: echo "Wipe it"
36 changes: 0 additions & 36 deletions pkg/blocks/expectstep.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ import (
"github.com/Netflix/go-expect"
"github.com/facebookincubator/ttpforge/pkg/logging"
"github.com/facebookincubator/ttpforge/pkg/outputs"
"go.uber.org/zap"
)

// ExpectStep represents an expect command.
Expand All @@ -52,7 +51,6 @@ type ExpectStep struct {
Executor string `yaml:"executor,omitempty"`
Expect *ExpectSpec `yaml:"expect,omitempty"`
Environment map[string]string `yaml:"env,omitempty"`
CleanupStep string `yaml:"cleanup,omitempty"`
Outputs map[string]outputs.Spec `yaml:"outputs,omitempty"`
}

Expand Down Expand Up @@ -263,40 +261,6 @@ func (s *ExpectStep) prepareCommand(ctx context.Context, execCtx TTPExecutionCon
return cmd
}

// Cleanup runs the cleanup command if specified.
//
// **Parameters:**
//
// execCtx: Execution context containing environment variables and working
// directory.
//
// **Returns:**
//
// *ActResult: A pointer to the action result.
// error: An error if cleanup fails.
func (s *ExpectStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) {
if s.CleanupStep == "" {
return &ActResult{}, nil
}

logging.L().Info("Running cleanup step")

envAsList := os.Environ()
cmd := s.prepareCommand(context.Background(), execCtx, envAsList, s.CleanupStep)

cmd.Stdout = os.Stdout
cmd.Stdin = os.Stdin
cmd.Stderr = os.Stderr

if err := cmd.Run(); err != nil {
logging.L().Error("Failed to run cleanup command", zap.Error(err))
return nil, fmt.Errorf("failed to run cleanup command: %w", err)
}

logging.L().Info("Cleanup step completed successfully")
return &ActResult{}, nil
}

// CanBeUsedInCompositeAction enables this action to be used in a composite
// action.
//
Expand Down
146 changes: 0 additions & 146 deletions pkg/blocks/expectstep_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -370,35 +370,6 @@ steps:
responses:
- prompt: "Enter a number:"
response: "30"
`,
},
{
name: "Test ExpectStep with CleanupStep",
script: `
print("Enter your name:")
name = input()
print(f"Hello, {name}!")
print("Enter your age:")
age = input()
print(f"You are {age} years old.")
`,
content: `
steps:
- name: run_expect_script
description: "Run an expect script to interact with the command."
expect:
inline: |
python3 interactive.py
chdir: "/tmp"
responses:
- prompt: "Enter your name:"
response: "John"
- prompt: "Enter your age:"
response: "30"
cleanup: |
pwd
cat interactive.py
rm interactive.py
`,
},
}
Expand Down Expand Up @@ -529,58 +500,6 @@ steps:
assert.Contains(t, normalizedOutput, actualDir)
assert.Contains(t, normalizedOutput, expectedSubstring2)
}

if tc.name == "Test ExpectStep with CleanupStep" {
// Mock the command execution
execCtx := NewTestTTPExecutionContext(tempDir)
console, err := expect.NewConsole(expectNoError(t), sendNoError(t), expect.WithStdout(os.Stdout), expect.WithStdin(os.Stdin))
require.NoError(t, err)
defer console.Close()

cmd := exec.Command("sh", "-c", "python3 "+scriptPath)
cmd.Stdin = console.Tty()
cmd.Stdout = console.Tty()
cmd.Stderr = console.Tty()

if expectStep.Chdir != "" {
cmd.Dir = expectStep.Chdir
}

err = cmd.Start()
require.NoError(t, err)

done := make(chan struct{})

// simulate console input
go func() {
defer close(done)
time.Sleep(1 * time.Second)
console.SendLine("John")
time.Sleep(1 * time.Second)
console.SendLine("30")
time.Sleep(1 * time.Second)
console.Tty().Close() // Close the Tty to signal EOF
}()

_, err = expectStep.Execute(execCtx)
require.NoError(t, err)
<-done

output, err := console.ExpectEOF()
require.NoError(t, err)

// Check the output of the command execution
normalizedOutput := strings.ReplaceAll(output, "\r\n", "\n")
expectedSubstring1 := "Hello, John!\n"
expectedSubstring2 := "You are 30 years old.\n"
assert.Contains(t, normalizedOutput, expectedSubstring1)
assert.Contains(t, normalizedOutput, expectedSubstring2)

// Execute cleanup step
result, err := expectStep.Cleanup(execCtx)
require.NoError(t, err)
assert.NotNil(t, result)
}
})
}
}
Expand Down Expand Up @@ -698,71 +617,6 @@ func (m MockCommandExecutor) Run() error {
return m.runFunc()
}

func TestCleanup(t *testing.T) {
testCases := []struct {
name string
cleanupStep string
mockRunOutput string
mockRunError string
wantErr bool
}{
{
name: "No Cleanup Step",
cleanupStep: "",
mockRunOutput: "",
mockRunError: "",
wantErr: false,
},
{
name: "Successful Cleanup",
cleanupStep: "echo 'cleaning up'",
mockRunOutput: "cleaning up",
mockRunError: "",
wantErr: false,
},
{
name: "Failed Cleanup",
cleanupStep: "exit 1",
mockRunOutput: "",
mockRunError: "exit status 1",
wantErr: true,
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
s := &ExpectStep{
CleanupStep: tc.cleanupStep,
Executor: "sh",
}

// Mock execCommand for testing
execCommand = func(ctx context.Context, name string, args ...string) *exec.Cmd {
cs := []string{"-test.run=TestHelperProcess", "--", name}
cs = append(cs, args...)
cmd := exec.CommandContext(ctx, os.Args[0], cs...)
cmd.Env = append(os.Environ(),
"GO_WANT_HELPER_PROCESS=1",
fmt.Sprintf("MOCK_RUN_OUTPUT=%s", tc.mockRunOutput),
fmt.Sprintf("MOCK_RUN_ERROR=%s", tc.mockRunError),
)
return cmd
}

execCtx := TTPExecutionContext{
Vars: &TTPExecutionVars{
WorkDir: ".",
},
}

_, err := s.Cleanup(execCtx)
if (err != nil) != tc.wantErr {
t.Errorf("Cleanup() error = %v, wantErr %v", err, tc.wantErr)
}
})
}
}

func TestHelperProcess(*testing.T) {
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
return
Expand Down

0 comments on commit c789c78

Please sign in to comment.