diff --git a/java/google/registry/tools/CreateLrpTokensCommand.java b/java/google/registry/tools/CreateLrpTokensCommand.java new file mode 100644 index 000000000..db80e9831 --- /dev/null +++ b/java/google/registry/tools/CreateLrpTokensCommand.java @@ -0,0 +1,152 @@ +// Copyright 2016 The Domain Registry Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.Sets.difference; +import static google.registry.model.ofy.ObjectifyService.ofy; +import static google.registry.model.registry.Registries.assertTldExists; +import static java.nio.charset.StandardCharsets.UTF_8; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; + +import com.google.common.base.Function; +import com.google.common.base.Splitter; +import com.google.common.collect.FluentIterable; +import com.google.common.collect.ImmutableSet; +import com.google.common.io.Files; +import com.google.common.io.LineReader; + +import com.googlecode.objectify.Key; +import com.googlecode.objectify.Work; + +import google.registry.model.domain.LrpToken; +import google.registry.tools.Command.GtechCommand; +import google.registry.tools.Command.RemoteApiCommand; +import google.registry.tools.params.PathParameter; + +import java.io.StringReader; +import java.nio.file.Path; +import java.util.Collection; +import java.util.Set; + +import javax.inject.Inject; + +/** + * Command to create one or more LRP tokens, given assignee(s) as either a parameter or a text file. + */ +@Parameters( + separators = " =", + commandDescription = "Create an LRP token for a given assignee (using -a) or import a text" + + " file of assignees for bulk token creation (using -i). Assignee/token pairs are printed" + + " to stdout, and should be piped to a file for distribution to assignees or for cleanup" + + " in the event of a command interruption.") +public final class CreateLrpTokensCommand implements RemoteApiCommand, GtechCommand { + + @Parameter( + names = {"-a", "--assignee"}, + description = "LRP token assignee") + private String assignee; + + @Parameter( + names = {"-t", "--tlds"}, + description = "Comma-delimited list of TLDs that the tokens to create will be valid on", + required = true) + private String tlds; + + @Parameter( + names = {"-i", "--input"}, + description = "Filename containing a list of assignees, newline-delimited", + validateWith = PathParameter.InputFile.class) + private Path assigneesFile; + + @Inject StringGenerator stringGenerator; + + private static final int TOKEN_LENGTH = 16; + private static final int BATCH_SIZE = 20; + + @Override + public void run() throws Exception { + checkArgument( + (assignee == null) == (assigneesFile != null), + "Exactly one of either assignee or filename must be specified."); + final Set validTlds = ImmutableSet.copyOf(Splitter.on(',').split(tlds)); + for (String tld : validTlds) { + assertTldExists(tld); + } + + LineReader reader = new LineReader( + (assignee == null) + ? Files.newReader(assigneesFile.toFile(), UTF_8) + : new StringReader(assignee)); + + String line = null; + do { + ImmutableSet.Builder tokensToSave = new ImmutableSet.Builder<>(); + for (String token : generateTokens(BATCH_SIZE)) { + line = reader.readLine(); + if (!isNullOrEmpty(line)) { + tokensToSave.add(new LrpToken.Builder() + .setAssignee(line) + .setToken(token) + .setValidTlds(validTlds) + .build()); + } + } + saveTokens(tokensToSave.build()); + } while (line != null); + } + + private void saveTokens(final ImmutableSet tokens) { + Collection savedTokens = ofy().transact(new Work>() { + @Override + public Collection run() { + return ofy().save().entities(tokens).now().values(); + }}); + for (LrpToken token : savedTokens) { + System.out.printf("%s,%s%n", token.getAssignee(), token.getToken()); + } + } + + /** + * This function generates at MOST {@code count} tokens, after filtering out any token strings + * that already exist. + * + *

