Add capability to filter flows when performing layer 2 detection

[pingpong.git] / Code / Projects / SmartPlugDetector / src / main / java / edu / uci / iotproject / detection / layer2 / Layer2ClusterMatcher.java

diff --git a/Code/Projects/SmartPlugDetector/src/main/java/edu/uci/iotproject/detection/layer2/Layer2ClusterMatcher.java b/Code/Projects/SmartPlugDetector/src/main/java/edu/uci/iotproject/detection/layer2/Layer2ClusterMatcher.java
index 79a945ffe790bf4a850d9471fe2e7d548b834732..5021c3119693bf84b5dc076e79cdf30e614ab815 100644 (file)
--- a/Code/Projects/SmartPlugDetector/src/main/java/edu/uci/iotproject/detection/layer2/Layer2ClusterMatcher.java
+++ b/Code/Projects/SmartPlugDetector/src/main/java/edu/uci/iotproject/detection/layer2/Layer2ClusterMatcher.java
@@ -11,6 +11,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * Attempts to detect members of a cluster (packet sequence mutations) in layer 2 flows.
@@ -27,9 +28,30 @@ public class Layer2ClusterMatcher extends AbstractClusterMatcher implements Laye
      */
     private final Map<Layer2Flow, Layer2SequenceMatcher[][]> mPerFlowSeqMatchers = new HashMap<>();
 
+    private final Function<Layer2Flow, Boolean> mFlowFilter;
 
+    /**
+     * Create a new {@link Layer2ClusterMatcher} that attempts to find occurrences of {@code cluster}'s members.
+     * @param cluster The sequence mutations that the new {@link Layer2ClusterMatcher} should search for.
+     */
     public Layer2ClusterMatcher(List<List<PcapPacket>> cluster) {
+        // Consider all flows if no flow filter specified.
+        this(cluster, flow -> true);
+    }
+
+    /**
+     * Create a new {@link Layer2ClusterMatcher} that attempts to find occurrences of {@code cluster}'s members.
+     * @param cluster The sequence mutations that the new {@link Layer2ClusterMatcher} should search for.
+     * @param flowFilter A filter that defines what {@link Layer2Flow}s the new {@link Layer2ClusterMatcher} should
+     *                   search for {@code cluster}'s members in. If {@code flowFilter} returns {@code true}, the flow
+     *                   will be included (searched). Note that {@code flowFilter} is only queried once for each flow,
+     *                   namely when the {@link Layer2FlowReassembler} notifies the {@link Layer2ClusterMatcher} about
+     *                   the new flow. This functionality may for example come in handy when one only wants to search
+     *                   for matches in the subset of flows that involves a specific (range of) MAC(s).
+     */
+    public Layer2ClusterMatcher(List<List<PcapPacket>> cluster, Function<Layer2Flow, Boolean> flowFilter) {
         super(cluster);
+        mFlowFilter = flowFilter;
     }
 
     @Override
@@ -113,10 +135,19 @@ public class Layer2ClusterMatcher extends AbstractClusterMatcher implements Laye
         return prunedCluster;
     }
 
+    private static final boolean DEBUG = false;
 
     @Override
     public void onNewFlow(Layer2FlowReassembler reassembler, Layer2Flow newFlow) {
-        // Subscribe to the new flow to get updates whenever a new packet pertaining to the flow is processed.
-        newFlow.addFlowObserver(this);
+        // New flow detected. Check if we should consider it when searching for cluster member matches.
+        if (mFlowFilter.apply(newFlow)) {
+            if (DEBUG) {
+                System.out.println(">>> ACCEPTING FLOW: " + newFlow + " <<<");
+            }
+            // Subscribe to the new flow to get updates whenever a new packet pertaining to the flow is processed.
+            newFlow.addFlowObserver(this);
+        } else if (DEBUG) {
+            System.out.println(">>> IGNORING FLOW: " + newFlow + " <<<");
+        }
     }
 }