Redacting credentials from Get Workflows and Search Workflows API (#504)

* Redacting credentials from Get Workflows API Signed-off-by: Joshua Palis <[email protected]> * addressing PR comments Signed-off-by: Joshua Palis <[email protected]> * Applying credential redaction to search workflows API Signed-off-by: Joshua Palis <[email protected]> * Only appying filter script on template index Signed-off-by: Joshua Palis <[email protected]> * Fixing search integ test Signed-off-by: Joshua Palis <[email protected]> --------- Signed-off-by: Joshua Palis <[email protected]>
opensearch-project · Feb 9, 2024 · f78f91b · f78f91b
1 parent b7d3737
commit f78f91b
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 18 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java
@@ -23,13 +23,17 @@
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
+import org.opensearch.script.Script;
+import org.opensearch.script.ScriptType;
 import org.opensearch.search.builder.SearchSourceBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS;
+import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
 import static org.opensearch.flowframework.util.RestHandlerUtils.getSourceContext;
 
@@ -88,6 +92,20 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         searchSourceBuilder.fetchSource(getSourceContext(request, searchSourceBuilder));
         searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true);
         searchSourceBuilder.timeout(flowFrameworkSettings.getRequestTimeout());
+
+        // Apply credential filter when searching templates
+        if (index.equals(GLOBAL_CONTEXT_INDEX)) {
+            searchSourceBuilder.scriptField(
+                "filter",
+                new Script(
+                    ScriptType.INLINE,
+                    "painless",
+                    "def filteredSource = new HashMap(params._source); def workflows = filteredSource.get(\"workflows\"); if (workflows != null) { def provision = workflows.get(\"provision\"); if (provision != null) { def nodes = provision.get(\"nodes\"); if (nodes != null) { for (node in nodes) { def userInputs = node.get(\"user_inputs\"); if (userInputs != null) { userInputs.remove(\"credential\"); } } } } } return filteredSource;",
+                    Collections.emptyMap()
+                )
+            );
+        }
+
         SearchRequest searchRequest = new SearchRequest().source(searchSourceBuilder).indices(index);
         return channel -> client.execute(actionType, searchRequest, search(channel));
     }

diff --git 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,24 +36,28 @@ public class GetWorkflowTransportAction extends HandledTransportAction<WorkflowR
     private final Logger logger = LogManager.getLogger(GetWorkflowTransportAction.class);
     private final FlowFrameworkIndicesHandler flowFrameworkIndicesHandler;
     private final Client client;
+    private final EncryptorUtils encryptorUtils;
 
     /**
      * Instantiates a new GetWorkflowTransportAction instance
      * @param transportService the transport service
      * @param actionFilters action filters
      * @param flowFrameworkIndicesHandler The Flow Framework indices handler
+     * @param encryptorUtils Encryptor utils
      * @param client the Opensearch Client
      */
     @Inject
     public GetWorkflowTransportAction(
         TransportService transportService,
         ActionFilters actionFilters,
         FlowFrameworkIndicesHandler flowFrameworkIndicesHandler,
-        Client client
+        Client client,
+        EncryptorUtils encryptorUtils
     ) {
         super(GetWorkflowAction.NAME, transportService, actionFilters, WorkflowRequest::new);
         this.flowFrameworkIndicesHandler = flowFrameworkIndicesHandler;
         this.client = client;
+        this.encryptorUtils = encryptorUtils;
     }
 
     @Override
