From 6d84c17449d1c21559be342885dca527a4100720 Mon Sep 17 00:00:00 2001
From: Steve Engledow <steve@offend.me.uk>
Date: Wed, 13 Nov 2019 14:10:38 +0000
Subject: [PATCH] Add MFA support

---
 client/common.go           | 15 ++++++++++++++-
 console/spinner/spinner.go |  8 ++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/client/common.go b/client/common.go
index 908ab352..26ff0a05 100644
--- a/client/common.go
+++ b/client/common.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/aws-cloudformation/rain/config"
+	"github.com/aws-cloudformation/rain/console"
 	"github.com/aws-cloudformation/rain/console/run"
 	"github.com/aws-cloudformation/rain/console/spinner"
 	"github.com/aws-cloudformation/rain/version"
@@ -21,6 +22,16 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws/external"
 )
 
+func MFAProvider() (string, error) {
+	spinner.Pause()
+	defer func() {
+		fmt.Println()
+		spinner.Resume()
+	}()
+
+	return console.Ask("MFA Token:"), nil
+}
+
 var awsCfg *aws.Config
 
 var defaultPython = []string{"/usr/bin/env", "python"}
@@ -77,7 +88,9 @@ func loadConfig() aws.Config {
 	var cfg aws.Config
 	var err error
 
-	configs := make([]external.Config, 0)
+	configs := []external.Config{
+		external.WithMFATokenFunc(MFAProvider),
+	}
 
 	if config.Profile != "" {
 		configs = append(configs, external.WithSharedConfigProfile(config.Profile))
diff --git a/console/spinner/spinner.go b/console/spinner/spinner.go
index 8174b0da..1962e250 100644
--- a/console/spinner/spinner.go
+++ b/console/spinner/spinner.go
@@ -71,6 +71,14 @@ func Stop() {
 	console.ClearLine()
 }
 
+func Pause() {
+	spinRunning = false
+}
+
+func Resume() {
+	spinRunning = true
+}
+
 // Update redisplays the spinner. Use this if your programme has updated
 // the screen and may have interfered with the spinners display location.
 func Update() {