Note that in the incredibly rare case that all generated tokens already exist, this function + * may return an empty set. + */ + private ImmutableSet generateTokens(int count) { + final ImmutableSet candidates = + ImmutableSet.copyOf(stringGenerator.createStrings(TOKEN_LENGTH, count)); + ImmutableSet> existingTokenKeys = FluentIterable.from(candidates) + .transform(new Function>() { + @Override + public Key apply(String input) { + return Key.create(LrpToken.class, input); + }}) + .toSet(); + ImmutableSet existingTokenStrings = FluentIterable + .from(ofy().load().keys(existingTokenKeys).values()) + .transform(new Function() { + @Override + public String apply(LrpToken input) { + return input.getToken(); + }}) + .toSet(); + return ImmutableSet.copyOf(difference(candidates, existingTokenStrings)); + } +} diff --git a/java/google/registry/tools/GtechTool.java b/java/google/registry/tools/GtechTool.java index e46095c0d..58b183d54 100644 --- a/java/google/registry/tools/GtechTool.java +++ b/java/google/registry/tools/GtechTool.java @@ -38,6 +38,7 @@ public final class GtechTool { .put("create_credit_balance", CreateCreditBalanceCommand.class) .put("create_domain", CreateDomainCommand.class) .put("create_host", CreateHostCommand.class) + .put("create_lrp_tokens", CreateLrpTokensCommand.class) .put("create_registrar_groups", CreateRegistrarGroupsCommand.class) .put("create_registrar", CreateRegistrarCommand.class) .put("create_sandbox_tld", CreateSandboxTldCommand.class) diff --git a/java/google/registry/tools/StringGenerator.java b/java/google/registry/tools/StringGenerator.java index 3834ae47e..dbbb06fc7 100644 --- a/java/google/registry/tools/StringGenerator.java +++ b/java/google/registry/tools/StringGenerator.java @@ -17,6 +17,10 @@ package google.registry.tools; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Strings.isNullOrEmpty; +import com.google.common.collect.ImmutableList; + +import java.util.Collection; + /** String generator. */ abstract class StringGenerator { @@ -41,4 +45,13 @@ abstract class StringGenerator { /** Generates a string of a specified length. */ abstract String createString(int length); + + /** Batch-generates an {@link ImmutableList} of strings of a specified length. */ + public Collection createStrings(int length, int count) { + ImmutableList.Builder listBuilder = new ImmutableList.Builder<>(); + for (int i = 0; i < count; i++) { + listBuilder.add(createString(length)); + } + return listBuilder.build(); + } } diff --git a/javatests/google/registry/tools/BUILD b/javatests/google/registry/tools/BUILD index 4325c63cd..6117e6887 100644 --- a/javatests/google/registry/tools/BUILD +++ b/javatests/google/registry/tools/BUILD @@ -28,6 +28,7 @@ java_library( "//third_party/java/joda_money", "//third_party/java/joda_time", "//third_party/java/json_simple", + "//third_party/java/jsr305_annotations", "//third_party/java/jsr330_inject", "//third_party/java/junit", "//third_party/java/mockito", diff --git a/javatests/google/registry/tools/CreateLrpTokensCommandTest.java b/javatests/google/registry/tools/CreateLrpTokensCommandTest.java new file mode 100644 index 000000000..d12173878 --- /dev/null +++ b/javatests/google/registry/tools/CreateLrpTokensCommandTest.java @@ -0,0 +1,204 @@ +// Copyright 2016 The Domain Registry Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.model.ofy.ObjectifyService.ofy; +import static google.registry.testing.DatastoreHelper.createTld; +import static google.registry.testing.DatastoreHelper.persistResource; +import static java.nio.charset.StandardCharsets.UTF_8; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.io.Files; + +import com.googlecode.objectify.Key; + +import google.registry.model.domain.LrpToken; +import google.registry.model.reporting.HistoryEntry; +import google.registry.tools.DeterministicStringGenerator.Rule; + +import java.io.File; +import java.io.IOException; +import java.util.Set; + +import javax.annotation.Nullable; + +import org.junit.Before; +import org.junit.Test; + +/** Unit tests for {@link CreateLrpTokensCommand}. */ +public class CreateLrpTokensCommandTest extends CommandTestCase { + + DeterministicStringGenerator stringGenerator = + new DeterministicStringGenerator("abcdefghijklmnopqrstuvwxyz"); + File assigneeFile; + String assigneeFilePath; + + @Before + public void init() throws IOException { + assigneeFile = tmpDir.newFile("lrp_assignees.txt"); + assigneeFilePath = assigneeFile.getPath(); + command.stringGenerator = stringGenerator; + createTld("tld"); + } + + @Test + public void testSuccess_oneAssignee() throws Exception { + runCommand("--assignee=domain.tld", "--tlds=tld"); + assertLrpTokens(createToken("abcdefghijklmnop", "domain.tld", ImmutableSet.of("tld"), null)); + assertInStdout("domain.tld,abcdefghijklmnop"); + } + + @Test + public void testSuccess_oneAssignee_tokenCollision() throws Exception { + LrpToken existingToken = persistResource(new LrpToken.Builder() + .setToken("abcdefghijklmnop") + .setAssignee("otherdomain.tld") + .setValidTlds(ImmutableSet.of("tld")) + .build()); + runCommand("--assignee=domain.tld", "--tlds=tld"); + assertLrpTokens( + existingToken, + createToken("qrstuvwxyzabcdef", "domain.tld", ImmutableSet.of("tld"), null)); + assertInStdout("domain.tld,qrstuvwxyzabcdef"); + } + + @Test + public void testSuccess_oneAssignee_byFile() throws Exception { + Files.write("domain.tld", assigneeFile, UTF_8); + runCommand("--input=" + assigneeFilePath, "--tlds=tld"); + assertLrpTokens(createToken("abcdefghijklmnop", "domain.tld", ImmutableSet.of("tld"), null)); + assertInStdout("domain.tld,abcdefghijklmnop"); + } + + @Test + public void testSuccess_emptyFile() throws Exception { + Files.write("", assigneeFile, UTF_8); + runCommand("--input=" + assigneeFilePath, "--tlds=tld"); + assertLrpTokens(); // no tokens exist + assertThat(getStdoutAsString()).isEmpty(); + } + + @Test + public void testSuccess_multipleAssignees_byFile() throws Exception { + Files.write("domain1.tld\ndomain2.tld\ndomain3.tld", assigneeFile, UTF_8); + runCommand("--input=" + assigneeFilePath, "--tlds=tld"); + + assertLrpTokens( + createToken("abcdefghijklmnop", "domain1.tld", ImmutableSet.of("tld"), null), + createToken("qrstuvwxyzabcdef", "domain2.tld", ImmutableSet.of("tld"), null), + createToken("ghijklmnopqrstuv", "domain3.tld", ImmutableSet.of("tld"), null)); + + assertInStdout("domain1.tld,abcdefghijklmnop"); + assertInStdout("domain2.tld,qrstuvwxyzabcdef"); + assertInStdout("domain3.tld,ghijklmnopqrstuv"); + } + + @Test + public void testSuccess_multipleAssignees_byFile_ignoreBlankLine() throws Exception { + Files.write("domain1.tld\n\ndomain2.tld", assigneeFile, UTF_8); + runCommand("--input=" + assigneeFilePath, "--tlds=tld"); + assertLrpTokens( + createToken("abcdefghijklmnop", "domain1.tld", ImmutableSet.of("tld"), null), + // Second deterministic token (qrstuvwxyzabcdef) still consumed but not assigned + createToken("ghijklmnopqrstuv", "domain2.tld", ImmutableSet.of("tld"), null)); + assertInStdout("domain1.tld,abcdefghijklmnop"); + assertInStdout("domain2.tld,ghijklmnopqrstuv"); + } + + @Test + public void testSuccess_largeFile() throws Exception { + int numberOfTokens = 67; + LrpToken[] expectedTokens = new LrpToken[numberOfTokens]; + // Prepend a counter to avoid collisions, 16-char alphabet will always generate the same string. + stringGenerator = + new DeterministicStringGenerator("abcdefghijklmnop", Rule.PREPEND_COUNTER); + command.stringGenerator = stringGenerator; + StringBuilder assigneeFileBuilder = new StringBuilder(); + for (int i = 0; i < numberOfTokens; i++) { + assigneeFileBuilder.append(String.format("domain%d.tld\n", i)); + expectedTokens[i] = + createToken( + String.format("%04d_abcdefghijklmnop", i), + String.format("domain%d.tld", i), + ImmutableSet.of("tld"), + null); + } + Files.write(assigneeFileBuilder, assigneeFile, UTF_8); + runCommand("--input=" + assigneeFilePath, "--tlds=tld"); + assertLrpTokens(expectedTokens); + for (int i = 0; i < numberOfTokens; i++) { + assertInStdout(String.format("domain%d.tld,%04d_abcdefghijklmnop", i, i)); + } + } + + @Test + public void testFailure_missingAssigneeOrFile() throws Exception { + thrown.expect( + IllegalArgumentException.class, + "Exactly one of either assignee or filename must be specified."); + runCommand("--tlds=tld"); + } + + @Test + public void testFailure_bothAssigneeAndFile() throws Exception { + thrown.expect( + IllegalArgumentException.class, + "Exactly one of either assignee or filename must be specified."); + runCommand("--assignee=domain.tld", "--tlds=tld", "--input=" + assigneeFilePath); + } + + @Test + public void testFailure_badTld() throws Exception { + thrown.expect(IllegalArgumentException.class, "TLD foo does not exist"); + runCommand("--assignee=domain.tld", "--tlds=foo"); + } + + private void assertLrpTokens(LrpToken... expected) throws Exception { + // Using ImmutableObject comparison here is tricky because updateTimestamp is not set on the + // expected LrpToken objects and will cause the assert to fail. + Iterable actual = ofy().load().type(LrpToken.class); + ImmutableMap.Builder actualTokenMapBuilder = new ImmutableMap.Builder<>(); + for (LrpToken token : actual) { + actualTokenMapBuilder.put(token.getToken(), token); + } + ImmutableMap actualTokenMap = actualTokenMapBuilder.build(); + assertThat(actualTokenMap).hasSize(expected.length); + for (LrpToken expectedToken : expected) { + LrpToken match = actualTokenMap.get(expectedToken.getToken()); + assertThat(match).isNotNull(); + assertThat(match.getAssignee()).isEqualTo(expectedToken.getAssignee()); + assertThat(match.getValidTlds()).containsExactlyElementsIn(expectedToken.getValidTlds()); + assertThat(match.getRedemptionHistoryEntry()) + .isEqualTo(expectedToken.getRedemptionHistoryEntry()); + } + } + + private LrpToken createToken( + String token, + String assignee, + Set validTlds, + @Nullable Key redemptionHistoryEntry) { + LrpToken.Builder tokenBuilder = new LrpToken.Builder() + .setAssignee(assignee) + .setValidTlds(validTlds) + .setToken(token); + if (redemptionHistoryEntry != null) { + tokenBuilder.setRedemptionHistoryEntry(redemptionHistoryEntry); + } + return tokenBuilder.build(); + } +} diff --git a/javatests/google/registry/tools/DeterministicStringGenerator.java b/javatests/google/registry/tools/DeterministicStringGenerator.java index f7dffbb52..708b635cb 100644 --- a/javatests/google/registry/tools/DeterministicStringGenerator.java +++ b/javatests/google/registry/tools/DeterministicStringGenerator.java @@ -33,6 +33,24 @@ import javax.inject.Named; class DeterministicStringGenerator extends StringGenerator { private Iterator iterator; + private final Rule rule; + private int counter = 0; + + /** String generation rules. */ + enum Rule { + + /** + * Simple string generation, cycling through sequential letters in the alphabet. May produce + * duplicates. + */ + DEFAULT, + + /** + * Same cyclical pattern as {@link Rule#DEFAULT}, prepending the iteration number and an + * underscore. Intended to avoid duplicates. + */ + PREPEND_COUNTER + } /** * Generates a string using sequential characters in the generator's alphabet, cycling back to the @@ -45,11 +63,22 @@ class DeterministicStringGenerator extends StringGenerator { for (int i = 0; i < length; i++) { password += iterator.next(); } - return password; + switch (rule) { + case PREPEND_COUNTER: + return String.format("%04d_%s", counter++, password); + case DEFAULT: + default: + return password; + } + } + + public DeterministicStringGenerator(@Named("alphabet") String alphabet, Rule rule) { + super(alphabet); + iterator = Iterators.cycle(charactersOf(alphabet)); + this.rule = rule; } public DeterministicStringGenerator(@Named("alphabet") String alphabet) { - super(alphabet); - iterator = Iterators.cycle(charactersOf(alphabet)); + this(alphabet, Rule.DEFAULT); } } diff --git a/javatests/google/registry/tools/GetLrpTokenCommandTest.java b/javatests/google/registry/tools/GetLrpTokenCommandTest.java index 70ffb4fd7..60a84e04c 100644 --- a/javatests/google/registry/tools/GetLrpTokenCommandTest.java +++ b/javatests/google/registry/tools/GetLrpTokenCommandTest.java @@ -26,8 +26,7 @@ import org.junit.Before; import org.junit.Test; /** Unit tests for {@link GetLrpTokenCommand}. */ -public class GetLrpTokenCommandTest - extends CommandTestCase { +public class GetLrpTokenCommandTest extends CommandTestCase { @Before public void before() {