From 35f04e834903f3d7f6624e3f1244b4f66b5d08e3 Mon Sep 17 00:00:00 2001
From: Owais Kazi <owaiskazi19@gmail.com>
Date: Wed, 24 Apr 2024 19:44:23 -0700
Subject: [PATCH] Hid the user field for Get API as well

Signed-off-by: Owais Kazi <owaiskazi19@gmail.com>
---
 .../transport/GetWorkflowTransportAction.java |  2 +-
 .../flowframework/util/EncryptorUtils.java    | 28 ++++++++++++++++---
 .../util/EncryptorUtilsTests.java             | 13 ++++++++-
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java
index 11c135210..2eb8b2e39 100644
--- a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java
+++ b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java
@@ -80,7 +80,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<GetW
                         listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.NOT_FOUND));
                     } else {
                         // Remove any credential from response
-                        Template template = encryptorUtils.redactTemplateCredentials(Template.parse(response.getSourceAsString()));
+                        Template template = encryptorUtils.redactTemplateSecuredFields(Template.parse(response.getSourceAsString()));
                         listener.onResponse(new GetWorkflowResponse(template));
                     }
                 }, exception -> {
diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
index df7e66e07..551ab7c9b 100644
--- a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
+++ b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -122,7 +122,7 @@ public Template decryptTemplateCredentials(Template template) {
     /**
      * Applies the given cipher function on template credentials
      * @param template the template to process
-     * @param cipher the encryption/decryption function to apply on credential values
+     * @param cipherFunction the encryption/decryption function to apply on credential values
      * @return template with encrypted credentials
      */
     private Template processTemplateCredentials(Template template, Function<String, String> cipherFunction) {
@@ -204,9 +204,29 @@ String decrypt(final String encryptedCredential) {
      * @param template the template
      * @return the redacted template
      */
-    public Template redactTemplateCredentials(Template template) {
+    public Template redactTemplateSecuredFields(Template template) {
+        Template updatedTemplate = null;
+
+        if (template.getUser() != null) {
+            updatedTemplate = new Template.Builder(template).name(template.name())
+                .description(template.description())
+                .useCase(template.useCase())
+                .templateVersion(template.templateVersion())
+                .user(null)
+                .uiMetadata(template.getUiMetadata())
+                .compatibilityVersion(template.compatibilityVersion())
+                .workflows(template.workflows())
+                .createdTime(template.createdTime())
+                .lastUpdatedTime(template.lastUpdatedTime())
+                .lastProvisionedTime(template.lastProvisionedTime())
+                .build();
+        } else {
+            updatedTemplate = template;
+        }
+
         Map<String, Workflow> processedWorkflows = new HashMap<>();
-        for (Map.Entry<String, Workflow> entry : template.workflows().entrySet()) {
+
+        for (Map.Entry<String, Workflow> entry : updatedTemplate.workflows().entrySet()) {
 
             List<WorkflowNode> processedNodes = new ArrayList<>();
             for (WorkflowNode node : entry.getValue().nodes()) {
@@ -227,7 +247,7 @@ public Template redactTemplateCredentials(Template template) {
             processedWorkflows.put(entry.getKey(), new Workflow(entry.getValue().userParams(), processedNodes, entry.getValue().edges()));
         }
 
-        return new Template.Builder(template).workflows(processedWorkflows).build();
+        return new Template.Builder(updatedTemplate).workflows(processedWorkflows).build();
     }
 
     /**
diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
index cae595430..c6ec15a92 100644
--- a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
+++ b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
@@ -200,10 +200,21 @@ public void testRedactTemplateCredential() {
         assertNotNull(node.userInputs().get(CREDENTIAL_FIELD));
 
         // Redact template with credential field
-        Template redactedTemplate = encryptorUtils.redactTemplateCredentials(testTemplate);
+        Template redactedTemplate = encryptorUtils.redactTemplateSecuredFields(testTemplate);
 
         // Validate the credential field has been removed
         WorkflowNode redactedNode = redactedTemplate.workflows().get("provision").nodes().get(0);
         assertNull(redactedNode.userInputs().get(CREDENTIAL_FIELD));
     }
+
+    public void testRedactTemplateUserField() {
+        // Confirm user is present in the non-redacted template
+        assertNotNull(testTemplate.getUser());
+
+        // Redact template with user field
+        Template redactedTemplate = encryptorUtils.redactTemplateSecuredFields(testTemplate);
+
+        // Validate the user field has been removed
+        assertNull(redactedTemplate.getUser());
+    }
 }