authhelper: Add Header based session management

Signed-off-by: Simon Bennetts <psiinon@gmail.com>
zaproxy · Feb 7, 2023 · 5618519 · 5618519
1 parent 10b9ba0
commit 5618519
Show file tree

Hide file tree

Showing 12 changed files with 1,264 additions and 29 deletions.
diff --git a/addOns/authhelper/CHANGELOG.md b/addOns/authhelper/CHANGELOG.md
@@ -4,6 +4,8 @@ All notable changes to this add-on will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## Unreleased
+### Added
+- Support for header based session management.
 
 ### Fixed
 - Code link in help.

diff --git a/addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/ExtensionAuthhelper.java b/addOns/authhelper/src/main/java/org/zaproxy/addon/authhelper/ExtensionAuthhelper.java
@@ -30,6 +30,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.parosproxy.paros.Constant;
+import org.parosproxy.paros.control.Control;
 import org.parosproxy.paros.control.Control.Mode;
 import org.parosproxy.paros.extension.ExtensionAdaptor;
 import org.parosproxy.paros.extension.ExtensionHook;
@@ -41,6 +42,7 @@
 import org.zaproxy.zap.authentication.ManualAuthenticationMethodType;
 import org.zaproxy.zap.authentication.PostBasedAuthenticationMethodType;
 import org.zaproxy.zap.authentication.PostBasedAuthenticationMethodType.PostBasedAuthenticationMethod;
+import org.zaproxy.zap.extension.sessions.ExtensionSessionManagement;
 import org.zaproxy.zap.model.Context;
 import org.zaproxy.zap.utils.Stats;
 import org.zaproxy.zap.utils.ZapXmlConfiguration;
@@ -51,11 +53,36 @@ public class ExtensionAuthhelper extends ExtensionAdaptor {
 
     private static final Logger LOGGER = LogManager.getLogger(ExtensionAuthhelper.class);
 
+    private static final HeaderBasedSessionManagementMethodType HEADER_BASED_TYPE =
+            new HeaderBasedSessionManagementMethodType();
+
     public ExtensionAuthhelper() {
         super();
         this.setI18nPrefix("authhelper");
     }
 
+    @Override
+    public void optionsLoaded() {
+        ExtensionSessionManagement extSm = getExtensionSessionManagement();
+        if (extSm != null) {
+            extSm.getSessionManagementMethodTypes().add(HEADER_BASED_TYPE);
+        }
+    }
+
+    @Override
+    public void unload() {
+        ExtensionSessionManagement extSm = getExtensionSessionManagement();
+        if (extSm != null) {
+            extSm.getSessionManagementMethodTypes().remove(HEADER_BASED_TYPE);
+        }
+    }
+
+    private static ExtensionSessionManagement getExtensionSessionManagement() {
+        return Control.getSingleton()
+                .getExtensionLoader()
+                .getExtension(ExtensionSessionManagement.class);
+    }
+
     @Override
     public void hook(ExtensionHook extensionHook) {
         extensionHook.addSessionListener(new AuthSessionChangedListener());