Adding final template encryption decryption test

Signed-off-by: Joshua Palis <[email protected]>
opensearch-project · Nov 27, 2023 · 5a72c4d · 5a72c4d
1 parent e320b7d
commit 5a72c4d
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 9 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -128,22 +128,22 @@ private Template processTemplateCredentials(Template template, Function<String,
                 if (node.userInputs().containsKey(CREDENTIAL_FIELD)) {
                     // Apply the cipher funcion on all values within credential field
                     @SuppressWarnings("unchecked")
-                    Map<String, String> credentials = (Map<String, String>) node.userInputs().get(CREDENTIAL_FIELD);
+                    Map<String, String> credentials = new HashMap<>((Map<String, String>) node.userInputs().get(CREDENTIAL_FIELD));
                     credentials.replaceAll((key, cred) -> cipherFunction.apply(cred));
 
                     // Replace credentials field in node user inputs
-                    Map<String, Object> encryptedUserInputs = new HashMap<>();
-                    encryptedUserInputs.putAll(node.userInputs());
-                    encryptedUserInputs.replace(CREDENTIAL_FIELD, credentials);
+                    Map<String, Object> processedUserInputs = new HashMap<>();
+                    processedUserInputs.putAll(node.userInputs());
+                    processedUserInputs.replace(CREDENTIAL_FIELD, credentials);
 
                     // build new node to add to processed nodes
-                    WorkflowNode encryptedWorkflowNode = new WorkflowNode(
+                    WorkflowNode processedWorkflowNode = new WorkflowNode(
                         node.id(),
                         node.type(),
                         node.previousNodeInputs(),
-                        encryptedUserInputs
+                        processedUserInputs
                     );
-                    processedNodes.add(encryptedWorkflowNode);
+                    processedNodes.add(processedWorkflowNode);
                 } else {
                     processedNodes.add(node);
                 }

diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
@@ -9,6 +9,7 @@
 package org.opensearch.flowframework.util;
 
 import com.google.common.collect.ImmutableMap;
+import org.opensearch.Version;
 import org.opensearch.action.get.GetRequest;
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.client.Client;
@@ -18,10 +19,18 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.TestHelpers;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.model.Template;
+import org.opensearch.flowframework.model.Workflow;
+import org.opensearch.flowframework.model.WorkflowNode;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 
+import java.util.List;
+import java.util.Map;
+
+import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.MASTER_KEY;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -35,6 +44,9 @@ public class EncryptorUtilsTests extends OpenSearchTestCase {
     private Client client;
     private EncryptorUtils encryptorUtils;
     private String testMasterKey;
+    private Template testTemplate;
+    private String testCredentialKey;
+    private String testCredentialValue;
 
     @Override
     public void setUp() throws Exception {
@@ -43,6 +55,30 @@ public void setUp() throws Exception {
         this.client = mock(Client.class);
         this.encryptorUtils = new EncryptorUtils(clusterService, client);
         this.testMasterKey = encryptorUtils.generateMasterKey();
+        this.testCredentialKey = "credential_key";
+        this.testCredentialValue = "12345";
+
+        Version templateVersion = Version.fromString("1.0.0");
+        List<Version> compatibilityVersions = List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0"));
+        WorkflowNode nodeA = new WorkflowNode(
+            "A",
+            "a-type",
+            Map.of(),
+            Map.of(CREDENTIAL_FIELD, Map.of(testCredentialKey, testCredentialValue))
+        );
+        List<WorkflowNode> nodes = List.of(nodeA);
+        Workflow workflow = new Workflow(Map.of("key", "value"), nodes, List.of());
+
+        this.testTemplate = new Template(
+            "test",
+            "description",
+            "use case",
+            templateVersion,
+            compatibilityVersions,
+            Map.of("provision", workflow),
+            Map.of(),
+            TestHelpers.randomUser()
+        );
 
         ClusterState clusterState = mock(ClusterState.class);
         Metadata metadata = mock(Metadata.class);
@@ -123,7 +159,35 @@ public void testInitializeMasterKeyFailure() {
         assertEquals("Encryption key has not been initialized", ex.getMessage());
     }
 
-    // TODO : test encrypting test template
-    // TODO : test decrypting test template
+    public void testEncryptDecryptTemplateCredential() {
+        encryptorUtils.setMasterKey(testMasterKey);
+
+        // Ecnrypt template with credential field
+        Template processedtemplate = encryptorUtils.encryptTemplateCredentials(testTemplate);
+
+        // Validate the encrytped field
+        WorkflowNode node = processedtemplate.workflows().get("provision").nodes().get(0);
+
+        @SuppressWarnings("unchecked")
+        Map<String, String> encryptedCredentialMap = (Map<String, String>) node.userInputs().get(CREDENTIAL_FIELD);
+        assertEquals(1, encryptedCredentialMap.size());
 
+        String encryptedCredential = encryptedCredentialMap.get(testCredentialKey);
+        assertNotNull(encryptedCredential);
+        assertNotEquals(testCredentialValue, encryptedCredential);
+
+        // Decrypt credential field
+        processedtemplate = encryptorUtils.decryptTemplateCredentials(processedtemplate);
+
+        // Validate the decrypted field
+        node = processedtemplate.workflows().get("provision").nodes().get(0);
+
+        @SuppressWarnings("unchecked")
+        Map<String, String> decryptedCredentialMap = (Map<String, String>) node.userInputs().get(CREDENTIAL_FIELD);
+        assertEquals(1, decryptedCredentialMap.size());
+
+        String decryptedCredential = decryptedCredentialMap.get(testCredentialKey);
+        assertNotNull(decryptedCredential);
+        assertEquals(testCredentialValue, decryptedCredential);
+    }
 }