Refactor expect step, increase test coverage, update pre-commit config (

#503)

Summary:
Pull Request resolved: #503

**Added:**

- `ssh-pw-expect.yaml`: Demonstrate expect step with SSH password.
- `expect.yaml`: Added YAML front matter for better structuring.
- Introduced `ExpectSpec` struct to encapsulate expect script and responses.
- Detailed comments for `ExpectSpec` attributes.
- Enhanced validation for `ExpectStep`, ensuring `ExpectSpec` is provided.
- Dynamic timeout handling in expect responses.
- `.semgrepignore` to avoid creating false positives with tests and examples

**Changed:**

- Refactored `expectstep.go`:
  - Introduced `ExpectSpec` for inline script and responses.
  - Improved validation and execution logic for expect steps.
  - Refactored `Execute` method to accommodate `ExpectSpec`.
  - Updated `prepareCommand` method to use `ExpectSpec` inline script.
  - Enhanced logging and error messages in `Execute` and `Validate` methods.
- Updated `expectstep_test.go`:
  - Mocked SSH interaction within expect steps.
  - Refactored test cases for better coverage.
  - Improved logging and error handling across `expectstep.go`.
- Updated tests to reflect changes in `ExpectStep` structure and logic.

**Removed:**

- `.pre-commit-config.yaml`: Removed `go-vet` hook for streamlined configuration.
- Removed unnecessary imports and streamlined existing ones.

Reviewed By: w51d

Differential Revision: D60238880

fbshipit-source-id: 8930ad1d6d59900a3be49a65682b3199c1312f5a

.pre-commit-config.yaml

@@ -49,16 +49,6 @@ repos:
         language: script
         entry: .hooks/go-licenses.sh check_forbidden
 
-      - id: go-vet
-        name: Run go vet
-        language: script
-        entry: .hooks/go-vet.sh
-        files: '\.go$'
-        always_run: true
-        pass_filenames: true
-        require_serial: true
-        log_file: /tmp/go-vet.log
-
       - id: go-copyright
         name: Ensure all go files have the copyright header
         language: script

.semgrepignore

@@ -0,0 +1,32 @@
+# Items added to this file will be ignored by Semgrep.
+#
+# This file uses .gitignore syntax:
+#
+# To ignore a file anywhere it occurs in your project, enter a
+# glob pattern here. E.g. "*.min.js".
+#
+# To ignore a directory anywhere it occurs in your project, add
+# a trailing slash to the file name. E.g. "dist/".
+#
+# To ignore a file or directory only relative to the project root,
+# include a slash anywhere except the last character. E.g.
+# "/dist/", or "src/generated".
+#
+# Some parts of .gitignore syntax are not supported, and patterns
+# using this syntax will be dropped from the ignore list:
+# - Explicit "include syntax", e.g. "!kept/".
+# - Multi-character expansion syntax, e.g. "*.py[cod]"
+# To include ignore patterns from another file, start a line
+# with ':include', followed by the path of the file. E.g.
+# ":include path/to/other/ignore/file".
+# UPDATE: this will not be be needed in osemgrep which supports
+# all of the .gitignore syntax (!kept/, *.py[cod])
+#
+# To ignore a file with a literal ':' character, escape it with
+# a backslash, e.g. "\:foo".
+
+# Tests
+*_test*
+
+# Examples
+example-ttps/*

example-ttps/actions/expect/expect.yaml

@@ -1,3 +1,4 @@
+---
 api_version: 2.0
 uuid: 7be9d2be-49be-4114-8618-306a62eedec2
 name: Complex Expect Step with Python Script
@@ -17,10 +18,11 @@ steps:
       echo 'age = input()' >> /tmp/interactive.py
       echo 'print(f"Hello {name}, you are {age} years old!")' >> /tmp/interactive.py
   - name: run_expect_script
-    inline: |
-      python3 /tmp/interactive.py
-    responses:
-      - prompt: "Enter your name:"
-        response: "John"
-      - prompt: "Enter your age:"
-        response: "30"
+    expect:
+      inline: |
+        python3 /tmp/interactive.py
+      responses:
+        - prompt: "Enter your name:"
+          response: "John"
+        - prompt: "Enter your age:"
+          response: "30"

example-ttps/actions/expect/ssh-pw-expect.yaml

@@ -0,0 +1,37 @@
+---
+api_version: 2.0
+uuid: 891a38dc-4b2f-4614-9960-a66a7e9499ab
+name: Expect step with SSH password
+description: |
+  This TTP demonstrates the usage of an expect step to automate interaction
+  with an SSH server using a password.
+args:
+  - name: ssh_host
+    description: The hostname or IP address of the SSH server
+    default: target-system
+  - name: ssh_user
+    description: The username to use for the SSH connection
+    default: bobbo
+  - name: ssh_password
+    description: The password to use for the SSH connection
+    default: "Password123!"
+requirements:
+  platforms:
+    - os: darwin
+    - os: windows
+    - os: linux
+steps:
+  - name: run_expect_script
+    expect:
+      inline: |
+        if command -v sshpass >/dev/null 2>&1; then
+          sshpass -p "{{ .Args.ssh_password }}" ssh {{ .Args.ssh_user }}@{{ .Args.ssh_host }}
+        else
+          echo "Error: sshpass is not installed. Please install it before running this script."
+          exit 1
+        fi
+      responses:
+        - prompt: "Welcome to Ubuntu"
+          response: "whoami"
+        - prompt: "{{ .Args.ssh_user }}"
+          response: "exit"

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/gjson v1.17.1
 	go.uber.org/zap v1.27.0
+	golang.org/x/crypto v0.24.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -32,10 +33,10 @@ require (
 	github.com/shopspring/decimal v1.4.0 // indirect
 	github.com/spf13/cast v1.6.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/crypto v0.24.0 // indirect
 	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.21.0 // indirect
 	golang.org/x/text v0.16.0 // indirect

go.sum

@@ -61,6 +61,8 @@ github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3k
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
@@ -105,6 +107,8 @@ golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.2.0/go.mod h1:TVmDHMZPmdnySmBfhjOoOdhjzdE1h4u1VwSiw2l1Nuc=
+golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA=
+golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=

pkg/blocks/expectstep.go

@@ -21,15 +21,13 @@ package blocks
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"os"
 	"os/exec"
 	"regexp"
-	"strings"
 	"time"
 
-	expect "github.com/Netflix/go-expect"
+	"github.com/Netflix/go-expect"
 	"github.com/facebookincubator/ttpforge/pkg/logging"
 	"github.com/facebookincubator/ttpforge/pkg/outputs"
 	"go.uber.org/zap"
@@ -50,15 +48,25 @@ import (
 type ExpectStep struct {
 	actionDefaults `yaml:",inline"`
 	Chdir          string                  `yaml:"chdir,omitempty"`
-	Responses      []Response              `yaml:"responses,omitempty"`
 	Timeout        int                     `yaml:"timeout,omitempty"`
 	Executor       string                  `yaml:"executor,omitempty"`
+	Expect         *ExpectSpec             `yaml:"expect,omitempty"`
 	Environment    map[string]string       `yaml:"env,omitempty"`
-	Inline         string                  `yaml:"inline"`
 	CleanupStep    string                  `yaml:"cleanup,omitempty"`
 	Outputs        map[string]outputs.Spec `yaml:"outputs,omitempty"`
 }
 
+// ExpectSpec represents the expect block in the expect step.
+//
+// **Attributes:**
+//
+// Inline: Inline script to execute.
+// Responses: List of expected prompts and responses.
+type ExpectSpec struct {
+	Inline    string     `yaml:"inline"`
+	Responses []Response `yaml:"responses"`
+}
+
 // Response represents a prompt-response pair.
 //
 // **Attributes:**
@@ -85,7 +93,7 @@ func NewExpectStep() *ExpectStep {
 //
 // bool: True if the step is nil or empty, false otherwise.
 func (s *ExpectStep) IsNil() bool {
-	return s.Inline == "" && len(s.Responses) == 0
+	return s.Expect == nil || (s.Expect.Inline == "" && len(s.Expect.Responses) == 0)
 }
 
 // Validate validates the step, checking for the necessary attributes and
@@ -100,31 +108,24 @@ func (s *ExpectStep) IsNil() bool {
 //
 // error: An error if validation fails.
 func (s *ExpectStep) Validate(_ TTPExecutionContext) error {
-	if len(s.Responses) == 0 {
-		err := errors.New("responses must be provided")
-		logging.L().Error(zap.Error(err))
-		return err
+	if s.Expect == nil {
+		return fmt.Errorf("expectStep is nil")
 	}
 
-	if s.Inline == "" {
-		err := errors.New("inline must be provided")
-		logging.L().Error(zap.Error(err))
-		return err
-	} else if s.Executor == "" {
-		logging.L().Debug("defaulting to bash since executor was not provided")
-		s.Executor = ExecutorBash
+	if len(s.Expect.Responses) == 0 {
+		return fmt.Errorf("responses must be provided")
 	}
 
-	if s.Executor == ExecutorBinary {
-		return nil
+	if s.Expect.Inline == "" {
+		return fmt.Errorf("inline must be provided")
+	} else if s.Executor == "" {
+		s.Executor = "bash"
 	}
 
 	if _, err := exec.LookPath(s.Executor); err != nil {
-		logging.L().Error(zap.Error(err))
-		return err
+		return fmt.Errorf("executor not found: %w", err)
 	}
 
-	logging.L().Debugw("command found in path", "executor", s.Executor)
 	return nil
 }
 
@@ -140,14 +141,17 @@ func (s *ExpectStep) Validate(_ TTPExecutionContext) error {
 // *ActResult: A pointer to the action result.
 // error: An error if execution fails.
 func (s *ExpectStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
-	if s == nil {
-		return nil, fmt.Errorf("expectStep is nil")
+	if s == nil || s.Expect == nil {
+		return nil, fmt.Errorf("expect block must be provided")
 	}
 
 	originalDir, err := os.Getwd()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get current directory: %w", err)
 	}
+	if err := s.Validate(execCtx); err != nil {
+		return nil, err
+	}
 	defer func() {
 		if err := os.Chdir(originalDir); err != nil {
 			fmt.Printf("failed to change back to original directory: %v\n", err)
@@ -166,8 +170,16 @@ func (s *ExpectStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
 	}
 	defer console.Close()
 
+	if s.Environment != nil {
+		for k, v := range s.Environment {
+			if err := os.Setenv(k, v); err != nil {
+				return nil, fmt.Errorf("failed to set environment variable: %w", err)
+			}
+		}
+	}
+
 	envAsList := os.Environ()
-	cmd := s.prepareCommand(context.Background(), execCtx, envAsList, s.Inline)
+	cmd := s.prepareCommand(context.Background(), execCtx, envAsList, s.Expect.Inline)
 	cmd.Stdin = console.Tty()
 	cmd.Stdout = console.Tty()
 	cmd.Stderr = console.Tty()
@@ -178,41 +190,48 @@ func (s *ExpectStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
 
 	done := make(chan error, 1)
 	go func() {
-		for _, response := range s.Responses {
+		defer close(done)
+		for _, response := range s.Expect.Responses {
+			logging.L().Infof("Waiting for prompt: %s\n", response.Prompt)
 			re := regexp.MustCompile(response.Prompt)
-			if _, err := console.Expect(expect.Regexp(re)); err != nil {
+			timeout := 120 // Default timeout is 120 seconds
+			if s.Timeout > 0 {
+				timeout = s.Timeout // Use the provided timeout if it is greater than 0
+			}
+			matched, err := console.Expect(expect.Regexp(re), expect.WithTimeout(time.Duration(timeout)*time.Second))
+			if err != nil {
 				done <- fmt.Errorf("failed to expect %q: %w", re, err)
 				return
 			}
+			logging.L().Infof("Matched prompt: %s\n", matched)
+			logging.L().Infof("Sending response: %s\n", response.Response)
 			if _, err := console.SendLine(response.Response); err != nil {
 				done <- fmt.Errorf("failed to send response: %w", err)
 				return
 			}
 		}
-		// Close the console to send EOF
+
+		logging.L().Info("Closing console TTY...")
 		if err := console.Tty().Close(); err != nil {
 			done <- fmt.Errorf("failed to close console Tty: %w", err)
 			return
 		}
+
+		logging.L().Info("Waiting for command to exit...")
+		if err := cmd.Wait(); err != nil {
+			done <- fmt.Errorf("command failed: %w", err)
+			return
+		}
 		done <- nil
 	}()
 
-	timeout := 30 * time.Second
-	if s.Timeout != 0 {
-		timeout = time.Duration(s.Timeout) * time.Second
-	}
-
 	select {
 	case err := <-done:
 		if err != nil {
-			return nil, fmt.Errorf("error in expect: %w", err)
+			return nil, err
 		}
-	case <-time.After(timeout):
-		return nil, fmt.Errorf("timeout waiting for expect")
-	}
-
-	if err := cmd.Wait(); err != nil {
-		return nil, fmt.Errorf("command wait failed: %w", err)
+	case <-time.After(120 * time.Second):
+		return nil, fmt.Errorf("command timed out")
 	}
 
 	if _, err := console.ExpectEOF(); err != nil {
@@ -240,7 +259,7 @@ func (s *ExpectStep) prepareCommand(ctx context.Context, execCtx TTPExecutionCon
 	cmd := exec.CommandContext(ctx, s.Executor, "-c", inline)
 	cmd.Env = envAsList
 	cmd.Dir = execCtx.WorkDir
-	cmd.Stdin = strings.NewReader(inline)
+
 	return cmd
 }
 
@@ -260,17 +279,21 @@ func (s *ExpectStep) Cleanup(execCtx TTPExecutionContext) (*ActResult, error) {
 		return &ActResult{}, nil
 	}
 
+	logging.L().Info("Running cleanup step")
+
 	envAsList := os.Environ()
 	cmd := s.prepareCommand(context.Background(), execCtx, envAsList, s.CleanupStep)
 
-	cmd.Stdin = os.Stdin
 	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
 }