Skip to content

Commit c789c78

Browse files
inesusvetfacebook-github-bot
authored andcommitted
Drop non-unique Cleanup logic from expect step (#516)
Summary: Pull Request resolved: #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 #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 #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
1 parent 50f6817 commit c789c78

3 files changed

Lines changed: 2 additions & 182 deletions

File tree

example-ttps/actions/expect/expect.yaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,5 @@ steps:
2626
response: "John"
2727
- prompt: "Enter your age:"
2828
response: "30"
29+
cleanup:
30+
inline: echo "Wipe it"

pkg/blocks/expectstep.go

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ import (
3030
"github.com/Netflix/go-expect"
3131
"github.com/facebookincubator/ttpforge/pkg/logging"
3232
"github.com/facebookincubator/ttpforge/pkg/outputs"
33-
"go.uber.org/zap"
3433
)
3534

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

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

266-
// Cleanup runs the cleanup command if specified.
267-
//
268-
// **Parameters:**
269-
//
270-
// execCtx: Execution context containing environment variables and working
271-
// directory.
272-
//
273-
// **Returns:**
274-
//
275-
// *ActResult: A pointer to the action result.
276-
// error: An error if cleanup fails.
277-
func (s *ExpectStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) {
278-
if s.CleanupStep == "" {
279-
return &ActResult{}, nil
280-
}
281-
282-
logging.L().Info("Running cleanup step")
283-
284-
envAsList := os.Environ()
285-
cmd := s.prepareCommand(context.Background(), execCtx, envAsList, s.CleanupStep)
286-
287-
cmd.Stdout = os.Stdout
288-
cmd.Stdin = os.Stdin
289-
cmd.Stderr = os.Stderr
290-
291-
if err := cmd.Run(); err != nil {
292-
logging.L().Error("Failed to run cleanup command", zap.Error(err))
293-
return nil, fmt.Errorf("failed to run cleanup command: %w", err)
294-
}
295-
296-
logging.L().Info("Cleanup step completed successfully")
297-
return &ActResult{}, nil
298-
}
299-
300264
// CanBeUsedInCompositeAction enables this action to be used in a composite
301265
// action.
302266
//

pkg/blocks/expectstep_test.go

Lines changed: 0 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -370,35 +370,6 @@ steps:
370370
responses:
371371
- prompt: "Enter a number:"
372372
response: "30"
373-
`,
374-
},
375-
{
376-
name: "Test ExpectStep with CleanupStep",
377-
script: `
378-
print("Enter your name:")
379-
name = input()
380-
print(f"Hello, {name}!")
381-
print("Enter your age:")
382-
age = input()
383-
print(f"You are {age} years old.")
384-
`,
385-
content: `
386-
steps:
387-
- name: run_expect_script
388-
description: "Run an expect script to interact with the command."
389-
expect:
390-
inline: |
391-
python3 interactive.py
392-
chdir: "/tmp"
393-
responses:
394-
- prompt: "Enter your name:"
395-
response: "John"
396-
- prompt: "Enter your age:"
397-
response: "30"
398-
cleanup: |
399-
pwd
400-
cat interactive.py
401-
rm interactive.py
402373
`,
403374
},
404375
}
@@ -529,58 +500,6 @@ steps:
529500
assert.Contains(t, normalizedOutput, actualDir)
530501
assert.Contains(t, normalizedOutput, expectedSubstring2)
531502
}
532-
533-
if tc.name == "Test ExpectStep with CleanupStep" {
534-
// Mock the command execution
535-
execCtx := NewTestTTPExecutionContext(tempDir)
536-
console, err := expect.NewConsole(expectNoError(t), sendNoError(t), expect.WithStdout(os.Stdout), expect.WithStdin(os.Stdin))
537-
require.NoError(t, err)
538-
defer console.Close()
539-
540-
cmd := exec.Command("sh", "-c", "python3 "+scriptPath)
541-
cmd.Stdin = console.Tty()
542-
cmd.Stdout = console.Tty()
543-
cmd.Stderr = console.Tty()
544-
545-
if expectStep.Chdir != "" {
546-
cmd.Dir = expectStep.Chdir
547-
}
548-
549-
err = cmd.Start()
550-
require.NoError(t, err)
551-
552-
done := make(chan struct{})
553-
554-
// simulate console input
555-
go func() {
556-
defer close(done)
557-
time.Sleep(1 * time.Second)
558-
console.SendLine("John")
559-
time.Sleep(1 * time.Second)
560-
console.SendLine("30")
561-
time.Sleep(1 * time.Second)
562-
console.Tty().Close() // Close the Tty to signal EOF
563-
}()
564-
565-
_, err = expectStep.Execute(execCtx)
566-
require.NoError(t, err)
567-
<-done
568-
569-
output, err := console.ExpectEOF()
570-
require.NoError(t, err)
571-
572-
// Check the output of the command execution
573-
normalizedOutput := strings.ReplaceAll(output, "\r\n", "\n")
574-
expectedSubstring1 := "Hello, John!\n"
575-
expectedSubstring2 := "You are 30 years old.\n"
576-
assert.Contains(t, normalizedOutput, expectedSubstring1)
577-
assert.Contains(t, normalizedOutput, expectedSubstring2)
578-
579-
// Execute cleanup step
580-
result, err := expectStep.Cleanup(execCtx)
581-
require.NoError(t, err)
582-
assert.NotNil(t, result)
583-
}
584503
})
585504
}
586505
}
@@ -698,71 +617,6 @@ func (m MockCommandExecutor) Run() error {
698617
return m.runFunc()
699618
}
700619

701-
func TestCleanup(t *testing.T) {
702-
testCases := []struct {
703-
name string
704-
cleanupStep string
705-
mockRunOutput string
706-
mockRunError string
707-
wantErr bool
708-
}{
709-
{
710-
name: "No Cleanup Step",
711-
cleanupStep: "",
712-
mockRunOutput: "",
713-
mockRunError: "",
714-
wantErr: false,
715-
},
716-
{
717-
name: "Successful Cleanup",
718-
cleanupStep: "echo 'cleaning up'",
719-
mockRunOutput: "cleaning up",
720-
mockRunError: "",
721-
wantErr: false,
722-
},
723-
{
724-
name: "Failed Cleanup",
725-
cleanupStep: "exit 1",
726-
mockRunOutput: "",
727-
mockRunError: "exit status 1",
728-
wantErr: true,
729-
},
730-
}
731-
732-
for _, tc := range testCases {
733-
t.Run(tc.name, func(t *testing.T) {
734-
s := &ExpectStep{
735-
CleanupStep: tc.cleanupStep,
736-
Executor: "sh",
737-
}
738-
739-
// Mock execCommand for testing
740-
execCommand = func(ctx context.Context, name string, args ...string) *exec.Cmd {
741-
cs := []string{"-test.run=TestHelperProcess", "--", name}
742-
cs = append(cs, args...)
743-
cmd := exec.CommandContext(ctx, os.Args[0], cs...)
744-
cmd.Env = append(os.Environ(),
745-
"GO_WANT_HELPER_PROCESS=1",
746-
fmt.Sprintf("MOCK_RUN_OUTPUT=%s", tc.mockRunOutput),
747-
fmt.Sprintf("MOCK_RUN_ERROR=%s", tc.mockRunError),
748-
)
749-
return cmd
750-
}
751-
752-
execCtx := TTPExecutionContext{
753-
Vars: &TTPExecutionVars{
754-
WorkDir: ".",
755-
},
756-
}
757-
758-
_, err := s.Cleanup(execCtx)
759-
if (err != nil) != tc.wantErr {
760-
t.Errorf("Cleanup() error = %v, wantErr %v", err, tc.wantErr)
761-
}
762-
})
763-
}
764-
}
765-
766620
func TestHelperProcess(*testing.T) {
767621
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
768622
return

0 commit comments

Comments
 (0)