Merge pull request #169 from cfkoehler/allow-single-child-inMultiFile…

…UnixCommandPlace Allow a single child and preserve parents data in MultiFileUnixCommandPlace
NationalSecurityAgency · Sep 29, 2021 · e579f64 · e579f64
2 parents c607eec + b2e983d
commit e579f64
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 8 deletions.
diff --git a/src/main/java/emissary/place/MultiFileUnixCommandPlace.java b/src/main/java/emissary/place/MultiFileUnixCommandPlace.java
@@ -44,6 +44,8 @@ public class MultiFileUnixCommandPlace extends MultiFileServerPlace implements I
     protected Executrix executrix;
     protected String logfilename;
     protected String charset = "UTF-8";
+    protected boolean singleOutputAsChild = false;
+    protected boolean preserveParentData = false;
 
     String placeDisplayName = "Some Place";
 
@@ -99,6 +101,8 @@ public MultiFileUnixCommandPlace(String configInfo, String loc) throws IOExcepti
      * <li>CUSTOM_FILE_TYPES: special mapping to set type by file extension</li>
      * <li>LOG_FILE_NAME: name of output file to translate into logger commands, default: [servicename].log from key</li>
      * <li>OUTPUT_CHARSET: charset of the process output, default UTF-8</li>
+     * <li>SINGLE_OUTPUT_AS_CHILD: If only one output file keep it as a child and do not replace the parent.</li>
+     * <li>PRESERVE_PARENT_DATA: Stops the parent from getting replaced by output data</li>
      * </ul>
      */
     @Override
@@ -130,6 +134,8 @@ public void configurePlace() {
         setTitleToFile = configG.findBooleanEntry("SET_TITLE_TO_FILENAME", true);
         placeDisplayName = configG.findStringEntry("SERVICE_DISPLAY_NAME", placeName);
         logfilename = configG.findStringEntry("LOG_FILE_NAME", KeyManipulator.getServiceName(keys.get(0)) + ".log");
+        singleOutputAsChild = configG.findBooleanEntry("SINGLE_OUTPUT_AS_CHILD", singleOutputAsChild);
+        preserveParentData = configG.findBooleanEntry("PRESERVE_PARENT_DATA", preserveParentData);
 
         for (String name : configG.findEntries("CUSTOM_FILE_TYPES")) {
             String tmp = configG.findStringEntry(name + "_EXT", null);
@@ -356,11 +362,13 @@ protected List<IBaseDataObject> sproutResults(IBaseDataObject parent, List<File>
 
             finishSprout(parent, fileCount, actualFileCount, newData);
 
-            try {
-                parent.setData(newData.toString().getBytes(charset));
-            } catch (UnsupportedEncodingException e) {
-                logger.debug("SproutResults charset problem", e);
-                parent.setData(newData.toString().getBytes());
+            if (!preserveParentData) {
+                try {
+                    parent.setData(newData.toString().getBytes(charset));
+                } catch (UnsupportedEncodingException e) {
+                    logger.debug("SproutResults charset problem", e);
+                    parent.setData(newData.toString().getBytes());
+                }
             }
         }
         return sprouts;
@@ -535,8 +543,8 @@ public List<IBaseDataObject> processHeavyDuty(IBaseDataObject tData) throws Reso
 
         addParentInformation(tData, entries);
 
-        // we'll replace parent with its one child when we can
-        if (executrix.getOutput().equals("FILE") && entries.size() == 1 && contentFile == null) {
+        // Replace parent if single child and singleOutputAsChild is false
+        if (executrix.getOutput().equals("FILE") && entries.size() == 1 && contentFile == null && !singleOutputAsChild) {
             IBaseDataObject d = entries.get(0);
             tData.setData(d.data());
             if (KEEP_PARENT_HASHES_FOR_SINGLE_CHILD) {

diff --git a/src/test/java/emissary/place/MultiFileUnixCommandPlaceTest.java b/src/test/java/emissary/place/MultiFileUnixCommandPlaceTest.java
@@ -38,6 +38,7 @@ public class MultiFileUnixCommandPlaceTest extends UnitTest {
     private static String W = "Президент Буш";
     private IBaseDataObject payload;
     private String FORM = "TEST";
+    private static final String PAYLOAD_STRING = "abcdefg";
 
     @Override
     @Before
@@ -66,7 +67,7 @@ public void setUp() throws Exception {
             }
         }
 
-        payload = DataObjectFactory.getInstance(new Object[] {"abcdefg".getBytes(), "myPayload", FORM});
+        payload = DataObjectFactory.getInstance(new Object[] {PAYLOAD_STRING.getBytes(), "myPayload", FORM});
 
         payload.putParameter("COPY_THIS", "copy value");
         payload.putParameter("IGNORE_THIS", "ignore value");
@@ -137,6 +138,21 @@ public void testMultiFileUnixCommandPlaceFileWithSingleChildHandling() throws Ex
         assertEquals("Parent should still have non-propagating metadata value", "ignore value", payload.getStringParameter("IGNORE_THIS"));
     }
 
+    @Test
+    public void testMultiFileUnixCommandPlaceFileWithSingleChildAsOutput() throws Exception {
+        assertNotNull("Place must be created", place);
+        place.setFileOutputCommand();
+        createScript(Executrix.OUTPUT_TYPE.FILE, 1);
+
+        place.preserveParentData = true;
+        place.singleOutputAsChild = true;
+
+        List<IBaseDataObject> att = place.processHeavyDuty(payload);
+        assertEquals("One Attachment should be created", 1, att.size());
+        assertEquals("Parent data should be preserved", PAYLOAD_STRING, new String(payload.data()));
+        assertEquals("Child payload should match script output", W, new String(att.get(0).data()).trim());
+    }
+
     @Test
     public void testMultiFileUnixCommandPlaceLogging() throws Exception {
         assertNotNull("Place must be created", place);