From 480768c3c92aeadeb5803847923146c50765a471 Mon Sep 17 00:00:00 2001
From: Michael Tibben <michael.tibben@99designs.com>
Date: Thu, 16 Feb 2023 21:08:26 +1100
Subject: [PATCH] Automatically select the best prompt driver

---
 cli/exec.go   |  6 +++++-
 cli/export.go |  2 +-
 cli/global.go | 23 ++++++++++++++++++++---
 cli/login.go  |  2 +-
 cli/rotate.go |  2 +-
 5 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/cli/exec.go b/cli/exec.go
index 1ff85f0d3..9842f31db 100644
--- a/cli/exec.go
+++ b/cli/exec.go
@@ -61,6 +61,10 @@ func (input ExecCommandInput) validate() error {
 	return nil
 }
 
+func CanExecUseTerminal(input ExecCommandInput) bool {
+	return !input.StartEcsServer && !input.StartEc2Server
+}
+
 func ConfigureExecCommand(app *kingpin.Application, a *AwsVault) {
 	input := ExecCommandInput{}
 
@@ -114,7 +118,7 @@ func ConfigureExecCommand(app *kingpin.Application, a *AwsVault) {
 		StringsVar(&input.Args)
 
 	cmd.Action(func(c *kingpin.ParseContext) (err error) {
-		input.Config.MfaPromptMethod = a.PromptDriver
+		input.Config.MfaPromptMethod = a.PromptDriver(CanExecUseTerminal(input))
 		input.Config.NonChainedGetSessionTokenDuration = input.SessionDuration
 		input.Config.AssumeRoleDuration = input.SessionDuration
 		input.Config.SSOUseStdout = input.UseStdout
diff --git a/cli/export.go b/cli/export.go
index a03b0d222..7644a677c 100644
--- a/cli/export.go
+++ b/cli/export.go
@@ -63,7 +63,7 @@ func ConfigureExportCommand(app *kingpin.Application, a *AwsVault) {
 		StringVar(&input.ProfileName)
 
 	cmd.Action(func(c *kingpin.ParseContext) (err error) {
-		input.Config.MfaPromptMethod = a.PromptDriver
+		input.Config.MfaPromptMethod = a.PromptDriver(true)
 		input.Config.NonChainedGetSessionTokenDuration = input.SessionDuration
 		input.Config.AssumeRoleDuration = input.SessionDuration
 		input.Config.SSOUseStdout = input.UseStdout
diff --git a/cli/global.go b/cli/global.go
index 5caaf0080..909b1f939 100644
--- a/cli/global.go
+++ b/cli/global.go
@@ -27,12 +27,30 @@ type AwsVault struct {
 	Debug          bool
 	KeyringConfig  keyring.Config
 	KeyringBackend string
-	PromptDriver   string
+	promptDriver   string
 
 	keyringImpl   keyring.Keyring
 	awsConfigFile *vault.ConfigFile
 }
 
+func (a *AwsVault) PromptDriver(canUseTerminal bool) string {
+	if a.promptDriver == "" {
+		if canUseTerminal {
+			return "terminal"
+		}
+		for _, driver := range prompt.Available() {
+			a.promptDriver = driver
+			if driver != "terminal" {
+				break
+			}
+		}
+	}
+
+	log.Println("Using prompt driver: " + a.promptDriver)
+
+	return a.promptDriver
+}
+
 func (a *AwsVault) Keyring() (keyring.Keyring, error) {
 	if a.keyringImpl == nil {
 		if a.KeyringBackend != "" {
@@ -89,9 +107,8 @@ func ConfigureGlobals(app *kingpin.Application) *AwsVault {
 		EnumVar(&a.KeyringBackend, backendsAvailable...)
 
 	app.Flag("prompt", fmt.Sprintf("Prompt driver to use %v", promptsAvailable)).
-		Default("terminal").
 		Envar("AWS_VAULT_PROMPT").
-		EnumVar(&a.PromptDriver, promptsAvailable...)
+		EnumVar(&a.promptDriver, promptsAvailable...)
 
 	app.Flag("keychain", "Name of macOS keychain to use, if it doesn't exist it will be created").
 		Default("aws-vault").
diff --git a/cli/login.go b/cli/login.go
index 64a8c3d65..e209cf98d 100644
--- a/cli/login.go
+++ b/cli/login.go
@@ -59,7 +59,7 @@ func ConfigureLoginCommand(app *kingpin.Application, a *AwsVault) {
 		StringVar(&input.ProfileName)
 
 	cmd.Action(func(c *kingpin.ParseContext) (err error) {
-		input.Config.MfaPromptMethod = a.PromptDriver
+		input.Config.MfaPromptMethod = a.PromptDriver(true)
 		input.Config.NonChainedGetSessionTokenDuration = input.SessionDuration
 		input.Config.AssumeRoleDuration = input.SessionDuration
 		input.Config.GetFederationTokenDuration = input.SessionDuration
diff --git a/cli/rotate.go b/cli/rotate.go
index a077aa97f..aa419f031 100644
--- a/cli/rotate.go
+++ b/cli/rotate.go
@@ -34,7 +34,7 @@ func ConfigureRotateCommand(app *kingpin.Application, a *AwsVault) {
 		StringVar(&input.ProfileName)
 
 	cmd.Action(func(c *kingpin.ParseContext) (err error) {
-		input.Config.MfaPromptMethod = a.PromptDriver
+		input.Config.MfaPromptMethod = a.PromptDriver(true)
 		keyring, err := a.Keyring()
 		if err != nil {
 			return err