refactors

frohoff · Apr 25, 2019 · 7adb3cf · 7adb3cf
1 parent 4543487
commit 7adb3cf
Showing 1 changed file with 50 additions and 30 deletions.
diff --git a/src/test/java/ysoserial/test/payloads/PayloadsTest.java b/src/test/java/ysoserial/test/payloads/PayloadsTest.java
@@ -3,6 +3,7 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.OutputStream;
 import java.io.PrintStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -17,7 +18,6 @@
 import org.jboss.shrinkwrap.resolver.api.maven.Maven;
 import org.junit.Assume;
 import org.junit.Test;
-import org.junit.internal.TextListener;
 import org.junit.runner.Description;
 import org.junit.runner.JUnitCore;
 import org.junit.runner.Result;
@@ -253,70 +253,90 @@ static Object deserializeWithDependencies ( byte[] serialized, final String[] de
     }
 
     public static void main(String[] args) {
+
         JUnitCore junit = new JUnitCore();
         PayloadListener listener = new PayloadListener();
         junit.addListener(listener);
         Result result = junit.run(PayloadsTest.class);
         System.exit(result.wasSuccessful() ? 0 : 1);
     }
 
+    public static class StdIo {
+
+        private static final PrintStream realOut = System.out;
+        private static final PrintStream realErr = System.err;
+
+        public static void restoreStreams() {
+            setStreams(realOut, realErr);
+        }
+
+        public static void setStreams(PrintStream out, PrintStream err) {
+            System.setOut(out);
+            System.setErr(err);
+        }
+
+        public static void setStreams(OutputStream out, OutputStream err) {
+            setStreams(new PrintStream(out), new PrintStream(err));
+        }
+    }
+
     public static class PayloadListener extends RunListener {
-        private static PrintStream realOut = System.out;
-        private static PrintStream realErr = System.err;
+        public enum Status {
+            SUCCESS,
+            FAILURE,
+            IGNORE,
+            ASSUMPTION_FAILURE
+        }
 
         private Map<Description, ByteArrayOutputStream> outs = new HashMap<Description, ByteArrayOutputStream>();
-        //private Map<Description, ByteArrayOutputStream> errs = new HashMap<Description, ByteArrayOutputStream>();
+        private Map<Description, ByteArrayOutputStream> errs = new HashMap<Description, ByteArrayOutputStream>();
+
+        private Map<Description, Status> statuses = new HashMap<Description, Status>();
+
+        private Map<Description, Failure> failures = new HashMap<Description, Failure>();
 
         @Override
         public void testStarted(Description description) throws Exception {
+            System.out.println(getPayload(description.getDisplayName()) + ": STARTED");
+
+            statuses.put(description, Status.SUCCESS);
+
             ByteArrayOutputStream out = new ByteArrayOutputStream();
-            //ByteArrayOutputStream err = new ByteArrayOutputStream();
+//            ByteArrayOutputStream err = new ByteArrayOutputStream();
 
             outs.put(description, out);
-            //errs.put(description, err);
+//            errs.put(description, err);
 
-            setStreams(new PrintStream(out), new PrintStream(out));
-            //System.setErr(new PrintStream((err)));
+            StdIo.setStreams(out, out);
         }
 
         @Override
         public void testFinished(Description description) throws Exception {
             outs.get(description).close();
             //errs.get(description).close();
 
-            restoreStreams();
-        }
-
-        private void restoreStreams() {
-            setStreams(realOut, realErr);
-        }
+            StdIo.restoreStreams();
 
-        private void setStreams(PrintStream realOut, PrintStream realErr) {
-            System.setOut(realOut);
-            System.setErr(realErr);
+            Status status = statuses.get(description);
+            System.out.println(getPayload(description.getDisplayName()) + ": " + status);
+            if (status == Status.FAILURE) System.err.println(outs.get(description).toString());
         }
 
         @Override
         public void testFailure(Failure failure) throws Exception {
-            restoreStreams();
-            System.out.println(failure.getDescription().getDisplayName() + ": TEST-FAILED");
-            System.err.println(outs.get(failure.getDescription()).toString());
+            statuses.put(failure.getDescription(), Status.FAILURE);
+            failures.put(failure.getDescription(), failure);
         }
 
         @Override
         public void testAssumptionFailure(Failure failure) {
-            restoreStreams();
-            System.out.println(failure.getDescription().getDisplayName() + ": ASSUMPTION-FAILED");
+            statuses.put(failure.getDescription(), Status.ASSUMPTION_FAILURE);
+            failures.put(failure.getDescription(), failure);
         }
 
-        @Override
-        public void testIgnored(Description description) throws Exception {
-            restoreStreams();
-            System.out.println(description.getDisplayName() + ": TEST-IGNORED");
+        // testPayload[payloadClass: class ysoserial.payloads.JavassistWeld1](ysoserial.test.payloads.PayloadsTest)
+        public static String getPayload(String displayName) {
+            return displayName.replaceAll(".*\\[\\S+: class (\\w+\\.)+(\\w+)\\].*", "$2");
         }
-
-
     }
-
-
 }