Repository: Versent/saml2aws
Pull request #369 (2019-10-28 00:25:42)

Fixes Okta MFA so that it works with TOTP
I wasn't able to get `saml2aws` to work with Okta using MFA for our organization. We have the "Okta Verify" MFA (in reality just a TOTP scheme) configured for our AWS account and I was getting the "response does not contain a valid SAML assertion" error from https://github.com/Versent/saml2aws/issues/118.

According to [Okta support](https://support.okta.com/help/s/article/Integrating-the-Amazon-Web-Services-Command-Line-Interface-Using-Okta) the "preferred solution" is gimme-aws-creds. I looked at [their code](https://github.com/Nike-Inc/gimme-aws-creds/blob/master/gimme_aws_creds/okta.py#L543-L548) and see they grabbed the state token from the page and resubmit to the `authn` endpoint.

Using that approach, I was able to get `saml2aws` to ask for my MFA code and return a valid AWS session.

Happy to update or add any additional comments/documentation. Thanks!


Patch:
diff --git a/pkg/creds/creds.go b/pkg/creds/creds.go
index 223d5e1..f4d2006 100644
--- a/pkg/creds/creds.go
+++ b/pkg/creds/creds.go
@@ -11,6 +11,7 @@ type LoginDetails struct {
 	MFAToken     string
 	DuoMFAOption string
 	URL          string
+	StateToken   string // used by Okta
 }
 
 // Validate validate the login details
diff --git a/pkg/provider/okta/README.md b/pkg/provider/okta/README.md
index 6ffe8d7..8f819ee 100644
--- a/pkg/provider/okta/README.md
+++ b/pkg/provider/okta/README.md
@@ -12,8 +12,4 @@ The path segments `/home/amazon_aws` in the above URL may vary.
 
 ## Features
 
-* Supports MFA (Okta Push, Okta TOTP, Duo, and Google Authenticator), when configured at *organization level*.
-
-## Limitations
-
-* Does **not** support application-level MFA, per [issue #118](https://github.com/Versent/saml2aws/issues/118#issuecomment-355688008)
\ No newline at end of file
+* Supports MFA (Okta Push, Okta TOTP, Duo, and Google Authenticator), when configured at *organization* or *application* level.
\ No newline at end of file
diff --git a/pkg/provider/okta/okta.go b/pkg/provider/okta/okta.go
index 44deec5..37e4ca9 100644
--- a/pkg/provider/okta/okta.go
+++ b/pkg/provider/okta/okta.go
@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"regexp"
 	"strings"
 	"time"
 
@@ -51,8 +52,9 @@ type Client struct {
 
 // AuthRequest represents an mfa okta request
 type AuthRequest struct {
-	Username string `json:"username"`
-	Password string `json:"password"`
+	Username   string `json:"username"`
+	Password   string `json:"password"`
+	StateToken string `json:"stateToken,omitempty"`
 }
 
 // VerifyRequest represents an mfa verify request
@@ -95,6 +97,9 @@ func (oc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 
 	//authenticate via okta api
 	authReq := AuthRequest{Username: loginDetails.Username, Password: loginDetails.Password}
+	if (loginDetails.StateToken != "") {
+		authReq = AuthRequest{StateToken: loginDetails.StateToken}
+	}
 	authBody := new(bytes.Buffer)
 	err = json.NewEncoder(authBody).Encode(authReq)
 	if err != nil {
@@ -160,7 +165,24 @@ func (oc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 
 	samlAssertion, ok := doc.Find("input[name=\"SAMLResponse\"]").Attr("value")
 	if !ok {
-		return samlAssertion, errors.Wrap(err, "unable to locate saml response")
+		req, err = http.NewRequest("GET", loginDetails.URL, nil)
+		if err != nil {
+			return samlAssertion, errors.Wrap(err, "error building app request")
+		}
+		res, err = oc.client.Do(req)
+		if err != nil {
+			return samlAssertion, errors.Wrap(err, "error retrieving app response")
+		}
+		body, err := ioutil.ReadAll(res.Body)
+		if err != nil {
+			return "", errors.Wrap(err, "error retrieving body from response")
+		}		
+		stateToken, err := getStateTokenFromOktaPageBody(string(body))
+		if err != nil {
+			return "", errors.Wrap(err, "error retrieving saml response")
+		}
+		loginDetails.StateToken = stateToken		
+		return oc.Authenticate(loginDetails)
 	}
 
 	logger.Debug("auth complete")
@@ -168,6 +190,15 @@ func (oc *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 	return samlAssertion, nil
 }
 
+func getStateTokenFromOktaPageBody(responseBody string) (string, error) {
+	re := regexp.MustCompile("var stateToken = '(.*)';")
+	match := re.FindStringSubmatch(responseBody)
+	if len(match) < 2 {
+		return "", errors.New("cannot find state token")
+	}
+	return strings.Replace(match[1], `\x2D`, "-", -1), nil
+}
+
 func parseMfaIdentifer(json string, arrayPosition int) string {
 	mfaProvider := gjson.Get(json, fmt.Sprintf("_embedded.factors.%d.provider", arrayPosition)).String()
 	factorType := strings.ToUpper(gjson.Get(json, fmt.Sprintf("_embedded.factors.%d.factorType", arrayPosition)).String())


Test patch:
diff --git a/pkg/provider/okta/okta_test.go b/pkg/provider/okta/okta_test.go
new file mode 100644
index 0000000..ef86f1b
--- /dev/null
+++ b/pkg/provider/okta/okta_test.go
@@ -0,0 +1,48 @@
+package okta
+
+import (
+	"testing"
+	"errors"
+	"github.com/stretchr/testify/assert"
+)
+
+type stateTokenTests struct {
+	title string
+	body string
+	stateToken string
+	err error
+}
+
+func TestGetStateTokenFromOktaPageBody(t *testing.T) {
+	tests := []stateTokenTests{
+		{
+			title: "State token in body gets returned",
+			body: "someJavascriptCode();\nvar stateToken = '123456789';\nsomeOtherJavaScriptCode();",
+			stateToken: "123456789",
+			err: nil,
+		},
+		{
+			title: "State token not in body casues error",
+			body: "someJavascriptCode();\nsomeOtherJavaScriptCode();",
+			stateToken: "",
+			err: errors.New("cannot find state token"),
+		},
+		{
+			title: "State token with hypen handled correctly",
+			body: "someJavascriptCode();\nvar stateToken = '12345\x2D6789';\nsomeOtherJavaScriptCode();",
+			stateToken: "12345-6789",
+			err: nil,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.title, func(t *testing.T) {
+			stateToken, err := getStateTokenFromOktaPageBody(test.body)
+			assert.Equal(t, test.stateToken, stateToken)
+			if test.err != nil {
+				assert.NotNil(t, err)
+				assert.Equal(t, test.err.Error(), err.Error())
+			}
+			
+		})
+	}
+}
\ No newline at end of file