1 package edu.uci.iotproject.analysis;
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;
9 import java.time.Instant;
11 import java.util.function.Function;
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.
17 * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
18 * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
20 public class TrafficLabeler implements PacketListener {
22 private final Map<UserAction, List<PcapPacket>> mActionToTrafficMap;
23 private final List<UserAction> mActionsSorted;
25 * The total number of packets labeled, i.e, the sum of the sizes of the values in {@link #mActionToTrafficMap}.
27 private long mPackets = 0;
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()));
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.
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
61 return listItem.getTimestamp().compareTo(key.getTimestamp());
64 // Associate the packet to the its corresponding user action (located during the binary search above).
65 mActionToTrafficMap.get(mActionsSorted.get(index)).add(packet);
68 // Ignore packet if it is not found to be in temporal proximity of a user action.
72 * Get the total number of packets labeled by this {@code TrafficLabeler}.
74 * @return the total number of packets labeled by this {@code TrafficLabeler}.
76 public long getTotalPacketCount() {
81 * Get the labeled traffic.
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}.
86 public Map<UserAction, List<PcapPacket>> getLabeledTraffic() {
87 return Collections.unmodifiableMap(mActionToTrafficMap);
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.
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}.
101 * @param mappingFunction A mapping function that converts a {@link List} of {@link PcapPacket} into some other type
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}.
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)));
115 * Get the labeled traffic reassembled as TCP connections (<b>note:</b> <em>discards</em> all non-TCP traffic).
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}.
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();
129 * Like {@link #getLabeledReassembledTcpTraffic()}, but uses the provided {@code ipHostnameMappings} to group
130 * {@link Conversation}s by hostname.
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.
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);