From 5a72c4d30fed757f4c5e9b6733a8363fae553951 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 27 Nov 2023 21:48:50 +0000 Subject: [PATCH] Adding final template encryption decryption test Signed-off-by: Joshua Palis --- .../flowframework/util/EncryptorUtils.java | 14 ++-- .../util/EncryptorUtilsTests.java | 68 ++++++++++++++++++- 2 files changed, 73 insertions(+), 9 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index 4ef815431..f954c504e 100644 --- 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 credentials = (Map) node.userInputs().get(CREDENTIAL_FIELD); + Map credentials = new HashMap<>((Map) node.userInputs().get(CREDENTIAL_FIELD)); credentials.replaceAll((key, cred) -> cipherFunction.apply(cred)); // Replace credentials field in node user inputs - Map encryptedUserInputs = new HashMap<>(); - encryptedUserInputs.putAll(node.userInputs()); - encryptedUserInputs.replace(CREDENTIAL_FIELD, credentials); + Map 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 index 00afb7514..c044783b8 100644 --- 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 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 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 encryptedCredentialMap = (Map) 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 decryptedCredentialMap = (Map) node.userInputs().get(CREDENTIAL_FIELD); + assertEquals(1, decryptedCredentialMap.size()); + + String decryptedCredential = decryptedCredentialMap.get(testCredentialKey); + assertNotNull(decryptedCredential); + assertEquals(testCredentialValue, decryptedCredential); + } }