Clean up token generation (#205)

* Clean up token generation

- Allow tokenLength of 0
- If specifying a token length of 0, throw an error if numTokens > 1

* Allow generation of 0-length strings

* Allow for --tokens option to generate specific tokens

* Revert String generators and disallow 0 'length' param

* Add verifyInput method and batch the listed tokens

* Check the number of tokens created

core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -35,6 +35,8 @@ import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -51,6 +53,7 @@ import java.util.Collection;
 import java.util.Deque;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import javax.inject.Named;
 import org.joda.time.DateTime;
@@ -64,6 +67,13 @@ import org.joda.time.DateTime;
 @NonFinalForTesting
 class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
+  @Parameter(
+      names = {"--tokens"},
+      description =
+          "Comma-separated list of exact tokens to generate, otherwise use --prefix or "
+              + "--domain_names_file to generate tokens if the exact strings do not matter")
+  private List<String> tokenStrings;
+
   @Parameter(
       names = {"-p", "--prefix"},
       description = "Allocation token prefix; defaults to blank")
@@ -132,24 +142,10 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
   @Override
   public void run() throws IOException {
-    checkArgument(
-        (numTokens > 0) ^ (domainNamesFile != null),
-        "Must specify either --number or --domain_names_file, but not both");
-
-    checkArgument(
-        !(UNLIMITED_USE.equals(tokenType) && CollectionUtils.isNullOrEmpty(tokenStatusTransitions)),
-        "For UNLIMITED_USE tokens, must specify --token_status_transitions");
-
-    // A list consisting solely of the empty string means user error when formatting parameters
-    checkArgument(
-        !ImmutableList.of("").equals(allowedClientIds),
-        "Either omit --allowed_client_ids if all registrars are allowed, or include a"
-            + " comma-separated list");
-
-    checkArgument(
-        !ImmutableList.of("").equals(allowedTlds),
-        "Either omit --allowed_tlds if all TLDs are allowed, or include a comma-separated list");
-
+    verifyInput();
+    if (tokenStrings != null) {
+      numTokens = tokenStrings.size();
+    }
     Deque<String> domainNames;
     if (domainNamesFile == null) {
       domainNames = null;
@@ -166,8 +162,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
     int tokensSaved = 0;
     do {
       ImmutableSet<AllocationToken> tokens =
-          generateTokens(BATCH_SIZE).stream()
-              .limit(numTokens - tokensSaved)
+          getNextTokenBatch(tokensSaved)
               .map(
                   t -> {
                     AllocationToken.Builder token =
@@ -189,6 +184,62 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
     } while (tokensSaved < numTokens);
   }
 
+  private void verifyInput() {
+    int inputMethods = 0;
+    if (numTokens > 0) {
+      inputMethods++;
+    }
+    if (domainNamesFile != null) {
+      inputMethods++;
+    }
+    if (tokenStrings != null && !tokenStrings.isEmpty()) {
+      inputMethods++;
+    }
+    checkArgument(
+        inputMethods == 1,
+        "Must specify exactly one of '--number', '--domain_names_file', and '--tokens'");
+
+    checkArgument(
+        tokenLength > 0,
+        "Token length should not be 0. To generate exact tokens, use the --tokens parameter.");
+
+    checkArgument(
+        !(UNLIMITED_USE.equals(tokenType) && CollectionUtils.isNullOrEmpty(tokenStatusTransitions)),
+        "For UNLIMITED_USE tokens, must specify --token_status_transitions");
+
+    // A list consisting solely of the empty string means user error when formatting parameters
+    checkArgument(
+        !ImmutableList.of("").equals(allowedClientIds),
+        "Either omit --allowed_client_ids if all registrars are allowed, or include a"
+            + " comma-separated list");
+
+    checkArgument(
+        !ImmutableList.of("").equals(allowedTlds),
+        "Either omit --allowed_tlds if all TLDs are allowed, or include a comma-separated list");
+
+    if (tokenStrings != null) {
+      verifyTokenStringsDoNotExist();
+    }
+  }
+
+  private void verifyTokenStringsDoNotExist() {
+    ImmutableSet<String> existingTokenStrings =
+        getExistingTokenStrings(ImmutableSet.copyOf(tokenStrings));
+    checkArgument(
+        existingTokenStrings.isEmpty(),
+        String.format(
+            "Cannot create specified tokens; the following tokens already exist: %s",
+            existingTokenStrings));
+  }
+
+  private Stream<String> getNextTokenBatch(int tokensSaved) {
+    if (tokenStrings != null) {
+      return Streams.stream(Iterables.limit(Iterables.skip(tokenStrings, tokensSaved), BATCH_SIZE));
+    } else {
+      return generateTokens(BATCH_SIZE).stream().limit(numTokens - tokensSaved);
+    }
+  }
+
   @VisibleForTesting
   int saveTokens(final ImmutableSet<AllocationToken> tokens) {
     Collection<AllocationToken> savedTokens =
@@ -210,14 +261,17 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
         stringGenerator.createStrings(tokenLength, count).stream()
             .map(s -> prefix + s)
             .collect(toImmutableSet());
+    ImmutableSet<String> existingTokenStrings = getExistingTokenStrings(candidates);
+    return ImmutableSet.copyOf(difference(candidates, existingTokenStrings));
+  }
+
+  private ImmutableSet<String> getExistingTokenStrings(ImmutableSet<String> candidates) {
     ImmutableSet<Key<AllocationToken>> existingTokenKeys =
         candidates.stream()
             .map(input -> Key.create(AllocationToken.class, input))
             .collect(toImmutableSet());
-    ImmutableSet<String> existingTokenStrings =
-        ofy().load().keys(existingTokenKeys).values().stream()
-            .map(AllocationToken::getToken)
-            .collect(toImmutableSet());
-    return ImmutableSet.copyOf(difference(candidates, existingTokenStrings));
+    return ofy().load().keys(existingTokenKeys).values().stream()
+        .map(AllocationToken::getToken)
+        .collect(toImmutableSet());
   }
 }