From 16c162bba5759b279524f21caaeb051870ede969 Mon Sep 17 00:00:00 2001 From: thc202 Date: Thu, 27 Dec 2018 08:56:20 +0000 Subject: [PATCH] Allow to listen for script output Add a listener for script output, which provides the script and the output being written. Part of #5113 - Allow to differentiate scripts' output --- .../zap/extension/script/ExtensionScript.java | 96 +++++++++++++++++-- .../script/ScriptOutputListener.java | 40 ++++++++ 2 files changed, 130 insertions(+), 6 deletions(-) create mode 100644 src/org/zaproxy/zap/extension/script/ScriptOutputListener.java diff --git a/src/org/zaproxy/zap/extension/script/ExtensionScript.java b/src/org/zaproxy/zap/extension/script/ExtensionScript.java index 35d277bd200..dfc6f5e4618 100644 --- a/src/org/zaproxy/zap/extension/script/ExtensionScript.java +++ b/src/org/zaproxy/zap/extension/script/ExtensionScript.java @@ -40,6 +40,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.CopyOnWriteArrayList; import javax.script.Invocable; import javax.script.ScriptContext; @@ -152,6 +154,15 @@ public class ExtensionScript extends ExtensionAdaptor implements CommandLineList */ private List trackedDirs = Collections.synchronizedList(new ArrayList<>()); + /** + * The script output listeners added to the extension. + * + * @see #addScriptOutputListener(ScriptOutputListener) + * @see #getWriters(ScriptWrapper) + * @see #removeScriptOutputListener(ScriptOutputListener) + */ + private List outputListeners = new CopyOnWriteArrayList<>(); + public ExtensionScript() { super(NAME); this.setOrder(EXTENSION_ORDER); @@ -1191,17 +1202,16 @@ public List getTemplates(ScriptType type) { * ever script. It also supports script specific writers. */ private Writer getWriters(ScriptWrapper script) { + Writer delegatee = this.writers; Writer writer = script.getWriter(); - if (writer == null) { - // No script specific writer, just return the std one - return this.writers; - } else { - // Return the script specific writer in addition to the std one + if (writer != null) { + // Use the script specific writer in addition to the std one MultipleWriters scriptWriters = new MultipleWriters(); scriptWriters.addWriter(writer); scriptWriters.addWriter(writers); - return scriptWriters; + delegatee = scriptWriters; } + return new ScriptWriter(script, delegatee, outputListeners); } /** @@ -1609,14 +1619,53 @@ public void removeListener(ScriptEventListener listener) { this.listeners.remove(listener); } + /** + * Adds the given writer. + *

+ * It will be written to each time a script writes some output. + * + * @param writer the writer to add. + * @see #removeWriter(Writer) + * @see #addScriptOutputListener(ScriptOutputListener) + */ public void addWriter(Writer writer) { this.writers.addWriter(writer); } + /** + * Removes the given writer. + * + * @param writer the writer to remove. + * @see #addWriter(Writer) + */ public void removeWriter(Writer writer) { this.writers.removeWriter(writer); } + /** + * Adds the given script output listener. + * + * @param listener the listener to add. + * @since TODO add version + * @throws NullPointerException if the given listener is {@code null}. + * @see #removeScriptOutputListener(ScriptOutputListener) + */ + public void addScriptOutputListener(ScriptOutputListener listener) { + outputListeners.add(Objects.requireNonNull(listener, "The parameter listener must not be null.")); + } + + /** + * Removes the given script output listener. + * + * @param listener the listener to remove. + * @since TODO add version + * @throws NullPointerException if the given listener is {@code null}. + * @see #addScriptOutputListener(ScriptOutputListener) + */ + public void removeScriptOutputListener(ScriptOutputListener listener) { + outputListeners.remove(Objects.requireNonNull(listener, "The parameter listener must not be null.")); + } + public ScriptUI getScriptUI() { return scriptUI; } @@ -1852,4 +1901,39 @@ public void sessionModeChanged(Mode mode) { } } + + /** + * A {@code Writer} that notifies {@link ScriptOutputListener}s when writing. + */ + private static class ScriptWriter extends Writer { + + private final ScriptWrapper script; + private final Writer delegatee; + private final List outputListeners; + + public ScriptWriter(ScriptWrapper script, Writer delegatee, List outputListeners) { + this.script = Objects.requireNonNull(script); + this.delegatee = Objects.requireNonNull(delegatee); + this.outputListeners = Objects.requireNonNull(outputListeners); + } + + @Override + public void write(char[] cbuf, int off, int len) throws IOException { + delegatee.write(cbuf, off, len); + if (!outputListeners.isEmpty()) { + String output = new String(cbuf, off, len); + outputListeners.forEach(e -> e.output(script, output)); + } + } + + @Override + public void flush() throws IOException { + delegatee.flush(); + } + + @Override + public void close() throws IOException { + delegatee.close(); + } + } } diff --git a/src/org/zaproxy/zap/extension/script/ScriptOutputListener.java b/src/org/zaproxy/zap/extension/script/ScriptOutputListener.java new file mode 100644 index 00000000000..4b1a1a75aa2 --- /dev/null +++ b/src/org/zaproxy/zap/extension/script/ScriptOutputListener.java @@ -0,0 +1,40 @@ +/* + * Zed Attack Proxy (ZAP) and its related class files. + * + * ZAP is an HTTP/HTTPS proxy for assessing web application security. + * + * Copyright 2018 The ZAP Development Team + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.zaproxy.zap.extension.script; + +/** + * A listener for scripts' output. + * + * @since TODO add version + * @see ExtensionScript#addScriptOutputListener(ScriptOutputListener) + */ +@FunctionalInterface +public interface ScriptOutputListener { + + /** + * Called each time a script writes some output. + *

+ * Can be called concurrently. + * + * @param script the source of the output. + * @param output the new output. + */ + void output(ScriptWrapper script, String output); +}