Code/Projects/PacketLevelSignatureExtractor/src/main/java/edu/uci/iotproject/analysis/TrafficLabeler.java

   1 package edu.uci.iotproject.analysis;
   2
   3 import edu.uci.iotproject.trafficreassembly.layer3.Conversation;
   4 import edu.uci.iotproject.DnsMap;
   5 import edu.uci.iotproject.trafficreassembly.layer3.TcpReassembler;
   6 import org.pcap4j.core.PacketListener;
   7 import org.pcap4j.core.PcapPacket;
   8
   9 import java.time.Instant;
  10 import java.util.*;
  11 import java.util.function.Function;
  12
  13 /**
  14  * A {@link PacketListener} that marks network traffic as (potentially) related to a user's actions by comparing the
  15  * timestamp of each packet to the timestamps of the provided list of user actions.
  16  *
  17  * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
  18  * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
  19  */
  20 public class TrafficLabeler implements PacketListener {
  21
  22     private final Map<UserAction, List<PcapPacket>> mActionToTrafficMap;
  23     private final List<UserAction> mActionsSorted;
  24     /**
  25      * The total number of packets labeled, i.e, the sum of the sizes of the values in {@link #mActionToTrafficMap}.
  26      */
  27     private long mPackets = 0;
  28
  29     public TrafficLabeler(List<UserAction> userActions) {
  30         // Init map with empty lists (no packets have been mapped to UserActions at the onset).
  31         mActionToTrafficMap = new HashMap<>();
  32         userActions.forEach(ua -> mActionToTrafficMap.put(ua, new ArrayList<>()));
  33         // Sort list of UserActions by timestamp in order to facilitate fast Packet-to-UserAction mapping.
  34         // For safety reasons, we create an internal copy of the list to prevent external code from changing the list's
  35         // contents as that would render our assumptions about order of elements invalid.
  36         // In addition, this also ensures that we do not break assumptions made by external code as we avoid reordering
  37         // the elements of the list passed from the external code.
  38         // If performance is to be favored over safety, assign userActions to mActionsSorted directly.
  39         mActionsSorted = new ArrayList<>();
  40         mActionsSorted.addAll(userActions);
  41         Collections.sort(mActionsSorted, (ua1, ua2) -> ua1.getTimestamp().compareTo(ua2.getTimestamp()));
  42     }
  43
  44
  45     @Override
  46     public void gotPacket(PcapPacket packet) {
  47         // Locate UserAction corresponding to packet, if any.
  48         int index = Collections.binarySearch(mActionsSorted, new UserAction(null, packet.getTimestamp()), (listItem, key) -> {
  49             // Start of inclusion interval is the time of the user action
  50             Instant intervalStart = listItem.getTimestamp();
  51             // End of inclusion interval is some arbitrary number of milliseconds after the user action.
  52             Instant intervalEnd = intervalStart.plusMillis(TriggerTrafficExtractor.INCLUSION_WINDOW_MILLIS);
  53             if (key.getTimestamp().isAfter(intervalStart) && key.getTimestamp().isBefore(intervalEnd)) {
  54                 // Packet lies within specified interval after the current UserAction, so we're done.
  55                 // Communicate termination to binarySearch by returning 0 which indicates equality.
  56                 return 0;
  57             }
  58             // If packet lies outside inclusion interval of current list item, continue search in lower or upper half of
  59             // list depending on whether the timestamp of the current list item is smaller or greater than that of the
  60             // packet.
  61             return listItem.getTimestamp().compareTo(key.getTimestamp());
  62         });
  63         if (index >= 0) {
  64             // Associate the packet to the its corresponding user action (located during the binary search above).
  65             mActionToTrafficMap.get(mActionsSorted.get(index)).add(packet);
  66             mPackets++;
  67         }
  68         // Ignore packet if it is not found to be in temporal proximity of a user action.
  69     }
  70
  71     /**
  72      * Get the total number of packets labeled by this {@code TrafficLabeler}.
  73      *
  74      * @return the total number of packets labeled by this {@code TrafficLabeler}.
  75      */
  76     public long getTotalPacketCount() {
  77         return mPackets;
  78     }
  79
  80     /**
  81      * Get the labeled traffic.
  82      *
  83      * @return A {@link Map} in which a {@link UserAction} points to a {@link List} of {@link PcapPacket}s believed to
  84      *         be related (occurring as a result of) that {@code UserAction}.
  85      */
  86     public Map<UserAction, List<PcapPacket>> getLabeledTraffic() {
  87         return Collections.unmodifiableMap(mActionToTrafficMap);
  88     }
  89
  90     /**
  91      * Like {@link #getLabeledTraffic()}, but allows the caller to supply a mapping function that is applied to
  92      * the traffic associated with each {@link UserAction} (the traffic label) before returning the labeled traffic.
  93      * This may for example be useful for a caller who wishes to perform some postprocessing of labeled traffic, e.g.,
  94      * in order to perform additional filtering or to transform the representation of labeled traffic.
  95      * <p>
  96      *     An example usecase is provided in {@link #getLabeledReassembledTcpTraffic()} which uses this function to
  97      *     build a {@link Map} in which a {@link UserAction} points to the reassembled TCP connections believed to have
  98      *     occurred as a result of that {@code UserAction}.
  99      * </p>
 100      *
 101      * @param mappingFunction A mapping function that converts a {@link List} of {@link PcapPacket} into some other type
 102      *                        {@code T}.
 103      * @param <T> The return type of {@code mappingFunction}.
 104      * @return A {@link Map} in which a {@link UserAction} points to the result of applying {@code mappingFunction} to
 105      *         the set of packets believed to be related (occurring as a result of) that {@code UserAction}.
 106      */
 107     public <T> Map<UserAction, T> getLabeledTraffic(Function<List<PcapPacket>, T> mappingFunction) {
 108         Map<UserAction, T> result = new HashMap<>();
 109         mActionToTrafficMap.forEach((ua, packets) -> result.put(ua, mappingFunction.apply(packets)));
 110         return result;
 111     }
 112
 113
 114     /**
 115      * Get the labeled traffic reassembled as TCP connections (<b>note:</b> <em>discards</em> all non-TCP traffic).
 116      *
 117      * @return A {@link Map} in which a {@link UserAction} points to a {@link List} of {@link Conversation}s believed to
 118      *         be related (occurring as a result of) that {@code UserAction}.
 119      */
 120     public Map<UserAction, List<Conversation>> getLabeledReassembledTcpTraffic() {
 121         return getLabeledTraffic(packets -> {
 122             TcpReassembler tcpReassembler = new TcpReassembler();
 123             packets.forEach(p -> tcpReassembler.gotPacket(p));
 124             return tcpReassembler.getTcpConversations();
 125         });
 126     }
 127
 128     /**
 129      * Like {@link #getLabeledReassembledTcpTraffic()}, but uses the provided {@code ipHostnameMappings} to group
 130      * {@link Conversation}s by hostname.
 131      *
 132      * @param ipHostnameMappings A {@link DnsMap} with IP to hostname mappings used for reverse DNS lookup.
 133      * @return A {@link Map} in which a {@link UserAction} points to the set of {@link Conversation}s believed to be
 134      *         related (occurring as a result of) that {@code UserAction}. More precisely, each {@code UserAction} in
 135      *         the returned {@code Map} points to <em>another</em> {@code Map} in which a hostname points to the set of
 136      *         {@code Conversation}s involving that hostname.
 137      */
 138     public Map<UserAction, Map<String, List<Conversation>>> getLabeledReassembledTcpTraffic(DnsMap ipHostnameMappings) {
 139         return getLabeledTraffic(packets -> {
 140             TcpReassembler tcpReassembler = new TcpReassembler();
 141             packets.forEach(p -> tcpReassembler.gotPacket(p));
 142             return TcpConversationUtils.groupConversationsByHostname(tcpReassembler.getTcpConversations(), ipHostnameMappings);
 143         });
 144     }
 145
 146 }