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
// 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.

View file

@ -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;

View file

@ -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<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
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<ImmutableList<String>> calls = new ArrayList<>();