Allow "encapsulated output" from the shell command

Since the ConsoleReader now controls the display of the prompt, there is no
longer any way for an external program interacting with the nomulus shell to
recognize when the last command issued has been completed.

To remedy this, we introduce an "--encapsulate_output" flag, which causes
standard output and standard error to be wrapped in a class that precedes
all normal lines of output with a prefix ("out: " or "err: ", accordingly)
and allows the command processor to insert a "SUCCESS" or "FAILURE"
line (with no special prefix) after completing the processing of a
command.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=196702338

java/google/registry/tools/ShellCommand.java

@@ -17,6 +17,7 @@ package google.registry.tools;
 import static com.google.common.base.StandardSystemProperty.USER_HOME;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.nio.charset.StandardCharsets.US_ASCII;
+import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.beust.jcommander.JCommander;
 import com.beust.jcommander.Parameter;
@@ -27,13 +28,18 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableTable;
+import com.google.common.escape.SourceCodeEscapers;
 import google.registry.util.Clock;
 import google.registry.util.SystemClock;
 import java.io.BufferedReader;
+import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.FilterOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.io.PrintStream;
 import java.io.StreamTokenizer;
 import java.io.StringReader;
 import java.util.Arrays;
@@ -64,6 +70,9 @@ public class ShellCommand implements Command {
   private static final String NON_ALERT_COLOR = "\u001b[32m"; // green foreground
   private static final String ALERT_COLOR = "\u001b[1;41;97m"; // red background
   private static final Duration IDLE_THRESHOLD = Duration.standardHours(1);
+  private static final String SUCCESS = "SUCCESS";
+  private static final String FAILURE = "FAILURE ";
+
 
   private final CommandRunner runner;
   private final BufferedReader lineReader;
@@ -77,6 +86,16 @@ public class ShellCommand implements Command {
               + "Will instead warn you and require re-running the command.")
   boolean dontExitOnIdle = false;
 
+  @Parameter(
+      names = {"--encapsulate_output"},
+      description =
+          "Encapsulate command standard output and error by combining the two streams to standard "
+              + "output and inserting a prefix ('out:' or 'err:') at the beginning of every line "
+              + "of normal output and adding a line consisting of either 'SUCCESS' or "
+              + "'FAILURE <exception-name> <error-message>' at the end of the output for a "
+              + "command, allowing the output to be easily parsed by wrapper scripts.")
+  boolean encapsulateOutput = false;
+
   public ShellCommand(CommandRunner runner) throws IOException {
     this.runner = runner;
     InputStream in = System.in;
@@ -136,6 +155,23 @@ public class ShellCommand implements Command {
     String line;
     DateTime lastTime = clock.nowUtc();
     while ((line = getLine()) != null) {
+      PrintStream orgStdout = null;
+      PrintStream orgStderr = null;
+      EncapsulatingOutputStream encapsulatedOutputStream = null;
+      EncapsulatingOutputStream encapsulatedErrorStream = null;
+
+
+      // Wrap standard output and error if requested. We have to do so here in run because the flags
+      // haven't been processed in the constructor.
+      if (encapsulateOutput) {
+        orgStdout = System.out;
+        orgStderr = System.err;
+        encapsulatedOutputStream = new EncapsulatingOutputStream(System.out, "out: ");
+        encapsulatedErrorStream = new EncapsulatingOutputStream(System.out, "err: ");
+        System.setOut(new PrintStream(encapsulatedOutputStream));
+        System.setErr(new PrintStream(encapsulatedErrorStream));
+      }
+
       // Make sure we're not idle for too long. Only relevant when we're "extra careful"
       if (!dontExitOnIdle
           && beExtraCareful
@@ -149,13 +185,32 @@ public class ShellCommand implements Command {
       if (lineArgs.length == 0) {
         continue;
       }
+      Exception lastError = null;
       try {
         runner.run(lineArgs);
       } catch (Exception e) {
+        lastError = e;
         System.err.println("Got an exception:\n" + e);
       }
+      try {
+        if (encapsulatedOutputStream != null) {
+          encapsulatedOutputStream.dumpLastLine();
+          encapsulatedErrorStream.dumpLastLine();
+          System.setOut(orgStdout);
+          System.setErr(orgStderr);
+          if (lastError == null) {
+            emitSuccess();
+          } else {
+            emitFailure(lastError);
+          }
+        }
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    if (!encapsulateOutput) {
+      System.err.println();
     }
-    System.err.println();
   }
 
   private String getLine() {
@@ -189,6 +244,26 @@ public class ShellCommand implements Command {
     return resultBuilder.build().toArray(new String[0]);
   }
 
+  /**
+   * Emit a success command separator.
+   *
+   * <p>Dumps the last line of output prior to doing this.
+   */
+  private void emitSuccess() throws IOException {
+    System.out.println(SUCCESS);
+    System.out.flush();
+  }
+
+  /**
+   * Emit a failure message obtained from the throwable.
+   *
+   * <p>Dumps the last line of output prior to doing this.
+   */
+  private void emitFailure(Throwable e) throws IOException {
+    System.out.println(FAILURE + e.getClass().getName() + " "
+                       + SourceCodeEscapers.javaCharEscaper().escape(e.getMessage()));
+  }
+
   @VisibleForTesting
   static class JCommanderCompletor implements Completor {
 
@@ -448,4 +523,51 @@ public class ShellCommand implements Command {
           .orElse(DEFAULT_PARAM_DOC);
     }
   }
+
+  /**
+   * Encapsulate output according to the protocol described in the documentation for the
+   * --encapsulate_output flag.
+   */
+  @VisibleForTesting
+  static class EncapsulatingOutputStream extends FilterOutputStream {
+
+    private final byte[] prefix;
+    private final ByteArrayOutputStream lastLine = new ByteArrayOutputStream();
+
+    // Flag to keep track of whether the last character written was a newline.  We initialize this
+    // to "true" because we always want the first line of output to be escaped with a leading space.
+    boolean lastWasNewline = true;
+
+    EncapsulatingOutputStream(OutputStream out, String identifier) {
+      super(out);
+      this.prefix = identifier.getBytes(UTF_8);
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+      lastLine.write(b);
+      if (b == '\n') {
+        out.write(prefix);
+        lastLine.writeTo(out);
+        out.flush();
+        lastLine.reset();
+      }
+    }
+
+    @Override
+    public void flush() throws IOException {
+      dumpLastLine();
+    }
+
+    /** Dump the accumulated last line of output, if there was one. */
+    public void dumpLastLine() throws IOException {
+      if (lastLine.size() > 0) {
+        out.write(prefix);
+        lastLine.writeTo(out);
+        out.write('\n');
+        out.flush();
+        lastLine.reset();
+      }
+    }
+  }
 }

javatests/google/registry/tools/ShellCommandTest.java

@@ -16,6 +16,7 @@ package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.JUnitBackports.assertThrows;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -28,12 +29,17 @@ import com.google.common.collect.ImmutableMap;
 import google.registry.testing.FakeClock;
 import google.registry.tools.ShellCommand.JCommanderCompletor;
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.PrintStream;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.List;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -44,8 +50,23 @@ public class ShellCommandTest {
   CommandRunner cli = mock(CommandRunner.class);
   FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ"));
 
+  PrintStream orgStdout;
+  PrintStream orgStderr;
+
   public ShellCommandTest() {}
 
+  @Before
+  public void setUp() {
+    orgStdout = System.out;
+    orgStderr = System.err;
+  }
+
+  @After
+  public void tearDown() {
+    System.setOut(orgStdout);
+    System.setErr(orgStderr);
+  }
+
   @Test
   public void testParsing() {
     assertThat(ShellCommand.parseCommand("foo bar 123 baz+ // comment \"string data\""))
@@ -226,6 +247,55 @@ public class ShellCommandTest {
         "testCommand --xorg ", 0, "", "Flag documentation: test organization\n  (PRIVATE, PUBLIC)");
   }
 
+  @Test
+  public void testEncapsulatedOutputStream_basicFuncionality() throws Exception {
+    ByteArrayOutputStream backing = new ByteArrayOutputStream();
+    PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
+    out.println("first line");
+    out.print("second line\ntrailing data");
+    out.flush();
+    assertThat(backing.toString())
+        .isEqualTo("out: first line\nout: second line\nout: trailing data\n");
+  }
+
+  @Test
+  public void testEncapsulatedOutputStream_emptyStream() throws Exception {
+    ByteArrayOutputStream backing = new ByteArrayOutputStream();
+    PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
+    out.flush();
+    assertThat(backing.toString()).isEqualTo("");
+  }
+
+  @Test
+  public void testEncapsulatedOutput_command() throws Exception {
+    RegistryToolEnvironment.ALPHA.setup();
+
+    // capture output (have to do this before the shell command is created)
+    ByteArrayOutputStream stdout = new ByteArrayOutputStream();
+    ByteArrayOutputStream stderr = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(stdout));
+    System.setErr(new PrintStream(stderr));
+    System.setIn(new ByteArrayInputStream("command1\n".getBytes(UTF_8)));
+
+    ShellCommand shellCommand =
+        new ShellCommand(
+            args -> {
+              System.out.println("first line");
+              System.err.println("second line");
+              System.out.print("fragmented ");
+              System.err.println("surprise!");
+              System.out.println("line");
+            });
+    shellCommand.encapsulateOutput = true;
+
+    shellCommand.run();
+    assertThat(stderr.toString()).isEmpty();
+    assertThat(stdout.toString())
+        .isEqualTo(
+            "out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n"
+                + "SUCCESS\n");
+  }
+
   @Parameters(commandDescription = "Test command")
   static class TestCommand implements Command {
     enum OrgType {