From 3f7f0960165382138fab372a108cf82edba2e0ff Mon Sep 17 00:00:00 2001
From: Matt Robinson <21823930+Hobbit44@users.noreply.github.com>
Date: Wed, 11 Dec 2024 01:42:25 +0000
Subject: [PATCH] Added support for MFA token from env for Onelogin TOTP

---
 cmd/saml2aws/main.go              |  2 +-
 pkg/provider/onelogin/onelogin.go | 11 ++++++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/cmd/saml2aws/main.go b/cmd/saml2aws/main.go
index fe6a96372..b25b61fb3 100644
--- a/cmd/saml2aws/main.go
+++ b/cmd/saml2aws/main.go
@@ -81,7 +81,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, GoogleApps). (env: SAML2AWS_MFA_TOKEN)").Envar("SAML2AWS_MFA_TOKEN").StringVar(&commonFlags.MFAToken)
+	app.Flag("mfa-token", "The current MFA token (supported in Keycloak, ADFS, GoogleApps, Onelogin TOTP). (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("policyfile", "The file containing the supplemental AssumeRole policy. (env: SAML2AWS_POLICY_FILE)").Envar("SAML2AWS_POLICY_FILE").StringVar(&commonFlags.PolicyFile)
 	app.Flag("policyarns", "The ARN of supplemental policies to restrict the token. (env: SAML2AWS_POLICY_ARNS)").Envar("SAML2AWS_POLICY_ARNS").StringVar(&commonFlags.PolicyARNs)
diff --git a/pkg/provider/onelogin/onelogin.go b/pkg/provider/onelogin/onelogin.go
index 951e4db21..48d76ac2b 100644
--- a/pkg/provider/onelogin/onelogin.go
+++ b/pkg/provider/onelogin/onelogin.go
@@ -159,7 +159,7 @@ func (c *Client) Authenticate(loginDetails *creds.LoginDetails) (string, error)
 		samlAssertion = authData.String()
 	case MessageMFARequired:
 		logger.Debug("Verifying MFA")
-		samlAssertion, err = verifyMFA(c, oauthToken, c.AppID, host, resp)
+		samlAssertion, err = verifyMFA(c, loginDetails.MFAToken, oauthToken, c.AppID, host, resp)
 		if err != nil {
 			return "", errors.Wrap(err, "error verifying MFA")
 		}
@@ -206,7 +206,7 @@ func addContentHeaders(r *http.Request) {
 
 // verifyMFA is used to either prompt to user for one time password or request approval using push notification.
 // For more details check https://developers.onelogin.com/api-docs/2/saml-assertions/verify-factor
-func verifyMFA(oc *Client, oauthToken, appID, host, resp string) (string, error) {
+func verifyMFA(oc *Client, mfaToken, oauthToken, appID, host, resp string) (string, error) {
 	stateToken := gjson.Get(resp, "state_token").String()
 	// choose an mfa option if there are multiple enabled
 	var option int
@@ -287,7 +287,12 @@ func verifyMFA(oc *Client, oauthToken, appID, host, resp string) (string, error)
 
 	switch mfaIdentifer {
 	case IdentifierSmsMfa, IdentifierTotpMfa, IdentifierYubiKey, IdentifierDuoSecurity:
-		verifyCode := prompter.StringRequired("Enter verification code")
+		var verifyCode string
+		if mfaIdentifer == IdentifierTotpMfa && mfaToken != "" {
+			verifyCode = mfaToken
+		} else {
+			verifyCode = prompter.StringRequired("Enter verification code")
+		}
 		var verifyBody bytes.Buffer
 		err := json.NewEncoder(&verifyBody).Encode(VerifyRequest{AppID: appID, DeviceID: mfaDeviceID, StateToken: stateToken, OTPToken: verifyCode})
 		if err != nil {