/*
 *  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.groovy;

import java.util.ArrayList;
import java.util.List;

import javax.script.ScriptEngine;
import javax.script.ScriptEngineFactory;
import javax.swing.ImageIcon;

import org.codehaus.groovy.jsr223.GroovyScriptEngineFactory;
import org.fife.ui.rsyntaxtextarea.SyntaxConstants;
import org.zaproxy.zap.control.ExtensionFactory;
import org.zaproxy.zap.extension.script.DefaultEngineWrapper;

public class GroovyEngineWrapper extends DefaultEngineWrapper {

	private final ScriptEngineFactory factory;

	private static ScriptEngine getScriptEngineFromFactory(ScriptEngineFactory factory){
		ClassLoader previousContextClassLoader = Thread.currentThread().getContextClassLoader();
		Thread.currentThread().setContextClassLoader(new AddOnClassLoaderWrapper(ExtensionFactory.getAddOnLoader()));
		try {
			return factory.getScriptEngine();
		} finally {
			Thread.currentThread().setContextClassLoader(previousContextClassLoader);
		}
	}

	public GroovyEngineWrapper() {
		this(new GroovyScriptEngineFactory());
	}

	private GroovyEngineWrapper(GroovyScriptEngineFactory factory) {
		super(factory.getScriptEngine());
		this.factory = factory;
	}

	@Override
	public ImageIcon getIcon() {
		return ExtensionGroovy.GROOVY_ICON;
	}

	@Override
	public String getSyntaxStyle() {
		return SyntaxConstants.SYNTAX_STYLE_GROOVY;
	}
	
	@Override
	public List<String> getExtensions() {
		List<String> list = new ArrayList<>();
		list.add("groovy");
		return list;
	}

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

	@Override
	public ScriptEngine getEngine() {
		return getScriptEngineFromFactory(this.factory);
	}
}

@@ -999,32 +999,90 @@
 tEngineFactory;%0A
+import org.codehaus.groovy.jsr223.GroovyScriptEngineImpl;%0A
 import org.fife.
@@ -1117,16 +1117,16 @@
 stants;%0A
-
 import o
@@ -1229,16 +1229,55 @@
 apper;%0A%0A
+import groovy.lang.GroovyClassLoader;%0A%0A
 public c
@@ -1377,17 +1377,16 @@
 actory;%0A
-%0A
 %09private
@@ -1390,430 +1390,44 @@
 ate 
-static ScriptEngine getScriptEngineFromFactory(ScriptEngineFactory factory)%7B%0A%09%09ClassLoader previousContextClassLoader = Thread.currentThread().getContextClassLoader();%0A%09%09Thread.currentThread().setContextClassLoader(new AddOnClassLoaderWrapper(ExtensionFactory.getAddOnLoader()));%0A%09%09try %7B%0A%09%09%09return factory.getScriptEngine();%0A%09%09%7D finally %7B%0A%09%09%09Thread.currentThread().setContextClassLoader(previousContextClassLoader);%0A%09%09%7D%0A%09%7D
+final GroovyClassLoader classLoader;
 %0A%0A%09p
@@ -1629,16 +1629,209 @@
 actory;%0A
+%09%09// Use AddOnLoader as parent class loader to allow access to (all) add-on classes.%0A%09%09this.classLoader = new GroovyClassLoader(new AddOnClassLoaderWrapper(ExtensionFactory.getAddOnLoader()));%0A
 %09%7D%0A%0A%09@Ov
@@ -2252,55 +2252,168 @@
 %7B%0A%09%09
-return getScriptEngineFromFactory(this.factory)
+GroovyScriptEngineImpl scriptEngine = (GroovyScriptEngineImpl) this.factory.getScriptEngine();%0A%09%09scriptEngine.setClassLoader(classLoader);%0A%09%09return scriptEngine
 ;%0A%09%7D

set the class loader to the engine