From 7adb3cfe638228ae0bd2baf7584dc7d3f2831f17 Mon Sep 17 00:00:00 2001
From: Chris Frohoff <cfrohoff@qualcomm.com>
Date: Fri, 30 Nov 2018 17:18:18 -0800
Subject: [PATCH] refactors

---
 .../ysoserial/test/payloads/PayloadsTest.java | 80 ++++++++++++-------
 1 file changed, 50 insertions(+), 30 deletions(-)

diff --git a/src/test/java/ysoserial/test/payloads/PayloadsTest.java b/src/test/java/ysoserial/test/payloads/PayloadsTest.java
index b71c14a..53e6bcd 100644
--- a/src/test/java/ysoserial/test/payloads/PayloadsTest.java
+++ b/src/test/java/ysoserial/test/payloads/PayloadsTest.java
@@ -3,6 +3,7 @@ package ysoserial.test.payloads;
 
 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 java.util.concurrent.Callable;
 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,6 +253,7 @@ public class PayloadsTest {
     }
 
     public static void main(String[] args) {
+
         JUnitCore junit = new JUnitCore();
         PayloadListener listener = new PayloadListener();
         junit.addListener(listener);
@@ -260,23 +261,53 @@ public class PayloadsTest {
         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
@@ -284,39 +315,28 @@ public class PayloadsTest {
             outs.get(description).close();
             //errs.get(description).close();
 
-            restoreStreams();
-        }
+            StdIo.restoreStreams();
 
-        private void restoreStreams() {
-            setStreams(realOut, realErr);
-        }
-
-        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");
         }
-
-
     }
-
-
 }