From 57ab33f9a646671e4f93d845d206713279280534 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 8 Feb 2024 21:13:38 +0000 Subject: [PATCH] Redacting credentials from Get Workflows API Signed-off-by: Joshua Palis --- .../transport/GetWorkflowTransportAction.java | 11 +++- .../flowframework/util/EncryptorUtils.java | 50 +++++++++++++++++++ .../GetWorkflowTransportActionTests.java | 7 ++- .../util/EncryptorUtilsTests.java | 9 ++++ 4 files changed, 74 insertions(+), 3 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java index e2a9b1931..8bb9c0b89 100644 --- a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java @@ -22,6 +22,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler; import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.util.EncryptorUtils; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -35,12 +36,14 @@ public class GetWorkflowTransportAction extends HandledTransportAction { logger.error("Failed to retrieve template from global context.", exception); diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index 890eaaf95..d99299b29 100644 --- a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java +++ b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java @@ -210,6 +210,56 @@ String decrypt(final String encryptedCredential) { return new String(decryptedResult.getResult(), StandardCharsets.UTF_8); } + /** + * Removes the credential fields from a template + * @param template the template + * @return the redacted template + */ + public Template redactTemplateCredentials(Template template) { + Template.Builder processedTemplateBuilder = new Template.Builder(); + + Map processedWorkflows = new HashMap<>(); + for (Map.Entry entry : template.workflows().entrySet()) { + + List processedNodes = new ArrayList<>(); + for (WorkflowNode node : entry.getValue().nodes()) { + if (node.userInputs().containsKey(CREDENTIAL_FIELD)) { + + // Remove credential field field in node user inputs + Map processedUserInputs = new HashMap<>(); + processedUserInputs.putAll(node.userInputs()); + processedUserInputs.remove(CREDENTIAL_FIELD); + + // build new node to add to processed nodes + WorkflowNode processedWorkflowNode = new WorkflowNode( + node.id(), + node.type(), + node.previousNodeInputs(), + processedUserInputs + ); + processedNodes.add(processedWorkflowNode); + } else { + processedNodes.add(node); + } + } + + // Add processed workflow nodes to processed workflows + processedWorkflows.put(entry.getKey(), new Workflow(entry.getValue().userParams(), processedNodes, entry.getValue().edges())); + } + + Template processedTemplate = processedTemplateBuilder.name(template.name()) + .description(template.description()) + .useCase(template.useCase()) + .templateVersion(template.templateVersion()) + .compatibilityVersion(template.compatibilityVersion()) + .workflows(processedWorkflows) + .uiMetadata(template.getUiMetadata()) + .user(template.getUser()) + .build(); + + return processedTemplate; + } + /** * Retrieves an existing master key or generates a new key to index * @param listener the action listener diff --git a/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java index 38ee192df..7b34c24f1 100644 --- a/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/GetWorkflowTransportActionTests.java @@ -13,6 +13,7 @@ import org.opensearch.action.get.GetResponse; import org.opensearch.action.support.ActionFilters; import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.xcontent.XContentFactory; @@ -25,6 +26,7 @@ import org.opensearch.flowframework.model.Workflow; import org.opensearch.flowframework.model.WorkflowEdge; import org.opensearch.flowframework.model.WorkflowNode; +import org.opensearch.flowframework.util.EncryptorUtils; import org.opensearch.index.get.GetResult; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; @@ -53,6 +55,7 @@ public class GetWorkflowTransportActionTests extends OpenSearchTestCase { private GetWorkflowTransportAction getTemplateTransportAction; private FlowFrameworkIndicesHandler flowFrameworkIndicesHandler; private Template template; + private EncryptorUtils encryptorUtils; @Override public void setUp() throws Exception { @@ -60,11 +63,13 @@ public void setUp() throws Exception { this.threadPool = mock(ThreadPool.class); this.client = mock(Client.class); this.flowFrameworkIndicesHandler = mock(FlowFrameworkIndicesHandler.class); + this.encryptorUtils = new EncryptorUtils(mock(ClusterService.class), client); this.getTemplateTransportAction = new GetWorkflowTransportAction( mock(TransportService.class), mock(ActionFilters.class), flowFrameworkIndicesHandler, - client + client, + encryptorUtils ); Version templateVersion = Version.fromString("1.0.0"); diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java index 790070a32..42245645e 100644 --- a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java +++ b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java @@ -190,4 +190,13 @@ public void testEncryptDecryptTemplateCredential() { assertNotNull(decryptedCredential); assertEquals(testCredentialValue, decryptedCredential); } + + public void testRedactTemplateCredential() { + // Redact template with credential field + Template processedTemplate = encryptorUtils.redactTemplateCredentials(testTemplate); + + // Validate the credential field has been removed + WorkflowNode node = processedTemplate.workflows().get("provision").nodes().get(0); + assertNull(node.userInputs().get(CREDENTIAL_FIELD)); + } }