Redacting credentials from Get Workflows API

Signed-off-by: Joshua Palis <[email protected]>

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);

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<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<>();
+                    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

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");

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));
+    }
 }