authhelper: skip MS login proof

Skip the MS login proof when possible.

Signed-off-by: thc202 <thc202@gmail.com>

Author: thc202

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

@@ -29,6 +29,7 @@
 import org.openqa.selenium.TimeoutException;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
+import org.openqa.selenium.support.ui.ExpectedCondition;
 import org.openqa.selenium.support.ui.ExpectedConditions;
 import org.openqa.selenium.support.ui.WebDriverWait;
 import org.parosproxy.paros.Constant;
@@ -51,6 +52,7 @@ public final class MsLoginAuthenticator implements Authenticator {
     private static final By PASSWORD_FIELD = By.id("i0118");
     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 enum State {
         START,
@@ -61,6 +63,9 @@ private enum State {
 
         POST_PASSWORD,
         STAY_SIGNED_IN,
+
+        PROOF_REDIRECT,
+        PROOF,
     }
 
     @Override
@@ -177,7 +182,13 @@ private Result authenticateImpl(
                             Constant.messages.getString(
                                     "authhelper.auth.method.diags.steps.ms.stepchoice"));
 
-                    // XXX There might be a proof step too…
+                    try {
+                        waitForElement(wd, PROOF_REDIRECT_FIELD);
+                        states.add(State.PROOF_REDIRECT);
+                        break;
+                    } catch (TimeoutException e) {
+                        // Ignore, there's still the next step to check.
+                    }
 
                     try {
                         waitForElement(wd, KMSI_FIELD);
@@ -206,15 +217,55 @@ private Result authenticateImpl(
                     states.add(State.SUBMIT);
                     states.add(State.POST_PASSWORD);
                     break;
+
+                case PROOF_REDIRECT:
+                    WebElement proofElement = wd.findElement(PROOF_REDIRECT_FIELD);
+                    proofElement.click();
+                    diags.recordStep(
+                            wd,
+                            Constant.messages.getString(
+                                    "authhelper.auth.method.diags.steps.ms.clickproofredirect"),
+                            proofElement);
+
+                    states.add(State.PROOF);
+                    break;
+
+                case PROOF:
+                    try {
+                        WebElement skipElement =
+                                waitForElement(
+                                        wd,
+                                        new ElementWithText(By.tagName("button"), "Skip setup"));
+                        skipElement.click();
+                        diags.recordStep(
+                                wd,
+                                Constant.messages.getString(
+                                        "authhelper.auth.method.diags.steps.ms.clickproofskip"),
+                                skipElement);
+
+                        states.add(State.POST_PASSWORD);
+                        break;
+                    } catch (TimeoutException e) {
+                        diags.recordStep(
+                                wd,
+                                Constant.messages.getString(
+                                        "authhelper.auth.method.diags.steps.ms.stepproofunknown"));
+                        LOGGER.debug(
+                                "Still in proof but no skip setup button found, assuming unsuccessful login.");
+                        break;
+                    }
             }
         } while (!states.isEmpty());
 
         return new Result(true, successful, userField, pwdField);
     }
 
     private WebElement waitForElement(WebDriver wd, By by) {
-        return new WebDriverWait(wd, DEFAULT_WAIT_UNTIL)
-                .until(ExpectedConditions.elementToBeClickable(by));
+        return waitForElement(wd, ExpectedConditions.elementToBeClickable(by));
+    }
+
+    private WebElement waitForElement(WebDriver wd, ExpectedCondition<WebElement> condition) {
+        return new WebDriverWait(wd, DEFAULT_WAIT_UNTIL).until(condition);
     }
 
     private static boolean isMsLoginFlow(WebDriver wd) {
@@ -230,4 +281,28 @@ private static boolean isMsLoginFlow(WebDriver wd, Duration duration) {
             return false;
         }
     }
+
+    private static class ElementWithText implements ExpectedCondition<WebElement> {
+
+        private final By locator;
+        private final String text;
+
+        ElementWithText(By locator, String text) {
+            this.locator = locator;
+            this.text = text;
+        }
+
+        @Override
+        public WebElement apply(WebDriver driver) {
+            return driver.findElements(locator).stream()
+                    .filter(e -> text.equalsIgnoreCase(e.getText()))
+                    .findFirst()
+                    .orElse(null);
+        }
+
+        @Override
+        public String toString() {
+            return String.format("element '%s' with text '%s' is not present", locator, text);
+        }
+    }
 }

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

@@ -66,10 +66,13 @@ authhelper.auth.method.diags.steps.finish = Finished Steps
 authhelper.auth.method.diags.steps.loginlink = Login Link
 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.clickproofredirect = [MS] Click Proof Redirect
+authhelper.auth.method.diags.steps.ms.clickproofskip = [MS] Click Proof Skip
 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
 authhelper.auth.method.diags.steps.ms.stepchoice = [MS] Step Choice
+authhelper.auth.method.diags.steps.ms.stepproofunknown = [MS] Step Proof Unknown
 authhelper.auth.method.diags.steps.ms.stepunknown = [MS] Step Unknown
 authhelper.auth.method.diags.steps.password = Auto Fill Password
 authhelper.auth.method.diags.steps.refresh = Auto Refresh