From cddf64fb62c9cd86a8e11767e450ee2617342a66 Mon Sep 17 00:00:00 2001
From: IndranilVyas <indranil.vyas@gmail.com>
Date: Thu, 28 Nov 2019 19:30:34 -0500
Subject: [PATCH] Add 2FA support

---
 cloudwatch/client.go | 9 ++++++++-
 main.go              | 3 ++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/cloudwatch/client.go b/cloudwatch/client.go
index 19aaca7b..89f4e454 100644
--- a/cloudwatch/client.go
+++ b/cloudwatch/client.go
@@ -5,10 +5,12 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 )
 
 //CW provides the APIo peration methods for making requests to AWS cloudwatch logs.
@@ -18,7 +20,7 @@ type CW struct {
 }
 
 // New creates a new instance of the CW client
-func New(awsProfile *string, awsRegion *string, log *log.Logger) *CW {
+func New(awsProfile *string, awsRegion *string, mfa *bool, log *log.Logger) *CW {
 	//workaround to figure out the user actual home dir within a SNAP (rather than the sandboxed one)
 	//and access the  .aws folder in its default location
 	if os.Getenv("SNAP_INSTANCE_NAME") != "" {
@@ -48,6 +50,11 @@ func New(awsProfile *string, awsRegion *string, log *log.Logger) *CW {
 	if awsRegion != nil {
 		opts.Config = aws.Config{Region: awsRegion}
 	}
+	if *mfa {
+		duration, _ := time.ParseDuration("1h")
+		opts.AssumeRoleTokenProvider = stscreds.StdinTokenProvider
+		opts.AssumeRoleDuration = duration
+	}
 
 	sess := session.Must(session.NewSessionWithOptions(opts))
 	return &CW{awsClwClient: cloudwatchlogs.New(sess),
diff --git a/main.go b/main.go
index 0f382f12..a2e66630 100644
--- a/main.go
+++ b/main.go
@@ -64,6 +64,7 @@ var (
 	grep  = tailCommand.Flag("grep", "Pattern to filter logs by. See http://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/FilterAndPatternSyntax.html for syntax.").
 		Short('g').Default("").String()
 	grepv = tailCommand.Flag("grepv", "Equivalent of grep --invert-match. Invert match pattern to filter logs by.").Short('v').Default("").String()
+	mfa = kp.Flag("mfa", "When target AWS profile has MFA enabled.").Short('m').Default("false").Bool()
 )
 
 func timestampToTime(timeStamp *string) (time.Time, error) {
@@ -173,7 +174,7 @@ func main() {
 		color.NoColor = true
 	}
 
-	c := cloudwatch.New(awsProfile, awsRegion, log)
+	c := cloudwatch.New(awsProfile, awsRegion, mfa, log)
 
 	switch cmd {
 	case "ls groups":