Repository: Versent/saml2aws
Pull request #176 (2018-05-15 17:47:42)

Add global mfa-token flag
This flag allows easier automation for running saml2aws with
--skip-prompt. Only Keycloak and ADFS are currently supported.

---
In my exact use case, [saml2aws-auto](https://github.com/Rukenshia/saml2aws-auto), I currently have to rely on passing the MFA token by sending it through `stdin`. This does, however, not work reliably on Windows. I think this addition is quite useful for `--skip-prompt` use cases.


Patch:
diff --git a/README.md b/README.md
index 4419364..b10daa3 100644
--- a/README.md
+++ b/README.md
@@ -79,6 +79,7 @@ Flags:
       --url=URL                The URL of the SAML IDP server used to login.
       --username=USERNAME      The username used to login.
       --password=PASSWORD      The password used to login.
+      --mfa-token=MFA-TOKEN    The current MFA token (supported in Keycloak, ADFS).
       --role=ROLE              The ARN of the role to assume.
       --aws-urn=AWS-URN        The URN used by SAML when you login.
       --skip-prompt            Skip prompting for parameters during login.
diff --git a/cmd/saml2aws/commands/login.go b/cmd/saml2aws/commands/login.go
index e83f92e..c884929 100644
--- a/cmd/saml2aws/commands/login.go
+++ b/cmd/saml2aws/commands/login.go
@@ -124,7 +124,7 @@ func resolveLoginDetails(account *cfg.IDPAccount, loginFlags *flags.LoginExecFla
 
 	// fmt.Printf("loginFlags %+v\n", loginFlags)
 
-	loginDetails := &creds.LoginDetails{URL: account.URL, Username: account.Username}
+	loginDetails := &creds.LoginDetails{URL: account.URL, Username: account.Username, MFAToken: loginFlags.CommonFlags.MFAToken}
 
 	fmt.Printf("Using IDP Account %s to access %s %s\n", loginFlags.CommonFlags.IdpAccount, account.Provider, account.URL)
 
diff --git a/cmd/saml2aws/main.go b/cmd/saml2aws/main.go
index 7d64919..1c8257f 100644
--- a/cmd/saml2aws/main.go
+++ b/cmd/saml2aws/main.go
@@ -57,6 +57,7 @@ func main() {
 	app.Flag("url", "The URL of the SAML IDP server used to login.").StringVar(&commonFlags.URL)
 	app.Flag("username", "The username used to login.").Envar("SAML2AWS_USERNAME").StringVar(&commonFlags.Username)
 	app.Flag("password", "The password used to login.").Envar("SAML2AWS_PASSWORD").StringVar(&commonFlags.Password)
+	app.Flag("mfa-token", "The current MFA token (supported in Keycloak, ADFS).").Envar("SAML2AWS_MFA_TOKEN").StringVar(&commonFlags.MFAToken)
 	app.Flag("role", "The ARN of the role to assume.").StringVar(&commonFlags.RoleArn)
 	app.Flag("aws-urn", "The URN used by SAML when you login.").StringVar(&commonFlags.AmazonWebservicesURN)
 	app.Flag("skip-prompt", "Skip prompting for parameters during login.").BoolVar(&commonFlags.SkipPrompt)
diff --git a/pkg/creds/creds.go b/pkg/creds/creds.go
index 3e695bf..45120f0 100644
--- a/pkg/creds/creds.go
+++ b/pkg/creds/creds.go
@@ -6,6 +6,7 @@ import "errors"
 type LoginDetails struct {
 	Username string
 	Password string
+	MFAToken string
 	URL      string
 }
 
diff --git a/pkg/flags/flags.go b/pkg/flags/flags.go
index 543c39a..eccc01e 100644
--- a/pkg/flags/flags.go
+++ b/pkg/flags/flags.go
@@ -9,6 +9,7 @@ type CommonFlags struct {
 	IdpAccount           string
 	IdpProvider          string
 	MFA                  string
+	MFAToken             string
 	URL                  string
 	Username             string
 	Password             string
diff --git a/pkg/provider/adfs/adfs.go b/pkg/provider/adfs/adfs.go
index 7e5110c..97e523c 100644
--- a/pkg/provider/adfs/adfs.go
+++ b/pkg/provider/adfs/adfs.go
@@ -96,7 +96,7 @@ func (ac *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 
 	switch ac.idpAccount.MFA {
 	case "VIP":
-		res, err = ac.vipMFA(authSubmitURL, res)
+		res, err = ac.vipMFA(authSubmitURL, loginDetails.MFAToken, res)
 		if err != nil {
 			return samlAssertion, errors.Wrap(err, "error retrieving mfa form results")
 		}
@@ -127,7 +127,7 @@ func (ac *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 
 // vipMFA when supplied with the the form response document attempt to extract the VIP mfa related field
 // then use that to trigger a submit of the MFA security token
-func (ac *Client) vipMFA(authSubmitURL string, res *http.Response) (*http.Response, error) {
+func (ac *Client) vipMFA(authSubmitURL string, mfaToken string, res *http.Response) (*http.Response, error) {
 
 	doc, err := goquery.NewDocumentFromResponse(res)
 	if err != nil {
@@ -142,10 +142,12 @@ func (ac *Client) vipMFA(authSubmitURL string, res *http.Response) (*http.Respon
 		return res, nil // if we didn't find the MFA flag then just continue
 	}
 
-	var token = prompter.RequestSecurityCode("000000")
+	if mfaToken == "" {
+		mfaToken = prompter.RequestSecurityCode("000000")
+	}
 
 	doc.Find("input").Each(func(i int, s *goquery.Selection) {
-		updateOTPFormData(otpForm, s, token)
+		updateOTPFormData(otpForm, s, mfaToken)
 	})
 
 	doc.Find("form").Each(func(i int, s *goquery.Selection) {
diff --git a/pkg/provider/keycloak/keycloak.go b/pkg/provider/keycloak/keycloak.go
index 4d4e1ea..d6a3062 100644
--- a/pkg/provider/keycloak/keycloak.go
+++ b/pkg/provider/keycloak/keycloak.go
@@ -68,7 +68,7 @@ func (kc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 			return "", errors.Wrap(err, "unable to locate IDP totp form submit URL")
 		}
 
-		doc, err = kc.postTotpForm(totpSubmitURL, doc)
+		doc, err = kc.postTotpForm(totpSubmitURL, loginDetails.MFAToken, doc)
 		if err != nil {
 			return "", errors.Wrap(err, "error posting totp form")
 		}
@@ -141,14 +141,16 @@ func (kc *Client) postLoginForm(authSubmitURL string, authForm url.Values) ([]by
 	return data, nil
 }
 
-func (kc *Client) postTotpForm(totpSubmitURL string, doc *goquery.Document) (*goquery.Document, error) {
+func (kc *Client) postTotpForm(totpSubmitURL string, mfaToken string, doc *goquery.Document) (*goquery.Document, error) {
 
 	otpForm := url.Values{}
 
-	var token = prompter.RequestSecurityCode("000000")
+	if mfaToken == "" {
+		mfaToken = prompter.RequestSecurityCode("000000")
+	}
 
 	doc.Find("input").Each(func(i int, s *goquery.Selection) {
-		updateOTPFormData(otpForm, s, token)
+		updateOTPFormData(otpForm, s, mfaToken)
 	})
 
 	req, err := http.NewRequest("POST", totpSubmitURL, strings.NewReader(otpForm.Encode()))


Test patch:
diff --git a/cmd/saml2aws/commands/login_test.go b/cmd/saml2aws/commands/login_test.go
index 6df8239..4fc480f 100644
--- a/cmd/saml2aws/commands/login_test.go
+++ b/cmd/saml2aws/commands/login_test.go
@@ -12,7 +12,7 @@ import (
 
 func TestResolveLoginDetailsWithFlags(t *testing.T) {
 
-	commonFlags := &flags.CommonFlags{URL: "https://id.example.com", Username: "wolfeidau", Password: "testtestlol", SkipPrompt: true}
+	commonFlags := &flags.CommonFlags{URL: "https://id.example.com", Username: "wolfeidau", Password: "testtestlol", MFAToken: "123456", SkipPrompt: true}
 	loginFlags := &flags.LoginExecFlags{CommonFlags: commonFlags}
 
 	idpa := &cfg.IDPAccount{
@@ -24,7 +24,7 @@ func TestResolveLoginDetailsWithFlags(t *testing.T) {
 	loginDetails, err := resolveLoginDetails(idpa, loginFlags)
 
 	assert.Empty(t, err)
-	assert.Equal(t, &creds.LoginDetails{Username: "wolfeidau", Password: "testtestlol", URL: "https://id.example.com"}, loginDetails)
+	assert.Equal(t, &creds.LoginDetails{Username: "wolfeidau", Password: "testtestlol", URL: "https://id.example.com", MFAToken: "123456"}, loginDetails)
 }
 
 func TestResolveRoleSingleEntry(t *testing.T) {
diff --git a/pkg/provider/keycloak/keycloak_test.go b/pkg/provider/keycloak/keycloak_test.go
index cec3de1..e5bbc48 100644
--- a/pkg/provider/keycloak/keycloak_test.go
+++ b/pkg/provider/keycloak/keycloak_test.go
@@ -10,10 +10,11 @@ import (
 
 	"github.com/PuerkitoBio/goquery"
 
-	"github.com/stretchr/testify/require"
 	"github.com/versent/saml2aws/mocks"
 	"github.com/versent/saml2aws/pkg/creds"
+	"github.com/versent/saml2aws/pkg/prompter"
 	"github.com/versent/saml2aws/pkg/provider"
+	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -78,13 +79,42 @@ func TestClient_postTotpForm(t *testing.T) {
 
 	doc, err := goquery.NewDocumentFromReader(bytes.NewReader(data))
 	require.Nil(t, err)
+
 	pr := &mocks.Prompter{}
+	prompter.SetPrompter(pr)
 
-	pr.Mock.On("RequestSecurityCode", "000000").Return("111222")
+	pr.Mock.On("RequestSecurityCode", "000000").Return("123456")
 
+	mfaToken := ""
 	kc := Client{client: &provider.HTTPClient{Client: http.Client{}}}
 
-	kc.postTotpForm(ts.URL, doc)
+	kc.postTotpForm(ts.URL, mfaToken, doc)
+
+	pr.Mock.AssertCalled(t, "RequestSecurityCode", "000000")
+}
+
+func TestClient_postTotpFormWithProvidedMFAToken(t *testing.T) {
+
+	data, err := ioutil.ReadFile("example/assertion.html")
+	require.Nil(t, err)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write(data)
+	}))
+	defer ts.Close()
+
+	doc, err := goquery.NewDocumentFromReader(bytes.NewReader(data))
+	require.Nil(t, err)
+
+	pr := &mocks.Prompter{}
+	prompter.SetPrompter(pr)
+
+	mfaToken := "123456"
+	kc := Client{client: &provider.HTTPClient{Client: http.Client{}}}
+
+	kc.postTotpForm(ts.URL, mfaToken, doc)
+
+	pr.Mock.AssertNumberOfCalls(t, "RequestSecurityCode", 0)
 }
 
 func TestClient_containsTotpForm(t *testing.T) {