Merge pull request #6931 from thc202/authhelper/ms-login-totp

authhelper: support TOTP in MS login

Author: psiinon

addOns/authhelper/CHANGELOG.md

@@ -5,13 +5,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## Unreleased
 ### Added
-- Handle account selection in Microsoft login.
+- Handle account selection and TOTP step in Microsoft login.
 
 ### Changed
 - Fail the Microsoft login if not able to perform all the expected steps.
 - Track GWT headers.
 - Handle additional exceptions when processing JSON authentication components.
 
+### Fixed
+- Do not include known authentication providers in context.
+
 ## [0.32.0] - 2025-11-07
 ### Changed
 - Track authentication headers with key in the name.

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/AuthUtils.java

@@ -428,6 +428,15 @@ public static <T> T ignoreSeleniumExceptions(Supplier<T> supplier) {
         return null;
     }
 
+    public static boolean isAuthProvider(HttpMessage msg) {
+        for (Authenticator authenticator : AUTHENTICATORS) {
+            if (authenticator.isOwnSite(msg)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Authenticate as the given user, by filling in and submitting the login form
      *

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/internal/ClientSideHandler.java

@@ -101,7 +101,8 @@ public void handleMessage(HttpMessageHandlerContext ctx, HttpMessage msg) {
                     && containsMaybeEncodedString(reqBody, authCreds.getUsername())
                     && containsMaybeEncodedString(reqBody, authCreds.getPassword())
                     && AuthUtils.getSessionManagementDetailsForContext(user.getContext().getId())
-                            != null) {
+                            != null
+                    && !AuthUtils.isAuthProvider(msg)) {
                 // The app is sending user creds to another site. Assume this is part of the valid
                 // auth flow and add to the context
                 try {

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/internal/auth/Authenticator.java

@@ -21,6 +21,7 @@
 
 import java.util.List;
 import org.openqa.selenium.WebDriver;
+import org.parosproxy.paros.network.HttpMessage;
 import org.zaproxy.addon.authhelper.AuthenticationDiagnostics;
 import org.zaproxy.addon.authhelper.internal.AuthenticationStep;
 import org.zaproxy.zap.authentication.UsernamePasswordAuthenticationCredentials;
@@ -33,6 +34,8 @@ public interface Authenticator {
     public static record Result(
             boolean isAttempted, boolean isSuccessful, boolean hasUserField, boolean hasPwdField) {}
 
+    boolean isOwnSite(HttpMessage msg);
+
     Result authenticate(
             AuthenticationDiagnostics diags,
             WebDriver wd,

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/internal/auth/DefaultAuthenticator.java

@@ -26,6 +26,7 @@
 import org.apache.logging.log4j.Logger;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
+import org.parosproxy.paros.network.HttpMessage;
 import org.zaproxy.addon.authhelper.AuthUtils;
 import org.zaproxy.addon.authhelper.AuthenticationDiagnostics;
 import org.zaproxy.addon.authhelper.internal.AuthenticationStep;
@@ -36,6 +37,12 @@ public class DefaultAuthenticator implements Authenticator {
 
     private static final Logger LOGGER = LogManager.getLogger(DefaultAuthenticator.class);
 
+    @Override
+    public boolean isOwnSite(HttpMessage msg) {
+        // Default does not own any site.
+        return false;
+    }
+
     @Override
     public Result authenticate(
             AuthenticationDiagnostics diags,

addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/internal/auth/MsLoginAuthenticator.java

@@ -23,6 +23,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.openqa.selenium.By;
@@ -33,16 +34,20 @@
 import org.openqa.selenium.support.ui.ExpectedConditions;
 import org.openqa.selenium.support.ui.WebDriverWait;
 import org.parosproxy.paros.Constant;
+import org.parosproxy.paros.network.HttpMessage;
 import org.zaproxy.addon.authhelper.AuthUtils;
 import org.zaproxy.addon.authhelper.AuthenticationDiagnostics;
 import org.zaproxy.addon.authhelper.internal.AuthenticationStep;
+import org.zaproxy.addon.commonlib.internal.TotpSupport;
 import org.zaproxy.zap.authentication.UsernamePasswordAuthenticationCredentials;
 import org.zaproxy.zap.model.Context;
 
 public final class MsLoginAuthenticator implements Authenticator {
 
     private static final String PARTIAL_LOGIN_URL = "login.microsoftonline";
 
+    private static final String LOGIN_URL = "https://" + PARTIAL_LOGIN_URL;
+
     private static final Logger LOGGER = LogManager.getLogger(MsLoginAuthenticator.class);
 
     private static final Duration PAGE_LOAD_WAIT_UNTIL = Duration.ofSeconds(5);
@@ -53,6 +58,8 @@ public final class MsLoginAuthenticator implements Authenticator {
     private static final By SUBMIT_BUTTON = By.id("idSIButton9");
     private static final By KMSI_FIELD = By.id("KmsiCheckboxField");
     private static final By PROOF_REDIRECT_FIELD = By.id("idSubmit_ProofUp_Redirect");
+    private static final By PROOF_TOTP_FIELD = By.id("idTxtBx_SAOTCC_OTC");
+    private static final By PROOF_TOTP_VERIFY_FIELD = By.id("idSubmit_SAOTCC_Continue");
     private static final By PROOF_DONE_FIELD = By.id("id__5");
 
     private enum State {
@@ -68,9 +75,15 @@ private enum State {
         STAY_SIGNED_IN,
 
         PROOF_REDIRECT,
+        PROOF_TOTP,
         PROOF,
     }
 
+    @Override
+    public boolean isOwnSite(HttpMessage msg) {
+        return msg.getRequestHeader().getURI().toString().startsWith(LOGIN_URL);
+    }
+
     @Override
     public Result authenticate(
             AuthenticationDiagnostics diags,
@@ -101,6 +114,7 @@ private Result authenticateImpl(
         boolean successful = false;
         boolean userField = false;
         boolean pwdField = false;
+        int totpRetries = 0;
 
         do {
             switch (states.remove()) {
@@ -117,13 +131,16 @@ private Result authenticateImpl(
 
                     if (findElement(wd, By.tagName("div"), "Pick an account") != null) {
                         states.add(State.ACCOUNT_SELECTION);
+                    } else if (findElement(wd, PROOF_TOTP_FIELD) != null) {
+                        userField = true;
+                        pwdField = true;
+                        states.add(State.PROOF_TOTP);
                     } else {
                         diags.recordStep(
                                 wd,
                                 Constant.messages.getString(
                                         "authhelper.auth.method.diags.steps.ms.missingusername"));
-                        LOGGER.debug(
-                                "Expected username field not found nor pick an account, failing login.");
+                        LOGGER.debug("Unexpected initial state, failing login.");
                     }
 
                     break;
@@ -246,6 +263,21 @@ private Result authenticateImpl(
                         // Ignore, there's still the next step to check.
                     }
 
+                    try {
+                        WebElement proofTotpElement =
+                                waitForElement(
+                                        wd,
+                                        new ElemenContainsText(
+                                                By.id("idDiv_SAOTCC_Description"),
+                                                "authenticator"));
+                        if (proofTotpElement != null) {
+                            states.add(State.PROOF_TOTP);
+                            break;
+                        }
+                    } catch (TimeoutException e) {
+                        // Ignore, there's still the next step to check.
+                    }
+
                     try {
                         waitForElement(wd, KMSI_FIELD);
                         states.add(State.STAY_SIGNED_IN);
@@ -286,6 +318,38 @@ private Result authenticateImpl(
                     states.add(State.PROOF);
                     break;
 
+                case PROOF_TOTP:
+                    totpRetries++;
+
+                    if (totpRetries > 2) {
+                        diags.recordStep(
+                                wd,
+                                Constant.messages.getString(
+                                        "authhelper.auth.method.diags.steps.ms.clickprooftotperror"));
+                        LOGGER.debug("TOTP proof failed, assuming unsuccessful login.");
+                        break;
+                    }
+
+                    WebElement proofTotpElement = wd.findElement(PROOF_TOTP_FIELD);
+                    WebElement proofTotpVerifyElement = wd.findElement(PROOF_TOTP_VERIFY_FIELD);
+                    proofTotpElement.clear();
+                    proofTotpElement.sendKeys(TotpSupport.getCode(credentials));
+                    proofTotpVerifyElement.click();
+                    diags.recordStep(
+                            wd,
+                            Constant.messages.getString(
+                                    "authhelper.auth.method.diags.steps.ms.clickprooftotpverify"),
+                            proofTotpVerifyElement);
+
+                    if (findElementContains(wd, By.id("idDiv_SAOTCC_ErrorMsg_OTC"), "try again")
+                            != null) {
+                        states.add(State.PROOF_TOTP);
+                    } else {
+                        states.add(State.POST_PASSWORD);
+                    }
+
+                    break;
+
                 case PROOF:
                     try {
                         waitForElement(wd, new ElementWithText(By.tagName("button"), "Skip setup"));
@@ -320,6 +384,17 @@ private static boolean isUserLoggedIn(WebDriver wd, WebElement userElement) {
                 .anyMatch(e -> "Signed in".equalsIgnoreCase(e.getText()));
     }
 
+    private static WebElement findElement(WebDriver wd, By by) {
+        return wd.findElements(by).stream().findFirst().orElse(null);
+    }
+
+    private static WebElement findElementContains(WebDriver wd, By by, String text) {
+        return wd.findElements(by).stream()
+                .filter(e -> StringUtils.containsIgnoreCase(e.getText(), text))
+                .findFirst()
+                .orElse(null);
+    }
+
     private static WebElement findElement(WebDriver wd, By by, String text) {
         return wd.findElements(by).stream()
                 .filter(e -> text.equalsIgnoreCase(e.getText()))
@@ -372,4 +447,28 @@ public String toString() {
             return String.format("element '%s' with text '%s' is not present", locator, text);
         }
     }
+
+    private static class ElemenContainsText implements ExpectedCondition<WebElement> {
+
+        private final By locator;
+        private final String text;
+
+        ElemenContainsText(By locator, String text) {
+            this.locator = locator;
+            this.text = text;
+        }
+
+        @Override
+        public WebElement apply(WebDriver driver) {
+            return driver.findElements(locator).stream()
+                    .filter(e -> StringUtils.containsIgnoreCase(e.getText(), text))
+                    .findFirst()
+                    .orElse(null);
+        }
+
+        @Override
+        public String toString() {
+            return String.format("element '%s' containing text '%s' is not present", locator, text);
+        }
+    }
 }

addOns/authhelper/src/main/resources/org/zaproxy/addon/authhelper/resources/Messages.properties

@@ -71,6 +71,8 @@ authhelper.auth.method.diags.steps.ms.clickbutton = [MS] Click Button
 authhelper.auth.method.diags.steps.ms.clickkmsi = [MS] Click KMSI
 authhelper.auth.method.diags.steps.ms.clickproofdone = [MS] Click Proof Done
 authhelper.auth.method.diags.steps.ms.clickproofredirect = [MS] Click Proof Redirect
+authhelper.auth.method.diags.steps.ms.clickprooftotperror = [MS] TOTP Error
+authhelper.auth.method.diags.steps.ms.clickprooftotpverify = [MS] TOTP Verify
 authhelper.auth.method.diags.steps.ms.missingbutton = [MS] Missing Button
 authhelper.auth.method.diags.steps.ms.missingpassword = [MS] Missing Password Field
 authhelper.auth.method.diags.steps.ms.missingusername = [MS] Missing Username Field