HADOOP-13109. Add ability to edit existing token file via dtutil -alias flag (Matthew Paduano via aw)

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/DtFileOperations.java

@@ -199,6 +199,29 @@ public final class DtFileOperations {
     doFormattedWrite(tokenFile, fileFormat, creds, conf);
   }
 
+  /** Alias a token from a file and save back to file in the local filesystem.
+   *  @param tokenFile a local File object to hold the input and output.
+   *  @param fileFormat a string equal to FORMAT_PB or FORMAT_JAVA, for output
+   *  @param alias overwrite service field of fetched token with this text.
+   *  @param service only apply alias to tokens matching this service text.
+   *  @param conf Configuration object passed along.
+   *  @throws IOException
+   */
+  public static void aliasTokenFile(File tokenFile, String fileFormat,
+      Text alias, Text service, Configuration conf) throws Exception {
+    Credentials newCreds = new Credentials();
+    Credentials creds = Credentials.readTokenStorageFile(tokenFile, conf);
+    for (Token<?> token : creds.getAllTokens()) {
+      newCreds.addToken(token.getService(), token);
+      if (token.getService().equals(service)) {
+        Token<?> aliasedToken = token.copyToken();
+        aliasedToken.setService(alias);
+        newCreds.addToken(alias, aliasedToken);
+      }
+    }
+    doFormattedWrite(tokenFile, fileFormat, newCreds, conf);
+  }
+
   /** Append tokens from list of files in local filesystem, saving to last file.
    *  @param tokenFiles list of local File objects.  Last file holds the output.
    *  @param fileFormat a string equal to FORMAT_PB or FORMAT_JAVA, for output

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/DtUtilShell.java

@@ -41,7 +41,7 @@ public class DtUtilShell extends CommandShell {
       DtFileOperations.FORMAT_PB + ")]";
   public static final String DT_USAGE = "hadoop dtutil " +
       "[-keytab <keytab_file> -principal <principal_name>] " +
-      "subcommand (help|print|get|append|cancel|remove|renew) " +
+      "subcommand (help|print|get|edit|append|cancel|remove|renew) " +
          FORMAT_SUBSTRING + " [-alias <alias>] filename...";
 
   // command line options
@@ -50,6 +50,7 @@ public class DtUtilShell extends CommandShell {
   private static final String PRINCIPAL = "-principal";
   private static final String PRINT = "print";
   private static final String GET = "get";
+  private static final String EDIT = "edit";
   private static final String APPEND = "append";
   private static final String CANCEL = "cancel";
   private static final String REMOVE = "remove";
@@ -127,6 +128,8 @@ public class DtUtilShell extends CommandShell {
           setSubCommand(new Print());
         } else if (command.equals(GET)) {
           setSubCommand(new Get(args[++i]));
+        } else if (command.equals(EDIT)) {
+          setSubCommand(new Edit());
         } else if (command.equals(APPEND)) {
           setSubCommand(new Append());
         } else if (command.equals(CANCEL)) {
@@ -172,10 +175,12 @@ public class DtUtilShell extends CommandShell {
 
   @Override
   public String getCommandUsage() {
-    return String.format("%n%s%n   %s%n   %s%n   %s%n   %s%n   %s%n   %s%n%n",
+    return String.format(
-                  DT_USAGE, (new Print()).getUsage(), (new Get()).getUsage(),
+        "%n%s%n   %s%n   %s%n   %s%n   %s%n   %s%n   %s%n   %s%n%n",
-                  (new Append()).getUsage(), (new Remove(true)).getUsage(),
+        DT_USAGE, (new Print()).getUsage(), (new Get()).getUsage(),
-                  (new Remove(false)).getUsage(), (new Renew()).getUsage());
+        (new Edit()).getUsage(), (new Append()).getUsage(),
+        (new Remove(true)).getUsage(), (new Remove(false)).getUsage(),
+        (new Renew()).getUsage());
   }
 
   private class Print extends SubCommand {
@@ -242,6 +247,38 @@ public class DtUtilShell extends CommandShell {
     }
   }
 
+  private class Edit extends SubCommand {
+    public static final String EDIT_USAGE =
+        "dtutil edit -service <service> -alias <alias> " +
+        FORMAT_SUBSTRING + "filename...";
+
+    @Override
+    public boolean validate() {
+      if (service == null) {
+        LOG.error("must pass -service field with dtutil edit command");
+        return false;
+      }
+      if (alias == null) {
+        LOG.error("must pass -alias field with dtutil edit command");
+        return false;
+      }
+      return true;
+    }
+
+    @Override
+    public void execute() throws Exception {
+      for (File tokenFile : tokenFiles) {
+        DtFileOperations.aliasTokenFile(
+            tokenFile, format, alias, service, getConf());
+      }
+    }
+
+    @Override
+    public String getUsage() {
+      return EDIT_USAGE;
+    }
+  }
+
   private class Append extends SubCommand {
     public static final String APPEND_USAGE =
         "dtutil append " + FORMAT_SUBSTRING + "filename...";

hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/TestDtUtilShell.java

@@ -157,6 +157,30 @@ public class TestDtUtilShell {
                 outContent.toString().contains(SERVICE.toString()));
   }
 
+  @Test
+  public void testEdit() throws Exception {
+    String oldService = SERVICE2.toString();
+    String newAlias = "newName:12345";
+    args = new String[] {"edit",
+        "-service", oldService, "-alias", newAlias, tokenFilename2};
+    rc = dt.run(args);
+    assertEquals("test simple edit exit code", 0, rc);
+    args = new String[] {"print", "-alias", oldService, tokenFilename2};
+    rc = dt.run(args);
+    assertEquals("test simple edit print old exit code", 0, rc);
+    assertTrue("test simple edit output kind old:\n" + outContent.toString(),
+               outContent.toString().contains(KIND.toString()));
+    assertTrue("test simple edit output service old:\n" + outContent.toString(),
+               outContent.toString().contains(oldService));
+    args = new String[] {"print", "-alias", newAlias, tokenFilename2};
+    rc = dt.run(args);
+    assertEquals("test simple edit print new exit code", 0, rc);
+    assertTrue("test simple edit output kind new:\n" + outContent.toString(),
+               outContent.toString().contains(KIND.toString()));
+    assertTrue("test simple edit output service new:\n" + outContent.toString(),
+               outContent.toString().contains(newAlias));
+  }
+
   @Test
   public void testAppend() throws Exception {
     args = new String[] {"append", tokenFilename, tokenFilename2};