add helper to find preferredData by regex and add to Analyzer (#594)

NationalSecurityAgency · Oct 2, 2023 · f580cd9 · f580cd9
1 parent ef13f6a
commit f580cd9
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 14 deletions.
diff --git a/src/main/java/emissary/analyze/Analyzer.java b/src/main/java/emissary/analyze/Analyzer.java
@@ -1,19 +1,25 @@
 package emissary.analyze;
 
 import emissary.core.IBaseDataObject;
+import emissary.core.IBaseDataObjectHelper;
 import emissary.place.ServiceProviderPlace;
 
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.List;
-import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * Base class for analyzers
  */
 public abstract class Analyzer extends ServiceProviderPlace {
-    protected List<String> PREFERRED_VIEWS;
-    protected String ANALYZED_DATA_NAME;
+    protected List<String> preferredViews;
+    protected String analyzedDataName;
+    protected List<Pattern> preferredViewPatternList;
+
+    protected boolean findPreferredViewByRegex;
+
 
     /**
      * Create the place
@@ -58,8 +64,12 @@ protected Analyzer() throws IOException {
      * </ul>
      */
     protected void configureAnalyzer() {
-        this.PREFERRED_VIEWS = configG.findEntries("PREFERRED_VIEW", "TEXT");
-        this.ANALYZED_DATA_NAME = configG.findStringEntry("ANALYZED_DATA_NAME", null);
+        this.preferredViews = configG.findEntries("PREFERRED_VIEW", "TEXT");
+        this.findPreferredViewByRegex = configG.findBooleanEntry("FIND_PREFERRED_VIEW_BY_REGEX", false);
+        if (findPreferredViewByRegex) {
+            preferredViewPatternList = preferredViews.stream().map(Pattern::compile).collect(Collectors.toList());
+        }
+        this.analyzedDataName = configG.findStringEntry("ANALYZED_DATA_NAME", null);
     }
 
     /**
@@ -68,14 +78,12 @@ protected void configureAnalyzer() {
      * @param payload the payload to pull data from
      */
     protected byte[] getPreferredData(final IBaseDataObject payload) {
-        final Set<String> altViewNames = payload.getAlternateViewNames();
 
-        for (final String view : this.PREFERRED_VIEWS) {
-            if (altViewNames.contains(view)) {
-                return payload.getAlternateView(view);
-            }
+        if (findPreferredViewByRegex) {
+            return IBaseDataObjectHelper.findPreferredDataByRegex(payload, preferredViewPatternList);
+        } else {
+            return IBaseDataObjectHelper.findPreferredData(payload, preferredViews);
         }
-        return payload.data();
     }
 
     /**
@@ -86,11 +94,11 @@ protected byte[] getPreferredData(final IBaseDataObject payload) {
      * @return true if the data was stored, false if not. See ANALYZED_DATA_NAME config element
      */
     protected boolean setPreferredData(final IBaseDataObject payload, final byte[] analyzedData) {
-        if (this.ANALYZED_DATA_NAME != null) {
-            if ("base".equals(this.ANALYZED_DATA_NAME)) {
+        if (this.analyzedDataName != null) {
+            if ("base".equals(this.analyzedDataName)) {
                 payload.setData(analyzedData);
             } else {
-                payload.addAlternateView(this.ANALYZED_DATA_NAME, analyzedData);
+                payload.addAlternateView(this.analyzedDataName, analyzedData);
             }
             return true;
         }

diff --git a/src/main/java/emissary/core/IBaseDataObjectHelper.java b/src/main/java/emissary/core/IBaseDataObjectHelper.java
@@ -12,10 +12,13 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.security.NoSuchAlgorithmException;
+import java.util.Collections;
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 /**
@@ -215,4 +218,40 @@ public static void addParentInformationToChildren(final IBaseDataObject parent,
             }
         }
     }
+
+    /**
+     * Search for the first preferred view by regular expression or use the primary data if none match
+     *
+     * @param payload the payload to pull data from
+     */
+    public static byte[] findPreferredDataByRegex(final IBaseDataObject payload, List<Pattern> preferredViewNamePatterns) {
+        return Optional.ofNullable(preferredViewNamePatterns).orElse(Collections.emptyList()).stream()
+                .map(preferredViewNamePattern -> findFirstAlternameViewNameByRegex(payload, preferredViewNamePattern))
+                .filter(Optional::isPresent)
+                .map(Optional::get)
+                .map(payload::getAlternateView)
+                .findFirst().orElse(payload.data());
+    }
+
+    private static Optional<String> findFirstAlternameViewNameByRegex(IBaseDataObject payload, Pattern preferredViewNamePattern) {
+        return payload.getAlternateViewNames().stream()
+                .filter(altViewName -> preferredViewNamePattern.matcher(altViewName).find())
+                .findFirst();
+    }
+
+    /**
+     * Search for the first preferred view that is present or use the primary data if none
+     *
+     * @param payload the payload to pull data from
+     */
+    public static byte[] findPreferredData(final IBaseDataObject payload, List<String> preferredViews) {
+        final Set<String> altViewNames = payload.getAlternateViewNames();
+
+        for (final String view : preferredViews) {
+            if (altViewNames.contains(view)) {
+                return payload.getAlternateView(view);
+            }
+        }
+        return payload.data();
+    }
 }
diff --git a/src/test/java/emissary/core/IBaseDataObjectHelperTest.java b/src/test/java/emissary/core/IBaseDataObjectHelperTest.java
@@ -22,6 +22,9 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -353,4 +356,37 @@ void testAddParentInformationToChildren() {
 
         assertTrue(true);
     }
+
+    @Test
+    void testFindPreferredDataByRegex() {
+        byte[] view1Bytes = "altview test bytes".getBytes(StandardCharsets.UTF_8);
+        ibdo1.addAlternateView("view1", view1Bytes);
+
+        ibdo1.addAlternateView("notpreferred", "test bytes".getBytes(StandardCharsets.UTF_8));
+
+        byte[] primaryViewBytes = "primary view test bytes".getBytes(StandardCharsets.UTF_8);
+        ibdo1.setData(primaryViewBytes);
+
+        byte[] preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, compileAsPatterns("view1", "view2", "view3"));
+        assertEquals(view1Bytes, preferredData);
+
+        preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, compileAsPatterns("view1"));
+        assertEquals(view1Bytes, preferredData, "Expected to find an alternate view and return alternate view bytes");
+
+        preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, compileAsPatterns("^view[12]"));
+        assertEquals(view1Bytes, preferredData, "Expected to find an alternate view by regex and return alternate view bytes");
+
+        preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, compileAsPatterns("view_1", "view2", "view3"));
+        assertEquals(primaryViewBytes, preferredData, "Expected to not find an alternate view and return the default primary view bytes");
+
+        preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, compileAsPatterns("^view[4]"));
+        assertEquals(primaryViewBytes, preferredData, "Expected to not find an alternate view by regex and return the default primary view bytes");
+
+        preferredData = IBaseDataObjectHelper.findPreferredDataByRegex(ibdo1, null);
+        assertEquals(primaryViewBytes, preferredData, "Expected a null pattern to return the default primary view bytes");
+    }
+
+    static List<Pattern> compileAsPatterns(String... values) {
+        return Stream.of(values).map(Pattern::compile).collect(Collectors.toList());
+    }
 }