Code/Projects/PacketLevelSignatureExtractor/src/main/java/edu/uci/iotproject/trafficreassembly/layer2/Layer2FlowReassembler.java

   1 package edu.uci.iotproject.trafficreassembly.layer2;
   2
   3 import edu.uci.iotproject.trafficreassembly.layer2.Layer2Flow;
   4 import edu.uci.iotproject.trafficreassembly.layer2.Layer2FlowReassemblerObserver;
   5 import org.pcap4j.core.PacketListener;
   6 import org.pcap4j.core.PcapPacket;
   7 import org.pcap4j.packet.EthernetPacket;
   8 import org.pcap4j.util.MacAddress;
   9
  10 import java.util.*;
  11
  12 /**
  13  * Reassembles traffic flows at layer 2, i.e., for each combination of hosts, creates a list of packets exchanged
  14  * between said hosts.
  15  *
  16  * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
  17  * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
  18  */
  19 public class Layer2FlowReassembler implements PacketListener {
  20
  21     /**
  22      * Maps a pair of MAC addresses to the packets exchanged between the two hosts.
  23      * The key is the concatenation of the two MAC addresses in hex string format, where the lexicographically smaller
  24      * MAC is at the front of the string.
  25      */
  26     private final Map<String, Layer2Flow> mFlows = new HashMap<>();
  27
  28     private final List<Layer2FlowReassemblerObserver> mObservers = new ArrayList<>();
  29
  30     private String mVpnClientMacAddress = null;
  31
  32     public Layer2FlowReassembler() { }
  33
  34     public Layer2FlowReassembler(String vpnClientMacAddress) {
  35         mVpnClientMacAddress = vpnClientMacAddress;
  36     }
  37
  38     @Override
  39     public void gotPacket(PcapPacket packet) {
  40         // TODO: update to 802.11 packet...?
  41         EthernetPacket ethPkt = packet.get(EthernetPacket.class);
  42
  43         MacAddress srcAddr = ethPkt.getHeader().getSrcAddr();
  44         MacAddress dstAddr = ethPkt.getHeader().getDstAddr();
  45
  46         String key = null;
  47         if (mVpnClientMacAddress != null) {
  48             if (srcAddr.toString().equals(mVpnClientMacAddress)) {
  49                 key = srcAddr.toString();
  50             } else if (dstAddr.toString().equals(mVpnClientMacAddress)) {
  51                 key = dstAddr.toString();
  52             } else {
  53                 return;
  54             }
  55         } else {
  56             key = keyFromAddresses(srcAddr, dstAddr);
  57         }
  58         // Create a new list if this pair of MAC addresses where not previously encountered and add packet to that list,
  59         // or simply add to an existing list if one is present.
  60         mFlows.computeIfAbsent(key, k -> {
  61             Layer2Flow newFlow = new Layer2Flow(srcAddr, dstAddr);
  62             // Inform observers of the new flow
  63             mObservers.forEach(o -> o.onNewFlow(this, newFlow));
  64             return newFlow;
  65         }).addPacket(packet);
  66     }
  67
  68     public void addObserver(Layer2FlowReassemblerObserver observer) {
  69         mObservers.add(observer);
  70     }
  71
  72     public void removeObserver(Layer2FlowReassemblerObserver observer) {
  73         mObservers.remove(observer);
  74     }
  75
  76     /**
  77      * Get the traffic flow between two local endpoints ({@link MacAddress}es).
  78      * @param addr1 The first endpoint.
  79      * @param addr2 The second endpoint
  80      * @return The traffic exchanged between the two endpoints.
  81      */
  82     public Layer2Flow getFlowForAddresses(MacAddress addr1, MacAddress addr2) {
  83         return mFlows.get(keyFromAddresses(addr1, addr2));
  84     }
  85
  86     /**
  87      * Get all traffic flows, i.e., a traffic flow for each unique pair of endpoints (MAC addresses).
  88      * @return All traffic flows.
  89      */
  90     public Collection<Layer2Flow> getFlows() {
  91         return mFlows.values();
  92     }
  93
  94     /**
  95      * Given two {@link MacAddress}es, generates the corresponding key string used in {@link #mFlows}.
  96      * @param addr1 The first address.
  97      * @param addr2 The second address.
  98      * @return the key string used in {@link #mFlows} corresponding to the two addresses.
  99      */
 100     private String keyFromAddresses(MacAddress addr1, MacAddress addr2) {
 101         String addr1Str = addr1.toString();
 102         String addr2Str = addr2.toString();
 103         return addr1Str.compareTo(addr2Str) < 0 ? addr1Str + addr2Str : addr2Str + addr1Str;
 104     }
 105 }