Support auto-resolving private key files based on token key ID

Lets a user specify a private key directory either with a command
line argument or via an environment variable. If a directory is
provided, the private key to use will be attempted auto-resolved
based on the key ID stored in the token. This only applies if the
key ID is comprised of exclusively path-safe characters.

5 files changed, 109 insertions, 17 deletions

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/ToolInvocation.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/ToolInvocation.java
index d1ff2687137..13df714a268 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/ToolInvocation.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/ToolInvocation.java
@@ -6,6 +6,7 @@ import org.apache.commons.cli.CommandLine;
 import java.io.InputStream;
 import java.io.PrintStream;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * @author vekterli
@@ -17,4 +18,10 @@ public record ToolInvocation(CommandLine arguments,
                              PrintStream stdError,
                              boolean debugMode) {
 
+    public void printIfDebug(Supplier<String> stringSupplier) {
+        if (debugMode) {
+            stdError.println(stringSupplier.get());
+        }
+    }
+
 }
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
index 87d3cb4d9f0..5cb40aa8f3b 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
@@ -19,7 +19,7 @@ public class CipherUtils {
      *
      * @param input source stream to read from
      * @param output destination stream to write to
-     * @param cipher an {@link AeadCipher} created with for either encryption or decryption
+     * @param cipher an {@link AeadCipher} created for either encryption or decryption
      * @throws IOException if any file operation fails
      */
     public static void streamEncipher(InputStream input, OutputStream output, AeadCipher cipher) throws IOException {
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
index 2cc724538d4..b22afb7e5fb 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
@@ -1,8 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.security.tool.crypto;
 
-import com.yahoo.security.KeyId;
-import com.yahoo.security.KeyUtils;
 import com.yahoo.security.SealedSharedKey;
 import com.yahoo.security.SharedKeyGenerator;
 import com.yahoo.vespa.security.tool.CliUtils;
@@ -12,11 +10,12 @@ import com.yahoo.vespa.security.tool.ToolInvocation;
 import org.apache.commons.cli.Option;
 
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Paths;
 import java.util.List;
 import java.util.Optional;
 
+import static com.yahoo.vespa.security.tool.crypto.ToolUtils.PRIVATE_KEY_DIR_OPTION;
+import static com.yahoo.vespa.security.tool.crypto.ToolUtils.PRIVATE_KEY_FILE_OPTION;
+
 /**
  * Tooling for decrypting a file using a private key that corresponds to the public key used
  * to originally encrypt the file.
@@ -28,7 +27,6 @@ import java.util.Optional;
 public class DecryptTool implements Tool {
 
     static final String OUTPUT_FILE_OPTION      = "output-file";
-    static final String PRIVATE_KEY_FILE_OPTION = "private-key-file";
     static final String EXPECTED_KEY_ID_OPTION  = "expected-key-id";
     static final String TOKEN_OPTION            = "token";
 
@@ -46,6 +44,13 @@ public class DecryptTool implements Tool {
                     .required(false)
                     .desc("Private key file in Base58 encoded format")
                     .build(),
+            Option.builder("d")
+                    .longOpt(PRIVATE_KEY_DIR_OPTION)
+                    .hasArg(true)
+                    .required(false)
+                    .desc("Private key file directory used for automatically looking up " +
+                          "private keys based on the key ID specified as part of a token.")
+                    .build(),
             Option.builder("e")
                     .longOpt(EXPECTED_KEY_ID_OPTION)
                     .hasArg(true)
@@ -83,17 +88,14 @@ public class DecryptTool implements Tool {
             if (leftoverArgs.length != 1) {
                 throw new IllegalArgumentException("Expected exactly 1 file argument to decrypt");
             }
-            var inputArg   = leftoverArgs[0];
-            var maybeKeyId = Optional.ofNullable(arguments.hasOption(EXPECTED_KEY_ID_OPTION)
-                                                 ? arguments.getOptionValue(EXPECTED_KEY_ID_OPTION)
-                                                 : null);
+            var inputArg    = leftoverArgs[0];
+            var maybeKeyId  = Optional.ofNullable(arguments.getOptionValue(EXPECTED_KEY_ID_OPTION));
             var outputArg   = CliUtils.optionOrThrow(arguments, OUTPUT_FILE_OPTION);
             var tokenString = CliUtils.optionOrThrow(arguments, TOKEN_OPTION);
             var sealedSharedKey = SealedSharedKey.fromTokenString(tokenString.strip());
             ToolUtils.verifyExpectedKeyId(sealedSharedKey, maybeKeyId);
 
-            var privKeyPath  = Paths.get(CliUtils.optionOrThrow(arguments, PRIVATE_KEY_FILE_OPTION));
-            var privateKey   = KeyUtils.fromBase58EncodedX25519PrivateKey(Files.readString(privKeyPath).strip());
+            var privateKey   = ToolUtils.resolvePrivateKeyFromInvocation(invocation, sealedSharedKey.keyId());
             var secretShared = SharedKeyGenerator.fromSealedKey(sealedSharedKey, privateKey);
             var cipher       = SharedKeyGenerator.makeAesGcmDecryptionCipher(secretShared);
 
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ResealTool.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ResealTool.java
index e9bc0ae8fee..83fdf6998df 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ResealTool.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ResealTool.java
@@ -12,11 +12,12 @@ import com.yahoo.vespa.security.tool.ToolInvocation;
 import org.apache.commons.cli.Option;
 
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Paths;
 import java.util.List;
 import java.util.Optional;
 
+import static com.yahoo.vespa.security.tool.crypto.ToolUtils.PRIVATE_KEY_DIR_OPTION;
+import static com.yahoo.vespa.security.tool.crypto.ToolUtils.PRIVATE_KEY_FILE_OPTION;
+
 /**
  * Tooling for resealing a token for another recipient. This allows for delegating
  * decryption to another party without having to reveal the private key of the original
@@ -26,7 +27,6 @@ import java.util.Optional;
  */
 public class ResealTool implements Tool {
 
-    static final String PRIVATE_KEY_FILE_OPTION     = "private-key-file";
     static final String EXPECTED_KEY_ID_OPTION      = "expected-key-id";
     static final String RECIPIENT_KEY_ID_OPTION     = "key-id";
     static final String RECIPIENT_PUBLIC_KEY_OPTION = "recipient-public-key";
@@ -38,6 +38,13 @@ public class ResealTool implements Tool {
                     .required(false)
                     .desc("Private key file in Base58 encoded format")
                     .build(),
+            Option.builder("d")
+                    .longOpt(PRIVATE_KEY_DIR_OPTION)
+                    .hasArg(true)
+                    .required(false)
+                    .desc("Private key file directory used for automatically looking up " +
+                          "private keys based on the key ID specified as part of a token.")
+                    .build(),
             Option.builder("e")
                     .longOpt(EXPECTED_KEY_ID_OPTION)
                     .hasArg(true)
@@ -90,8 +97,7 @@ public class ResealTool implements Tool {
 
             var recipientPubKey = KeyUtils.fromBase58EncodedX25519PublicKey(CliUtils.optionOrThrow(arguments, RECIPIENT_PUBLIC_KEY_OPTION).strip());
             var recipientKeyId  = KeyId.ofString(CliUtils.optionOrThrow(arguments, RECIPIENT_KEY_ID_OPTION));
-            var privKeyPath     = Paths.get(CliUtils.optionOrThrow(arguments, PRIVATE_KEY_FILE_OPTION));
-            var privateKey      = KeyUtils.fromBase58EncodedX25519PrivateKey(Files.readString(privKeyPath).strip());
+            var privateKey      = ToolUtils.resolvePrivateKeyFromInvocation(invocation, sealedSharedKey.keyId());
             var secretShared    = SharedKeyGenerator.fromSealedKey(sealedSharedKey, privateKey);
             var resealedShared  = SharedKeyGenerator.reseal(secretShared, recipientPubKey, recipientKeyId);
 
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ToolUtils.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ToolUtils.java
index 32e9c6679f6..2a25832708c 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ToolUtils.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/ToolUtils.java
@@ -2,15 +2,31 @@
 package com.yahoo.vespa.security.tool.crypto;
 
 import com.yahoo.security.KeyId;
+import com.yahoo.security.KeyUtils;
 import com.yahoo.security.SealedSharedKey;
+import com.yahoo.vespa.security.tool.CliUtils;
+import com.yahoo.vespa.security.tool.ToolInvocation;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.PosixFilePermission;
+import java.security.interfaces.XECPrivateKey;
 import java.util.Optional;
+import java.util.regex.Pattern;
 
 /**
  * @author vekterli
  */
 public class ToolUtils {
 
+    static final String PRIVATE_KEY_FILE_OPTION = "private-key-file";
+    static final String PRIVATE_KEY_DIR_OPTION  = "private-key-dir";
+    static final String PRIVATE_KEY_DIR_ENV_VAR = "VESPA_CRYPTO_CLI_PRIVATE_KEY_DIR";
+
+    static final Pattern SAFE_KEY_ID_PATTERN = Pattern.compile("^[a-zA-Z0-9_-]+$");
+
     static void verifyExpectedKeyId(SealedSharedKey sealedSharedKey, Optional<String> maybeKeyId) {
         if (maybeKeyId.isPresent()) {
             var myKeyId = KeyId.ofString(maybeKeyId.get());
@@ -22,4 +38,65 @@ public class ToolUtils {
         }
     }
 
+    private static void verifyKeyIdIsPathSafe(KeyId keyId) {
+        String keyIdStr = keyId.asString();
+        if (!SAFE_KEY_ID_PATTERN.matcher(keyIdStr).matches()) {
+            throw new IllegalArgumentException("The token key ID is not comprised of path-safe characters; refusing " +
+                                               "to auto-deduce key file name");
+        }
+    }
+
+    private static void verifyPrivateKeyFileNotWorldReadable(Path keyPath) throws IOException {
+        var privKeyPerms = Files.getPosixFilePermissions(keyPath);
+        if (privKeyPerms.contains(PosixFilePermission.OTHERS_READ)) {
+            throw new IllegalArgumentException("Private key file '%s' is insecurely world-readable; refusing to read it"
+                                               .formatted(keyPath.toAbsolutePath()));
+        }
+    }
+
+    private static XECPrivateKey attemptResolvePrivateKeyFromDir(Path privKeyDirPath, KeyId tokenKeyId) throws IOException {
+        if (!Files.isDirectory(privKeyDirPath)) {
+            throw new IllegalArgumentException("'%s' is not a valid directory".formatted(privKeyDirPath.toAbsolutePath()));
+        }
+        verifyKeyIdIsPathSafe(tokenKeyId);
+        var keyPath = privKeyDirPath.resolve(tokenKeyId.asString() + ".key");
+        if (!Files.exists(keyPath)) {
+            // We've verified the key ID contents, so we know it's safe to print here
+            throw new IllegalArgumentException("Could not find a private key file matching token key ID '%s'"
+                                               .formatted(tokenKeyId.asString()));
+        }
+        verifyPrivateKeyFileNotWorldReadable(keyPath);
+        return KeyUtils.fromBase58EncodedX25519PrivateKey(Files.readString(keyPath).strip());
+    }
+
+    public static XECPrivateKey resolvePrivateKeyFromInvocation(ToolInvocation invocation, KeyId tokenKeyId) throws IOException {
+        var arguments = invocation.arguments();
+        var envVars   = invocation.envVars();
+
+        if (arguments.hasOption(PRIVATE_KEY_FILE_OPTION)) {
+            if (arguments.hasOption(PRIVATE_KEY_DIR_OPTION)) {
+                throw new IllegalArgumentException("--%s and --%s cannot be specified at the same time"
+                                                   .formatted(PRIVATE_KEY_FILE_OPTION, PRIVATE_KEY_DIR_OPTION));
+            }
+            var privKeyFilePath = Paths.get(arguments.getOptionValue(PRIVATE_KEY_FILE_OPTION));
+            invocation.printIfDebug(() -> "Using private key file '%s'".formatted(privKeyFilePath));
+            if (!Files.exists(privKeyFilePath)) {
+                throw new IllegalArgumentException("Specified private key file '%s' does not exist"
+                                                   .formatted(privKeyFilePath.toAbsolutePath()));
+            }
+            verifyPrivateKeyFileNotWorldReadable(privKeyFilePath);
+            return KeyUtils.fromBase58EncodedX25519PrivateKey(Files.readString(privKeyFilePath).strip());
+        } else if (arguments.hasOption(PRIVATE_KEY_DIR_OPTION) || envVars.containsKey(PRIVATE_KEY_DIR_ENV_VAR)) {
+            // Explicitly provided command line directory is preferred over env var, if set
+            var privKeyDirPath = Paths.get(arguments.hasOption(PRIVATE_KEY_DIR_OPTION)
+                                           ? arguments.getOptionValue(PRIVATE_KEY_DIR_OPTION)
+                                           : envVars.get(PRIVATE_KEY_DIR_ENV_VAR));
+            invocation.printIfDebug(() -> "Using private key lookup directory '%s'".formatted(privKeyDirPath));
+            return attemptResolvePrivateKeyFromDir(privKeyDirPath, tokenKeyId);
+        } else {
+            throw new IllegalArgumentException("No private key specified. Must specify either --%s or --%s"
+                                               .formatted(PRIVATE_KEY_FILE_OPTION, PRIVATE_KEY_DIR_OPTION));
+        }
+    }
+
 }