From ba9ca3d0e10d2d8c38c3ce5f385ff186a56fd34c Mon Sep 17 00:00:00 2001
From: Nikita Podshivalov <psihachina@gmail.com>
Date: Tue, 15 Nov 2022 18:12:53 +0300
Subject: [PATCH] fixed work ssm-plugin in non-interactive env

---
 pkg/ssmsession/ssmsession.go | 44 ++++++++++++++++++++++++++++++++++--
 tests/e2e/ecs_apps_test.go   |  4 ++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/pkg/ssmsession/ssmsession.go b/pkg/ssmsession/ssmsession.go
index f13898d7..7684aeac 100644
--- a/pkg/ssmsession/ssmsession.go
+++ b/pkg/ssmsession/ssmsession.go
@@ -1,12 +1,16 @@
 package ssmsession
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
+	"log"
 	"os"
 	"os/exec"
 	"strings"
+	"syscall"
 
+	"github.com/Netflix/go-expect"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/hazelops/ize/pkg/term"
 )
@@ -34,16 +38,33 @@ func NewSSMPluginCommand(region string) SSMPluginCommand {
 }
 
 func (s SSMPluginCommand) Start(ssmSession *ecs.Session) error {
+	var output bytes.Buffer
 	response, err := json.Marshal(ssmSession)
 	if err != nil {
 		return fmt.Errorf("marshal session response: %w", err)
 	}
+
+	c, err := expect.NewConsole(expect.WithStdout(&output), expect.WithStdin(os.Stdin))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer c.Close()
+
 	cmd := exec.Command(ssmPluginBinaryName, []string{string(response), s.region, startSessionAction}...)
-	out, _, _, err := s.runner.Run(cmd)
+	cmd.Stdin = c.Tty()
+	cmd.Stdout = c.Tty()
+	cmd.Stderr = c.Tty()
+
+	go func() {
+		c.ExpectEOF()
+	}()
+
+	_, _, _, err = s.Run(cmd)
 	if err != nil {
 		return fmt.Errorf("start session: %w", err)
 	}
-	if strings.Contains(out, "ERROR") {
+	fmt.Println(strings.TrimSpace(output.String()))
+	if strings.Contains(output.String(), "ERROR") {
 		return fmt.Errorf("exit status: 1")
 	}
 
@@ -63,3 +84,22 @@ func (s SSMPluginCommand) StartInteractive(ssmSession *ecs.Session) error {
 
 	return nil
 }
+
+func (s SSMPluginCommand) Run(cmd *exec.Cmd) (stdout, stderr string, exitCode int, err error) {
+
+	if err = cmd.Start(); err != nil {
+		return
+	}
+
+	err = cmd.Wait()
+
+	if err != nil {
+		if err2, ok := err.(*exec.ExitError); ok {
+			if s, ok := err2.Sys().(syscall.WaitStatus); ok {
+				err = nil
+				exitCode = s.ExitStatus()
+			}
+		}
+	}
+	return
+}
diff --git a/tests/e2e/ecs_apps_test.go b/tests/e2e/ecs_apps_test.go
index 3d2c3706..6a44e3fb 100644
--- a/tests/e2e/ecs_apps_test.go
+++ b/tests/e2e/ecs_apps_test.go
@@ -351,7 +351,7 @@ func TestIzeExecGoblin(t *testing.T) {
 
 	ize := NewBinary(t, izeBinary, examplesRootDir)
 
-	stdout, stderr, err := ize.RunPty("--plain-text-output", "exec", "goblin", "--", "sh -c \"echo $APP_NAME\"")
+	stdout, stderr, err := ize.RunRaw("--plain-text-output", "exec", "goblin", "--", "sh -c \"echo $APP_NAME\"")
 
 	if err != nil {
 		t.Errorf("error: %s", err)
@@ -379,7 +379,7 @@ func TestIzeExecSquibby(t *testing.T) {
 
 	ize := NewBinary(t, izeBinary, examplesRootDir)
 
-	stdout, stderr, err := ize.RunPty("--plain-text-output", "exec", "squibby", "--", "sh -c \"echo $APP_NAME\"")
+	stdout, stderr, err := ize.RunRaw("--plain-text-output", "exec", "squibby", "--", "sh -c \"echo $APP_NAME\"")
 
 	if err != nil {
 		t.Errorf("error: %s", err)