From af03f541377a0d1cf730a0475c8b086adf39eb07 Mon Sep 17 00:00:00 2001
From: thc202 <thc202@gmail.com>
Date: Wed, 14 Nov 2018 23:37:11 +0000
Subject: [PATCH] zest: rely on script context writer

Change ZestZapRunner to not set the script console writer as output
writer if the script context writer is already set, while both end up in
the script console the latter allows the core to intercept the writes.
For the cases the context writer is not set the writes are not
intercepted.
Update changes in ZapAddOn.xml file.

Related to zaproxy/zaproxy#5113 - Allow to differentiate scripts' output
---
 .../zaproxy/zap/extension/zest/ZapAddOn.xml   |  1 +
 .../zap/extension/zest/ZestZapRunner.java     | 23 +++++++++++++++----
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/org/zaproxy/zap/extension/zest/ZapAddOn.xml b/src/org/zaproxy/zap/extension/zest/ZapAddOn.xml
index e766f655454..30d5c47e755 100644
--- a/src/org/zaproxy/zap/extension/zest/ZapAddOn.xml
+++ b/src/org/zaproxy/zap/extension/zest/ZapAddOn.xml
@@ -7,6 +7,7 @@
 	<url>https://github.com/zaproxy/zap-core-help/wiki/HelpAddonsZestZest</url>
 	<changes>
 	<![CDATA[
+	Rely on script context writer for script output.<br>
 	]]>
 	</changes>
 	<dependencies>
diff --git a/src/org/zaproxy/zap/extension/zest/ZestZapRunner.java b/src/org/zaproxy/zap/extension/zest/ZestZapRunner.java
index 0843fb3f56c..d87429ee8ac 100644
--- a/src/org/zaproxy/zap/extension/zest/ZestZapRunner.java
+++ b/src/org/zaproxy/zap/extension/zest/ZestZapRunner.java
@@ -21,6 +21,7 @@
 package org.zaproxy.zap.extension.zest;
 
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -69,6 +70,8 @@ public class ZestZapRunner extends ZestBasicRunner implements ScannerListener {
 
 	private static final int ZEST_HISTORY_REFERENCE_TYPE = HistoryReference.TYPE_ZEST_SCRIPT;
 	private static final int FAIL_ACTION_PLUGIN_ID = 50004;
+
+	private static Field fieldOutputWriter;
 	
 	private ExtensionZest extension;
 	private ZestScriptWrapper wrapper = null;
@@ -126,8 +129,8 @@ public String run(ZestScript script, Map<String, String> params) throws ZestAsse
 			this.target = null;
 			if (wrapper.getWriter() != null) {
 				super.setOutputWriter(wrapper.getWriter());
-			} else if (extension.getExtScript().getScriptUI() != null) {
-				super.setOutputWriter(extension.getExtScript().getScriptUI().getOutputWriter());
+			} else if (scriptUI != null && !hasOutputWriter()) {
+				super.setOutputWriter(scriptUI.getOutputWriter());
 			}
 			this.setDebug(this.wrapper.isDebug());
 			
@@ -135,6 +138,18 @@ public String run(ZestScript script, Map<String, String> params) throws ZestAsse
     	}
 	}
 
+	private boolean hasOutputWriter() {
+		try {
+			if (fieldOutputWriter == null) {
+				fieldOutputWriter = ZestBasicRunner.class.getDeclaredField("outputWriter");
+				fieldOutputWriter.setAccessible(true);
+			}
+			return fieldOutputWriter.get(this) != null;
+		} catch (IllegalAccessException | NoSuchFieldException e) {
+			return false;
+		}
+	}
+
 	@Override
 	public String run (ZestScript script, ZestRequest target, Map<String, String> params) 
 			throws ZestAssertFailException, ZestActionFailException, IOException,
@@ -142,8 +157,8 @@ public String run (ZestScript script, ZestRequest target, Map<String, String> pa
     	log.debug("Run script " + script.getTitle());
 		if (wrapper.getWriter() != null) {
 			super.setOutputWriter(wrapper.getWriter());
-		} else if (extension.getExtScript().getScriptUI() != null) {
-			super.setOutputWriter(extension.getExtScript().getScriptUI().getOutputWriter());
+		} else if (scriptUI != null && !hasOutputWriter()) {
+			super.setOutputWriter(scriptUI.getOutputWriter());
 		}
 		this.setDebug(this.wrapper.isDebug());
 		String result = super.run(script, target, params);