Support pass start session response with env variable

Author: Yangtao-Hua

src/sessionmanagerplugin/session/session.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/aws/session-manager-plugin/src/config"
@@ -162,7 +163,14 @@ func ValidateInputAndStartSession(args []string, out io.Writer) {
 	for argsIndex := 1; argsIndex < len(args); argsIndex++ {
 		switch argsIndex {
 		case 1:
-			response = []byte(args[1])
+			if strings.HasPrefix(args[1], "AWS_SSM_START_SESSION_RESPONSE") == true {
+				response = []byte(os.Getenv(args[1]))
+				if err = os.Unsetenv(args[1]); err != nil {
+					log.Errorf("Failed to remove temporary session env parameter: %v", err)
+				}
+			} else {
+				response = []byte(args[1])
+			}
 		case 2:
 			region = args[2]
 		case 3:

src/sessionmanagerplugin/session/session_test.go

@@ -17,6 +17,7 @@ package session
 import (
 	"bytes"
 	"fmt"
+	"os"
 	"sync"
 	"testing"
 	"time"
@@ -61,6 +62,47 @@ func TestValidateInputAndStartSession(t *testing.T) {
 	assert.Contains(t, buffer.String(), "Cannot perform start session: Some error")
 }
 
+func TestValidateInputAndStartSessionWithEnvVariableParameter(t *testing.T) {
+	var buffer bytes.Buffer
+	sessionResponse := "{\"SessionId\": \"user-012345\", \"TokenValue\": \"Session-Token\", \"StreamUrl\": \"wss://ssmmessages.us-east-1.amazonaws.com/v1/data-channel/user-012345?role=publish_subscribe\"}"
+	os.Setenv("AWS_SSM_START_SESSION_RESPONSE", sessionResponse)
+	args := []string{"session-manager-plugin",
+		"AWS_SSM_START_SESSION_RESPONSE",
+		"us-east-1", "StartSession", "", "{\"Target\": \"i-0123abc\"}", "https://ssm.us-east-1.amazonaws.com"}
+	parameterPassed := false
+	startSession = func(session *Session, log log.T) error {
+		if session.TokenValue == "Session-Token" && session.SessionId == "user-012345" {
+			parameterPassed = true
+		}
+		return nil
+	}
+
+	ValidateInputAndStartSession(args, &buffer)
+	var _, envVariableExist = os.LookupEnv("AWS_SSM_START_SESSION_RESPONSE")
+	assert.False(t, envVariableExist)
+	assert.True(t, parameterPassed)
+}
+
+func TestValidateInputAndStartSessionWithWrongEnvVariableName(t *testing.T) {
+	var buffer bytes.Buffer
+	sessionResponse := "{\"SessionId\": \"user-012345\", \"TokenValue\": \"Session-Token\", \"StreamUrl\": \"wss://ssmmessages.us-east-1.amazonaws.com/v1/data-channel/user-012345?role=publish_subscribe\"}"
+	os.Setenv("WRONG_ENV_NAME", sessionResponse)
+	args := []string{"session-manager-plugin",
+		"WRONG_ENV_NAME",
+		"us-east-1", "StartSession", "", "{\"Target\": \"i-0123abc\"}", "https://ssm.us-east-1.amazonaws.com"}
+	startSessionInvoked := false
+	startSession = func(session *Session, log log.T) error {
+		startSessionInvoked = true
+		return nil
+	}
+
+	ValidateInputAndStartSession(args, &buffer)
+	var _, envVariableExist = os.LookupEnv("WRONG_ENV_NAME")
+	assert.Contains(t, buffer.String(), "Cannot perform start session: invalid character 'W'")
+	assert.True(t, envVariableExist)
+	assert.False(t, startSessionInvoked)
+}
+
 func TestExecute(t *testing.T) {
 	sessionMock := &Session{}
 	sessionMock.DataChannel = mockDataChannel