@@ -75,7 +80,9 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<GetW
                             )
                         );
                     } else {
-                        listener.onResponse(new GetWorkflowResponse(Template.parse(response.getSourceAsString())));
+                        // Remove any credential from response
+                        Template template = encryptorUtils.redactTemplateCredentials(Template.parse(response.getSourceAsString()));
+                        listener.onResponse(new GetWorkflowResponse(template));
                     }
                 }, exception -> {
                     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
@@ -210,6 +210,50 @@ String decrypt(final String encryptedCredential) {
         return new String(decryptedResult.getResult(), StandardCharsets.UTF_8);
     }
 
+    // TODO : Improve redactTemplateCredentials to redact different fields
+    /**
+     * Removes the credential fields from a template
+     * @param template the template
+     * @return the redacted template
+     */
+    public Template redactTemplateCredentials(Template template) {
+        Template.Builder redactedTemplateBuilder = new Template.Builder();
+
+        Map<String, Workflow> processedWorkflows = new HashMap<>();
+        for (Map.Entry<String, Workflow> entry : template.workflows().entrySet()) {
+
+            List<WorkflowNode> processedNodes = new ArrayList<>();
+            for (WorkflowNode node : entry.getValue().nodes()) {
+                if (node.userInputs().containsKey(CREDENTIAL_FIELD)) {
+
+                    // Remove credential field field in node user inputs
+                    Map<String, Object> processedUserInputs = new HashMap<>(node.userInputs());
+                    processedUserInputs.remove(CREDENTIAL_FIELD);
+
+                    // build new node to add to processed nodes
+                    processedNodes.add(new WorkflowNode(node.id(), node.type(), node.previousNodeInputs(), processedUserInputs));
+                } 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 = redactedTemplateBuilder.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/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -8,6 +8,8 @@
  */
 package org.opensearch.flowframework.rest;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
@@ -26,7 +28,6 @@
 import org.junit.ComparisonFailure;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -36,11 +37,11 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
-import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 
 public class FlowFrameworkRestApiIT extends FlowFrameworkRestTestCase {
+    private static final Logger logger = LogManager.getLogger(FlowFrameworkRestApiIT.class);
 
     private static AtomicBoolean waitToStart = new AtomicBoolean(true);
 
@@ -69,19 +70,6 @@ public void testSearchWorkflows() throws Exception {
         String termIdQuery = "{\"query\":{\"ids\":{\"values\":[\"" + workflowId + "\"]}}}";
         SearchResponse searchResponse = searchWorkflows(client(), termIdQuery);
         assertEquals(1, searchResponse.getHits().getTotalHits().value);
-
-        String searchHitSource = searchResponse.getHits().getAt(0).getSourceAsString();
-        Template searchHitTemplate = Template.parse(searchHitSource);
-
-        // Confirm that credentials have been encrypted within the search response
-        List<WorkflowNode> provisionNodes = searchHitTemplate.workflows().get(PROVISION_WORKFLOW).nodes();
-        for (WorkflowNode node : provisionNodes) {
-            if (node.type().equals("create_connector")) {
-                @SuppressWarnings("unchecked")
-                Map<String, String> credentialMap = new HashMap<>((Map<String, String>) node.userInputs().get(CREDENTIAL_FIELD));
-                assertTrue(credentialMap.values().stream().allMatch(x -> x != "12345"));
-            }
-        }
     }
 
     public void testFailedUpdateWorkflow() throws Exception {

diff --git 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,18 +55,21 @@ public class GetWorkflowTransportActionTests extends OpenSearchTestCase {
     private GetWorkflowTransportAction getTemplateTransportAction;
     private FlowFrameworkIndicesHandler flowFrameworkIndicesHandler;
     private Template template;
+    private EncryptorUtils encryptorUtils;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
         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
@@ -190,4 +190,17 @@ public void testEncryptDecryptTemplateCredential() {
         assertNotNull(decryptedCredential);
         assertEquals(testCredentialValue, decryptedCredential);
     }
+
+    public void testRedactTemplateCredential() {
+        // Confirm credentials are present in the non-redacted template
+        WorkflowNode node = testTemplate.workflows().get("provision").nodes().get(0);
+        assertNotNull(node.userInputs().get(CREDENTIAL_FIELD));
+
+        // Redact template with credential field
+        Template redactedTemplate = encryptorUtils.redactTemplateCredentials(testTemplate);
+
+        // Validate the credential field has been removed
+        WorkflowNode redactedNode = redactedTemplate.workflows().get("provision").nodes().get(0);
+        assertNull(redactedNode.userInputs().get(CREDENTIAL_FIELD));
+    }
 }