From c5ae02722522de415fbcd41e8f133bb7c57d047e Mon Sep 17 00:00:00 2001
From: Michael Jumper <mjumper@apache.org>
Date: Fri, 21 Jan 2022 15:23:41 -0800
Subject: [PATCH] GUACAMOLE-641: Add user- and gateway-specific tokens.

---
 .../vault/ksm/secret/KsmSecretService.java    | 90 +++++++++++++------
 1 file changed, 65 insertions(+), 25 deletions(-)

diff --git a/extensions/guacamole-vault/modules/guacamole-vault-ksm/src/main/java/org/apache/guacamole/vault/ksm/secret/KsmSecretService.java b/extensions/guacamole-vault/modules/guacamole-vault-ksm/src/main/java/org/apache/guacamole/vault/ksm/secret/KsmSecretService.java
index 6fd9f203c..d7b4deb50 100644
--- a/extensions/guacamole-vault/modules/guacamole-vault-ksm/src/main/java/org/apache/guacamole/vault/ksm/secret/KsmSecretService.java
+++ b/extensions/guacamole-vault/modules/guacamole-vault-ksm/src/main/java/org/apache/guacamole/vault/ksm/secret/KsmSecretService.java
@@ -71,44 +71,84 @@ public class KsmSecretService implements VaultSecretService {
         return ksm.getSecret(name);
     }
 
+    /**
+     * Adds contextual parameter tokens for the secrets in the given record to
+     * the given map of existing tokens. The values of each token are
+     * determined from secrets within the record. Depending on the record, this
+     * will be a subset of the username, password, private key, and passphrase.
+     *
+     * @param tokens
+     *     The map of parameter tokens that any new tokens should be added to.
+     *
+     * @param prefix
+     *     The prefix that should be prepended to each added token.
+     *
+     * @param record
+     *     The record to retrieve secrets from when generating tokens. This may
+     *     be null.
+     */
+    private void addRecordTokens(Map<String, Future<String>> tokens, String prefix,
+            KeeperRecord record) {
+
+        if (record == null)
+            return;
+
+        // Username of server-related record
+        String username = recordService.getUsername(record);
+        if (username != null)
+            tokens.put(prefix + "USERNAME", CompletableFuture.completedFuture(username));
+
+        // Password of server-related record
+        String password = recordService.getPassword(record);
+        if (password != null)
+            tokens.put(prefix + "PASSWORD", CompletableFuture.completedFuture(password));
+
+        // Key passphrase of server-related record
+        String passphrase = recordService.getPassphrase(record);
+        if (passphrase != null)
+            tokens.put(prefix + "PASSPHRASE", CompletableFuture.completedFuture(passphrase));
+
+        // Private key of server-related record
+        String privateKey = recordService.getPrivateKey(record);
+        if (privateKey != null)
+            tokens.put(prefix + "KEY", CompletableFuture.completedFuture(privateKey));
+
+    }
+    
     @Override
     public Map<String, Future<String>> getTokens(GuacamoleConfiguration config,
             TokenFilter filter) throws GuacamoleException {
 
         Map<String, Future<String>> tokens = new HashMap<>();
-
-        // TODO: Verify protocol before assuming meaning of "hostname"
-        // parameter
-
         Map<String, String> parameters = config.getParameters();
 
         // Retrieve and define server-specific tokens, if any
         String hostname = parameters.get("hostname");
-        if (hostname != null && !hostname.isEmpty()) {
-            KeeperRecord record = ksm.getRecordByHost(filter.filter(hostname));
-            if (record != null) {
+        if (hostname != null && !hostname.isEmpty())
+            addRecordTokens(tokens, "KEEPER_SERVER_",
+                    ksm.getRecordByHost(filter.filter(hostname)));
 
-                // Username of server-related record
-                String username = recordService.getUsername(record);
-                if (username != null)
-                    tokens.put("KEEPER_SERVER_USERNAME", CompletableFuture.completedFuture(username));
+        // Retrieve and define user-specific tokens, if any
+        String username = parameters.get("username");
+        if (username != null && !username.isEmpty())
+            addRecordTokens(tokens, "KEEPER_USER_",
+                    ksm.getRecordByLogin(filter.filter(username)));
 
-                // Password of server-related record
-                String password = recordService.getPassword(record);
-                if (password != null)
-                    tokens.put("KEEPER_SERVER_PASSWORD", CompletableFuture.completedFuture(password));
+        // Tokens specific to RDP
+        if ("rdp".equals(config.getProtocol())) {
+        
+            // Retrieve and define gateway server-specific tokens, if any
+            String gatewayHostname = parameters.get("gateway-hostname");
+            if (gatewayHostname != null && !gatewayHostname.isEmpty())
+                addRecordTokens(tokens, "KEEPER_GATEWAY_",
+                        ksm.getRecordByHost(filter.filter(gatewayHostname)));
 
-                // Key passphrase of server-related record
-                String passphrase = recordService.getPassphrase(record);
-                if (passphrase != null)
-                    tokens.put("KEEPER_SERVER_PASSPHRASE", CompletableFuture.completedFuture(passphrase));
+            // Retrieve and define gateway user-specific tokens, if any
+            String gatewayUsername = parameters.get("gateway-username");
+            if (gatewayUsername != null && !gatewayUsername.isEmpty())
+                addRecordTokens(tokens, "KEEPER_GATEWAY_USER_",
+                        ksm.getRecordByLogin(filter.filter(gatewayUsername)));
 
-                // Private key of server-related record
-                String privateKey = recordService.getPrivateKey(record);
-                if (privateKey != null)
-                    tokens.put("KEEPER_SERVER_KEY", CompletableFuture.completedFuture(privateKey));
-
-            }
         }
 
         return tokens;