Repository: Versent/saml2aws
Pull request #406 (2020-01-17 11:27:16)

Allow passing MFA token for GoogleApps
Uses token from --mfa-token flag instead of prompting if flag present.


Patch:
diff --git a/README.md b/README.md
index 69307be..e1ef188 100644
--- a/README.md
+++ b/README.md
@@ -135,7 +135,7 @@ Flags:
       --password=PASSWORD      The password used to login. (env:
                                SAML2AWS_PASSWORD)
       --mfa-token=MFA-TOKEN    The current MFA token (supported in Keycloak,
-                               ADFS). (env: SAML2AWS_MFA_TOKEN)
+                               ADFS, GoogleApps). (env: SAML2AWS_MFA_TOKEN)
       --role=ROLE              The ARN of the role to assume. (env:
                                SAML2AWS_ROLE)
       --aws-urn=AWS-URN        The URN used by SAML when you login. (env:
diff --git a/cmd/saml2aws/main.go b/cmd/saml2aws/main.go
index 57b5dc1..87058ef 100644
--- a/cmd/saml2aws/main.go
+++ b/cmd/saml2aws/main.go
@@ -59,7 +59,7 @@ func main() {
 	app.Flag("url", "The URL of the SAML IDP server used to login. (env: SAML2AWS_URL)").Envar("SAML2AWS_URL").StringVar(&commonFlags.URL)
 	app.Flag("username", "The username used to login. (env: SAML2AWS_USERNAME)").Envar("SAML2AWS_USERNAME").StringVar(&commonFlags.Username)
 	app.Flag("password", "The password used to login. (env: SAML2AWS_PASSWORD)").Envar("SAML2AWS_PASSWORD").StringVar(&commonFlags.Password)
-	app.Flag("mfa-token", "The current MFA token (supported in Keycloak, ADFS). (env: SAML2AWS_MFA_TOKEN)").Envar("SAML2AWS_MFA_TOKEN").StringVar(&commonFlags.MFAToken)
+	app.Flag("mfa-token", "The current MFA token (supported in Keycloak, ADFS, GoogleApps). (env: SAML2AWS_MFA_TOKEN)").Envar("SAML2AWS_MFA_TOKEN").StringVar(&commonFlags.MFAToken)
 	app.Flag("role", "The ARN of the role to assume. (env: SAML2AWS_ROLE)").Envar("SAML2AWS_ROLE").StringVar(&commonFlags.RoleArn)
 	app.Flag("aws-urn", "The URN used by SAML when you login. (env: SAML2AWS_AWS_URN)").Envar("SAML2AWS_AWS_URN").StringVar(&commonFlags.AmazonWebservicesURN)
 	app.Flag("skip-prompt", "Skip prompting for parameters during login.").BoolVar(&commonFlags.SkipPrompt)
diff --git a/pkg/provider/googleapps/googleapps.go b/pkg/provider/googleapps/googleapps.go
index 1202e8b..2af5a87 100644
--- a/pkg/provider/googleapps/googleapps.go
+++ b/pkg/provider/googleapps/googleapps.go
@@ -61,7 +61,7 @@ func (kc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 	authForm.Set("Passwd", loginDetails.Password)
 	authForm.Set("rawidentifier", loginDetails.Username)
 
-	responseDoc, err := kc.loadChallengePage(passwordURL+"?hl=en&loc=US", authURL, authForm)
+	responseDoc, err := kc.loadChallengePage(passwordURL+"?hl=en&loc=US", authURL, authForm, loginDetails)
 	if err != nil {
 		return "", errors.Wrap(err, "error loading challenge page")
 	}
@@ -91,7 +91,7 @@ func (kc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 		captchaForm.Set("Passwd", loginDetails.Password)
 		captchaForm.Set("logincaptcha", captcha)
 
-		responseDoc, err = kc.loadChallengePage(captchaURL+"?hl=en&loc=US", captchaURL, captchaForm)
+		responseDoc, err = kc.loadChallengePage(captchaURL+"?hl=en&loc=US", captchaURL, captchaForm, loginDetails)
 		if err != nil {
 			return "", errors.Wrap(err, "error loading challenge page")
 		}
@@ -186,7 +186,7 @@ func (kc *Client) loadLoginPage(submitURL string, referer string, authForm url.V
 	return loginURL, loginForm, err
 }
 
-func (kc *Client) loadChallengePage(submitURL string, referer string, authForm url.Values) (*goquery.Document, error) {
+func (kc *Client) loadChallengePage(submitURL string, referer string, authForm url.Values, loginDetails *creds.LoginDetails) (*goquery.Document, error) {
 
 	req, err := http.NewRequest("POST", submitURL, strings.NewReader(authForm.Encode()))
 	if err != nil {
@@ -236,7 +236,10 @@ func (kc *Client) loadChallengePage(submitURL string, referer string, authForm u
 		switch {
 		case strings.Contains(secondActionURL, "challenge/totp/"): // handle TOTP challenge
 
-			var token = prompter.RequestSecurityCode("000000")
+			var token = loginDetails.MFAToken
+			if token == "" {
+				token = prompter.RequestSecurityCode("000000")
+			}
 
 			responseForm.Set("Pin", token)
 			responseForm.Set("TrustDevice", "on") // Don't ask again on this computer
@@ -303,7 +306,7 @@ func (kc *Client) loadChallengePage(submitURL string, referer string, authForm u
 
 		u.Path = skipActionURL
 
-		return kc.loadAlternateChallengePage(u.String(), submitURL, skipResponseForm)
+		return kc.loadAlternateChallengePage(u.String(), submitURL, skipResponseForm, loginDetails)
 
 	}
 
@@ -311,7 +314,7 @@ func (kc *Client) loadChallengePage(submitURL string, referer string, authForm u
 
 }
 
-func (kc *Client) loadAlternateChallengePage(submitURL string, referer string, authForm url.Values) (*goquery.Document, error) {
+func (kc *Client) loadAlternateChallengePage(submitURL string, referer string, authForm url.Values, loginDetails *creds.LoginDetails) (*goquery.Document, error) {
 
 	req, err := http.NewRequest("POST", submitURL, strings.NewReader(authForm.Encode()))
 	if err != nil {
@@ -365,7 +368,7 @@ func (kc *Client) loadAlternateChallengePage(submitURL string, referer string, a
 	u, _ := url.Parse(submitURL)
 	u.Path = newActionURL
 
-	return kc.loadChallengePage(u.String(), submitURL, responseForm)
+	return kc.loadChallengePage(u.String(), submitURL, responseForm, loginDetails)
 }
 
 func (kc *Client) postJSON(submitURL string, values map[string]string, referer string) (*http.Response, error) {


Test patch:
diff --git a/pkg/provider/googleapps/googleapps_test.go b/pkg/provider/googleapps/googleapps_test.go
index c9c9bdc..b5ec1e7 100644
--- a/pkg/provider/googleapps/googleapps_test.go
+++ b/pkg/provider/googleapps/googleapps_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/PuerkitoBio/goquery"
 	"github.com/stretchr/testify/require"
+	"github.com/versent/saml2aws/pkg/creds"
 	"github.com/versent/saml2aws/pkg/provider"
 )
 
@@ -81,10 +82,10 @@ func TestChallengePage(t *testing.T) {
 	defer ts.Close()
 
 	kc := Client{client: &provider.HTTPClient{Client: http.Client{}}}
-	// loginDetails := &creds.LoginDetails{URL: ts.URL, Username: "test", Password: "test123"}
+	loginDetails := &creds.LoginDetails{URL: ts.URL, Username: "test", Password: "test123"}
 	authForm := url.Values{}
 
-	challengeDoc, err := kc.loadChallengePage(ts.URL, "https://accounts.google.com/signin/challenge/sl/password", authForm)
+	challengeDoc, err := kc.loadChallengePage(ts.URL, "https://accounts.google.com/signin/challenge/sl/password", authForm, loginDetails)
 	require.Nil(t, err)
 	require.NotNil(t, challengeDoc)
 }