mirror of
https://github.com/google/nomulus.git
synced 2025-05-13 07:57:13 +02:00
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:
parent
7bf0b059a6
commit
013558c814
3 changed files with 108 additions and 18 deletions
|
@ -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.
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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<>();
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue