diff --git a/java/google/registry/tools/RegistryCli.java b/java/google/registry/tools/RegistryCli.java index 0b0caf157..ac820d114 100644 --- a/java/google/registry/tools/RegistryCli.java +++ b/java/google/registry/tools/RegistryCli.java @@ -104,10 +104,9 @@ final class RegistryCli implements AutoCloseable, CommandRunner { // Create the "help" and "shell" commands (these are special in that they don't have a default // constructor). jcommander.addCommand("help", new HelpCommand(jcommander)); - ShellCommand shellCommand = null; if (isFirstUse) { isFirstUse = false; - shellCommand = new ShellCommand(this); + ShellCommand shellCommand = new ShellCommand(this); // We have to build the completions based on the jcommander *before* we add the "shell" // command - to avoid completion for the "shell" command itself. shellCommand.buildCompletions(jcommander); @@ -139,11 +138,6 @@ final class RegistryCli implements AutoCloseable, CommandRunner { checkState(RegistryToolEnvironment.get() == environment, "RegistryToolEnvironment argument pre-processing kludge failed."); - // We have to set the prompt here, because the environment wasn't set until this point - if (shellCommand != null) { - shellCommand.setPrompt(String.format("nom@%s > ", environment)); - } - // JCommander stores sub-commands as nested JCommander objects containing a list of user objects // to be populated. Extract the subcommand by getting the JCommander wrapper and then // retrieving the first (and, by virtue of our usage, only) object from it. diff --git a/java/google/registry/tools/ShellCommand.java b/java/google/registry/tools/ShellCommand.java index f0a5dc734..a56b68aa7 100644 --- a/java/google/registry/tools/ShellCommand.java +++ b/java/google/registry/tools/ShellCommand.java @@ -19,11 +19,15 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static java.nio.charset.StandardCharsets.US_ASCII; import com.beust.jcommander.JCommander; +import com.beust.jcommander.Parameter; import com.beust.jcommander.ParameterDescription; import com.beust.jcommander.Parameters; 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 google.registry.util.Clock; +import google.registry.util.SystemClock; import java.io.BufferedReader; import java.io.File; import java.io.IOException; @@ -41,6 +45,8 @@ import jline.ConsoleReader; import jline.ConsoleReaderInputStream; import jline.FileNameCompletor; import jline.History; +import org.joda.time.DateTime; +import org.joda.time.Duration; /** * Implements a tiny shell interpreter for the nomulus tool. @@ -52,10 +58,22 @@ import jline.History; public class ShellCommand implements Command { private static final String HISTORY_FILE = ".nomulus_history"; + private static final String RESET = "\u001b[0m"; + 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 final CommandRunner runner; private final BufferedReader lineReader; private final ConsoleReader consoleReader; + private final Clock clock; + + @Parameter( + names = {"--dont_exit_on_idle"}, + description = + "Prevents the shell from exiting on PROD after the 1 hour idle delay. " + + "Will instead warn you and require re-running the command.") + boolean dontExitOnIdle = false; public ShellCommand(CommandRunner runner) throws IOException { this.runner = runner; @@ -71,20 +89,29 @@ public class ShellCommand implements Command { consoleReader = null; } this.lineReader = new BufferedReader(new InputStreamReader(in, US_ASCII)); + this.clock = new SystemClock(); } @VisibleForTesting - ShellCommand(InputStream in, CommandRunner runner) { + ShellCommand(BufferedReader bufferedReader, Clock clock, CommandRunner runner) { this.runner = runner; - this.lineReader = new BufferedReader(new InputStreamReader(in, US_ASCII)); + this.lineReader = bufferedReader; + this.clock = clock; this.consoleReader = null; } - public ShellCommand setPrompt(String prompt) { - if (consoleReader != null) { - consoleReader.setDefaultPrompt(prompt); + private void setPrompt(RegistryToolEnvironment environment, boolean alert) { + if (consoleReader == null) { + return; + } + if (alert) { + consoleReader.setDefaultPrompt( + String.format("nom@%s%s%s > ", ALERT_COLOR, environment, RESET)); + } else { + consoleReader.setDefaultPrompt( + String.format( + "nom@%s%s%s > ", NON_ALERT_COLOR, Ascii.toLowerCase(environment.toString()), RESET)); } - return this; } public ShellCommand buildCompletions(JCommander jcommander) { @@ -101,8 +128,21 @@ public class ShellCommand implements Command { /** Run the shell until the user presses "Ctrl-D". */ @Override public void run() { + // On Production we want to be extra careful - to prevent accidental use. + boolean beExtraCareful = (RegistryToolEnvironment.get() == RegistryToolEnvironment.PRODUCTION); + setPrompt(RegistryToolEnvironment.get(), beExtraCareful); String line; + DateTime lastTime = clock.nowUtc(); while ((line = getLine()) != null) { + // Make sure we're not idle for too long. Only relevant when we're "extra careful" + if (!dontExitOnIdle + && beExtraCareful + && lastTime.plus(IDLE_THRESHOLD).isBefore(clock.nowUtc())) { + throw new RuntimeException( + "Been idle for too long, while in 'extra careful' mode. " + + "The last command was saved in history. Please rerun the shell and try again."); + } + lastTime = clock.nowUtc(); String[] lineArgs = parseCommand(line); if (lineArgs.length == 0) { continue; diff --git a/javatests/google/registry/tools/ShellCommandTest.java b/javatests/google/registry/tools/ShellCommandTest.java index 6709207f6..894418c94 100644 --- a/javatests/google/registry/tools/ShellCommandTest.java +++ b/javatests/google/registry/tools/ShellCommandTest.java @@ -16,8 +16,8 @@ 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.US_ASCII; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.beust.jcommander.JCommander; import com.beust.jcommander.MissingCommandException; @@ -25,10 +25,15 @@ import com.beust.jcommander.Parameter; import com.beust.jcommander.Parameters; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import google.registry.testing.FakeClock; import google.registry.tools.ShellCommand.JCommanderCompletor; -import java.io.ByteArrayInputStream; +import java.io.BufferedReader; +import java.io.IOException; +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.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -37,6 +42,7 @@ import org.junit.runners.JUnit4; public class ShellCommandTest { CommandRunner cli = mock(CommandRunner.class); + FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ")); public ShellCommandTest() {} @@ -49,12 +55,25 @@ public class ShellCommandTest { assertThat(ShellCommand.parseCommand("")).isEqualTo(new String[0]); } + private ShellCommand createShellCommand( + CommandRunner commandRunner, Duration delay, String... commands) throws Exception { + ArrayDeque queue = new ArrayDeque(ImmutableList.copyOf(commands)); + BufferedReader bufferedReader = mock(BufferedReader.class); + when(bufferedReader.readLine()).thenAnswer((x) -> { + clock.advanceBy(delay); + if (queue.isEmpty()) { + throw new IOException(); + } + return queue.poll(); + }); + return new ShellCommand(bufferedReader, clock, commandRunner); + } + @Test - public void testCommandProcessing() { - String testData = "test1 foo bar\ntest2 foo bar\n"; + public void testCommandProcessing() throws Exception { MockCli cli = new MockCli(); ShellCommand shellCommand = - new ShellCommand(new ByteArrayInputStream(testData.getBytes(US_ASCII)), cli); + createShellCommand(cli, Duration.ZERO, "test1 foo bar", "test2 foo bar"); shellCommand.run(); assertThat(cli.calls) .containsExactly( @@ -62,6 +81,43 @@ public class ShellCommandTest { .inOrder(); } + @Test + public void testNoIdleWhenInAlpha() throws Exception { + RegistryToolEnvironment.ALPHA.setup(); + MockCli cli = new MockCli(); + ShellCommand shellCommand = + createShellCommand(cli, Duration.standardDays(1), "test1 foo bar", "test2 foo bar"); + shellCommand.run(); + } + + @Test + public void testNoIdleWhenInSandbox() throws Exception { + RegistryToolEnvironment.SANDBOX.setup(); + MockCli cli = new MockCli(); + ShellCommand shellCommand = + createShellCommand(cli, Duration.standardDays(1), "test1 foo bar", "test2 foo bar"); + shellCommand.run(); + } + + @Test + public void testIdleWhenOverHourInProduction() throws Exception { + RegistryToolEnvironment.PRODUCTION.setup(); + MockCli cli = new MockCli(); + ShellCommand shellCommand = + createShellCommand(cli, Duration.standardMinutes(61), "test1 foo bar", "test2 foo bar"); + RuntimeException exception = assertThrows(RuntimeException.class, shellCommand::run); + assertThat(exception).hasMessageThat().contains("Been idle for too long"); + } + + @Test + public void testNoIdleWhenUnderHourInProduction() throws Exception { + RegistryToolEnvironment.PRODUCTION.setup(); + MockCli cli = new MockCli(); + ShellCommand shellCommand = + createShellCommand(cli, Duration.standardMinutes(59), "test1 foo bar", "test2 foo bar"); + shellCommand.run(); + } + static class MockCli implements CommandRunner { public ArrayList> calls = new ArrayList<>();