Make it VERY clear when nomulus shell is on PROD

We don't want people to accidentally run commands on prod thinking they were on
Alpha / Sandbox.

To do that - we add 2 safeguards:

1) when on prod, the shell has a strong RED "PRODUCTION" in the commandline, while on alpha/sandbox it's green.

2) if a prod shell is idle for > 1h, it exits. So don't accidentally use a prod shell from a long time ago.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=191931731
This commit is contained in:
guyben 2018-04-06 13:28:16 -07:00 committed by Ben McIlwain
parent 7bf0b059a6
commit 013558c814
3 changed files with 108 additions and 18 deletions

View file

@ -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 // Create the "help" and "shell" commands (these are special in that they don't have a default
// constructor). // constructor).
jcommander.addCommand("help", new HelpCommand(jcommander)); jcommander.addCommand("help", new HelpCommand(jcommander));
ShellCommand shellCommand = null;
if (isFirstUse) { if (isFirstUse) {
isFirstUse = false; 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" // We have to build the completions based on the jcommander *before* we add the "shell"
// command - to avoid completion for the "shell" command itself. // command - to avoid completion for the "shell" command itself.
shellCommand.buildCompletions(jcommander); shellCommand.buildCompletions(jcommander);
@ -139,11 +138,6 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
checkState(RegistryToolEnvironment.get() == environment, checkState(RegistryToolEnvironment.get() == environment,
"RegistryToolEnvironment argument pre-processing kludge failed."); "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 // 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 // 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. // retrieving the first (and, by virtue of our usage, only) object from it.

View file

@ -19,11 +19,15 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.nio.charset.StandardCharsets.US_ASCII; import static java.nio.charset.StandardCharsets.US_ASCII;
import com.beust.jcommander.JCommander; import com.beust.jcommander.JCommander;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.ParameterDescription; import com.beust.jcommander.ParameterDescription;
import com.beust.jcommander.Parameters; import com.beust.jcommander.Parameters;
import com.google.common.annotations.VisibleForTesting; import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Ascii;
import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableTable; import com.google.common.collect.ImmutableTable;
import google.registry.util.Clock;
import google.registry.util.SystemClock;
import java.io.BufferedReader; import java.io.BufferedReader;
import java.io.File; import java.io.File;
import java.io.IOException; import java.io.IOException;
@ -41,6 +45,8 @@ import jline.ConsoleReader;
import jline.ConsoleReaderInputStream; import jline.ConsoleReaderInputStream;
import jline.FileNameCompletor; import jline.FileNameCompletor;
import jline.History; import jline.History;
import org.joda.time.DateTime;
import org.joda.time.Duration;
/** /**
* Implements a tiny shell interpreter for the nomulus tool. * Implements a tiny shell interpreter for the nomulus tool.
@ -52,10 +58,22 @@ import jline.History;
public class ShellCommand implements Command { public class ShellCommand implements Command {
private static final String HISTORY_FILE = ".nomulus_history"; 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 CommandRunner runner;
private final BufferedReader lineReader; private final BufferedReader lineReader;
private final ConsoleReader consoleReader; 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 { public ShellCommand(CommandRunner runner) throws IOException {
this.runner = runner; this.runner = runner;
@ -71,20 +89,29 @@ public class ShellCommand implements Command {
consoleReader = null; consoleReader = null;
} }
this.lineReader = new BufferedReader(new InputStreamReader(in, US_ASCII)); this.lineReader = new BufferedReader(new InputStreamReader(in, US_ASCII));
this.clock = new SystemClock();
} }
@VisibleForTesting @VisibleForTesting
ShellCommand(InputStream in, CommandRunner runner) { ShellCommand(BufferedReader bufferedReader, Clock clock, CommandRunner runner) {
this.runner = runner; this.runner = runner;
this.lineReader = new BufferedReader(new InputStreamReader(in, US_ASCII)); this.lineReader = bufferedReader;
this.clock = clock;
this.consoleReader = null; this.consoleReader = null;
} }
public ShellCommand setPrompt(String prompt) { private void setPrompt(RegistryToolEnvironment environment, boolean alert) {
if (consoleReader != null) { if (consoleReader == null) {
consoleReader.setDefaultPrompt(prompt); 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) { public ShellCommand buildCompletions(JCommander jcommander) {
@ -101,8 +128,21 @@ public class ShellCommand implements Command {
/** Run the shell until the user presses "Ctrl-D". */ /** Run the shell until the user presses "Ctrl-D". */
@Override @Override
public void run() { 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; String line;
DateTime lastTime = clock.nowUtc();
while ((line = getLine()) != null) { 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); String[] lineArgs = parseCommand(line);
if (lineArgs.length == 0) { if (lineArgs.length == 0) {
continue; continue;

View file

@ -16,8 +16,8 @@ package google.registry.tools;
import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertThat;
import static google.registry.testing.JUnitBackports.assertThrows; 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.mock;
import static org.mockito.Mockito.when;
import com.beust.jcommander.JCommander; import com.beust.jcommander.JCommander;
import com.beust.jcommander.MissingCommandException; import com.beust.jcommander.MissingCommandException;
@ -25,10 +25,15 @@ import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters; import com.beust.jcommander.Parameters;
import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMap;
import google.registry.testing.FakeClock;
import google.registry.tools.ShellCommand.JCommanderCompletor; 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.ArrayList;
import java.util.List; import java.util.List;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.junit.Test; import org.junit.Test;
import org.junit.runner.RunWith; import org.junit.runner.RunWith;
import org.junit.runners.JUnit4; import org.junit.runners.JUnit4;
@ -37,6 +42,7 @@ import org.junit.runners.JUnit4;
public class ShellCommandTest { public class ShellCommandTest {
CommandRunner cli = mock(CommandRunner.class); CommandRunner cli = mock(CommandRunner.class);
FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ"));
public ShellCommandTest() {} public ShellCommandTest() {}
@ -49,12 +55,25 @@ public class ShellCommandTest {
assertThat(ShellCommand.parseCommand("")).isEqualTo(new String[0]); assertThat(ShellCommand.parseCommand("")).isEqualTo(new String[0]);
} }
private ShellCommand createShellCommand(
CommandRunner commandRunner, Duration delay, String... commands) throws Exception {
ArrayDeque<String> queue = new ArrayDeque<String>(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 @Test
public void testCommandProcessing() { public void testCommandProcessing() throws Exception {
String testData = "test1 foo bar\ntest2 foo bar\n";
MockCli cli = new MockCli(); MockCli cli = new MockCli();
ShellCommand shellCommand = ShellCommand shellCommand =
new ShellCommand(new ByteArrayInputStream(testData.getBytes(US_ASCII)), cli); createShellCommand(cli, Duration.ZERO, "test1 foo bar", "test2 foo bar");
shellCommand.run(); shellCommand.run();
assertThat(cli.calls) assertThat(cli.calls)
.containsExactly( .containsExactly(
@ -62,6 +81,43 @@ public class ShellCommandTest {
.inOrder(); .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 { static class MockCli implements CommandRunner {
public ArrayList<ImmutableList<String>> calls = new ArrayList<>(); public ArrayList<ImmutableList<String>> calls = new ArrayList<>();