authhelper: skip logout statements in CSA

Disable the statements after the logout comment in the CSA scripts.

Signed-off-by: thc202 <[email protected]>

Author: thc202

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

@@ -91,6 +91,9 @@
 import org.zaproxy.zap.users.User;
 import org.zaproxy.zap.utils.Pair;
 import org.zaproxy.zap.utils.Stats;
+import org.zaproxy.zest.core.v1.ZestComment;
+import org.zaproxy.zest.core.v1.ZestScript;
+import org.zaproxy.zest.core.v1.ZestStatement;
 
 public class AuthUtils {
 
@@ -146,6 +149,8 @@ public class AuthUtils {
 
     private static final String INPUT_TAG = "input";
 
+    private static final String RECORDING_LOGOUT = "ZAP Recording LOGOUT";
+
     private static final HttpRequestConfig REDIRECT_NOTIFIER_CONFIG =
             HttpRequestConfig.builder()
                     .setRedirectionValidator(
@@ -1362,4 +1367,16 @@ public static boolean isRelevantToAuthDiags(HttpMessage msg) {
                 || host.contains("mozilla")
                 || host.contains("safebrowsing-cache"));
     }
+
+    public static void disableLogoutStatements(ZestScript zestScript) {
+        boolean disable = false;
+        for (ZestStatement stmt : zestScript.getStatements()) {
+            if (disable) {
+                stmt.setEnabled(false);
+            } else if (stmt instanceof ZestComment comment
+                    && RECORDING_LOGOUT.equals(comment.getComment())) {
+                disable = true;
+            }
+        }
+    }
 }

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

@@ -149,7 +149,7 @@ public void insertDiagnostics(ZestScript zestScript) {
 
         for (int i = 0; i < zestScript.getStatements().size(); i++) {
             ZestStatement stmt = zestScript.getStatements().get(i);
-            if (stmt instanceof ZestClientElementClear) {
+            if (!stmt.isEnabled() || stmt instanceof ZestClientElementClear) {
                 continue;
             }
 

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

@@ -414,6 +414,7 @@ public WebSession authenticate(
                         zestScript.add(
                                 new ZestActionSleep(TimeUnit.SECONDS.toMillis(getLoginPageWait())));
                         removeCloseStatements(zestScript);
+                        AuthUtils.disableLogoutStatements(zestScript);
                     } else {
                         LOGGER.warn("Expected authScript to be a Zest script");
                         return null;

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

@@ -81,6 +81,7 @@ public void browserLaunched(SeleniumScriptUtils ssUtils) {
             paramsValues.put(USERNAME, credentials.getParam(USERNAME));
             paramsValues.put(PASSWORD, credentials.getParam(PASSWORD));
             ZestScript zs = csaMethod.getZestScript();
+            AuthUtils.disableLogoutStatements(zs);
             runner.setup(user, zs);
             runner.run(zs, paramsValues);
 

addOns/authhelper/src/test/java/org/zaproxy/addon/authhelper/AuthUtilsUnitTest.java

@@ -93,6 +93,9 @@
 import org.zaproxy.zap.testutils.TestUtils;
 import org.zaproxy.zap.users.User;
 import org.zaproxy.zap.utils.Pair;
+import org.zaproxy.zest.core.v1.ZestActionPrint;
+import org.zaproxy.zest.core.v1.ZestComment;
+import org.zaproxy.zest.core.v1.ZestScript;
 
 class AuthUtilsUnitTest extends TestUtils {
 
@@ -983,6 +986,43 @@ void shouldReportRelevantResponseHeaderTypeToAuthDiags(String type, String resul
         assertThat(res, is(equalTo(Boolean.parseBoolean(result))));
     }
 
+    @Test
+    void shouldDisableLogoutStatements() {
+        // Given
+        ZestScript zs = new ZestScript();
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestComment("ZAP Recording LOGOUT"));
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestActionPrint());
+
+        // When
+        AuthUtils.disableLogoutStatements(zs);
+
+        // Then
+        assertThat(zs.getStatements().get(0).isEnabled(), is(equalTo(true)));
+        assertThat(zs.getStatements().get(1).isEnabled(), is(equalTo(true)));
+        assertThat(zs.getStatements().get(2).isEnabled(), is(equalTo(true)));
+        assertThat(zs.getStatements().get(3).isEnabled(), is(equalTo(false)));
+        assertThat(zs.getStatements().get(4).isEnabled(), is(equalTo(false)));
+    }
+
+    @Test
+    void shouldNotDisableStatementsWhenNoLogoutCommentPresent() {
+        // Given
+        ZestScript zs = new ZestScript();
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestActionPrint());
+        zs.add(new ZestActionPrint());
+
+        // When
+        AuthUtils.disableLogoutStatements(zs);
+
+        // Then
+        zs.getStatements().forEach(e -> assertThat(e.isEnabled(), is(equalTo(true))));
+    }
+
     static class BrowserTest extends TestUtils {
 
         private static final String HTML_SHADOM_DOM =