package org.zaproxy.zap.userauth.session;

import java.net.HttpCookie;
import java.util.Iterator;
import java.util.List;

import org.apache.log4j.Logger;
import org.parosproxy.paros.Constant;
import org.parosproxy.paros.control.Control;
import org.parosproxy.paros.extension.ExtensionHook;
import org.parosproxy.paros.model.Model;
import org.parosproxy.paros.network.HttpMessage;
import org.zaproxy.zap.extension.httpsessions.ExtensionHttpSessions;
import org.zaproxy.zap.extension.httpsessions.HttpSession;
import org.zaproxy.zap.extension.httpsessions.HttpSessionTokensSet;
import org.zaproxy.zap.model.Context;

/**
 * The type corresponding to a {@link SessionManagementMethod} for web applications that use cookies
 * for session management.
 */
public class CookieBasedSessionManagementMethodType extends SessionManagementMethodType {

	/** The Constant METHOD_NAME. */
	private static final String METHOD_NAME = Constant.messages.getString("sessionmanagement.method.cb.name");

	/**
	 * The implementation for a {@link SessionManagementMethod} that for web applications that use
	 * cookies for session management.
	 */
	public static class CookieBasedSessionManagementMethod implements SessionManagementMethod {

		private static final Logger log = Logger.getLogger(CookieBasedSessionManagementMethod.class);

		private int contextId;

		private Context context;

		private static ExtensionHttpSessions extHttpSessions;

		public CookieBasedSessionManagementMethod(int contextId) {
			this.contextId = contextId;
		}

		@Override
		public String toString() {
			return CookieBasedSessionManagementMethodType.METHOD_NAME;
		}

		@Override
		public boolean isConfigured() {
			// Always configured
			return true;
		}

		@Override
		public WebSession extractWebSession(HttpMessage msg) {
			// TODO: Based on the already identified http session. This is based on the assumption
			// that the HttpSessions Extension has already identified the session
			return msg.getHttpSession();
			// return identifyMatchingWebSession(getSessionsForContext(), msg);
		}

		@Override
		public void processMessageToMatchSession(HttpMessage message, WebSession session)
				throws UnsupportedWebSessionException {

			if (session != null && message.getHttpSession() != session) {
				if (log.isDebugEnabled()) {
					log.debug("Modifying message to match session: " + session);
				}

				// make sure it's the right type
				if (!(session instanceof HttpSession)) {
					throw new UnsupportedWebSessionException(
							"The WebSession type provided is unsupported. Cookie based session management only supports "
									+ HttpSession.class + " type of WebSession.");
				}

				CookieBasedSessionManagementHelper.processMessageToMatchSession(message,
						(HttpSession) session);
			}
		}

		@Override
		public WebSession identifyMatchingWebSession(List<WebSession> sessions, HttpMessage msg)
				throws UnsupportedWebSessionException {
			// TODO: Proper implementation. Now it's hacked and does not work in all scenarios
			return CookieBasedSessionManagementHelper.getMatchingHttpSession(getSessionsForContext(), msg
					.getRequestHeader().getHttpCookies(), getHttpSessionsExtension()
					.getHttpSessionTokensSetForContext(getContext()));
		}

		private ExtensionHttpSessions getHttpSessionsExtension() {
			if (extHttpSessions == null) {
				extHttpSessions = (ExtensionHttpSessions) Control.getSingleton().getExtensionLoader()
						.getExtension(ExtensionHttpSessions.class);
				if (extHttpSessions == null)
					log.error("An error occured while loading the ExtensionHttpSessions.");
			}
			return extHttpSessions;
		}

		private Context getContext() {
			if (context == null) {
				context = Model.getSingleton().getSession().getContext(contextId);
			}
			return context;
		}

		private List<HttpSession> getSessionsForContext() {
			return getHttpSessionsExtension().getHttpSessionsForContext(getContext());
		}

		@Override
		public SessionManagementMethod duplicate() {
			return new CookieBasedSessionManagementMethod(contextId);
		}

		@Override
		public void clearWebSessionIdentifiers(HttpMessage msg) {
			HttpSessionTokensSet tokens = getHttpSessionsExtension().getHttpSessionTokensSetForContext(
					getContext());
			if (tokens == null) {
				log.info("No tokens to clear.");
				return;
			}
			List<HttpCookie> requestCookies = msg.getRequestHeader().getHttpCookies();
			Iterator<HttpCookie> it = requestCookies.iterator();
			while (it.hasNext())
				if (tokens.isSessionToken(it.next().getName()))
					it.remove();
			msg.setCookies(requestCookies);
		}
	}

	@Override
	public CookieBasedSessionManagementMethod createSessionManagementMethod(int contextId) {
		return new CookieBasedSessionManagementMethod(contextId);
	}

	@Override
	public String getName() {
		return METHOD_NAME;
	}

	@Override
	public AbstractSessionManagementMethodOptionsPanel buildOptionsPanel(Context uiSharedContext) {
		// No need for a configuration panel yet
		return null;
	}

	@Override
	public boolean hasOptionsPanel() {
		return false;
	}

	@Override
	public boolean isTypeForMethod(SessionManagementMethod method) {
		return (method instanceof CookieBasedSessionManagementMethod);
	}

	@Override
	public void hook(ExtensionHook extensionHook) {
		// No need for hooking anything
	}
}
@@ -36,16 +36,52 @@
 ssion;%0A%0A
+import java.lang.ref.WeakReference;%0A
 import j
@@ -1418,16 +1418,30 @@
  static 
+WeakReference%3C
 Extensio
@@ -1453,16 +1453,17 @@
 Sessions
+%3E
  extHttp
@@ -3386,34 +3386,108 @@
 null
-) %7B%0A%09%09%09%09extHttpSessions = 
+ %7C%7C extHttpSessions.get() == null) %7B%0A%09%09%09%09extHttpSessions = new WeakReference%3CExtensionHttpSessions%3E(
 (Ext
@@ -3513,16 +3513,23 @@
  Control
+%0A%09%09%09%09%09%09
 .getSing
@@ -3556,23 +3556,16 @@
 Loader()
-%0A%09%09%09%09%09%09
 .getExte
@@ -3598,16 +3598,17 @@
 s.class)
+)
 ;%0A%09%09%09%09if
@@ -3708,32 +3708,32 @@
 ssions.%22);%0A%09%09%09%7D%0A
-
 %09%09%09return extHtt
@@ -3741,16 +3741,22 @@
 Sessions
+.get()
 ;%0A%09%09%7D%0A%0A%09

Store WeakReference to HttpSessionsExtension.