Put checks for explicit termination of conversation to use in FlowPatternFinder.
[pingpong.git] / Code / Projects / SmartPlugDetector / src / main / java / edu / uci / iotproject / FlowPatternFinder.java
1 package edu.uci.iotproject;
2
3 import edu.uci.iotproject.comparison.ComparisonFunctions;
4 import edu.uci.iotproject.comparison.CompleteMatchPatternComparisonResult;
5 import edu.uci.iotproject.comparison.PatternComparisonTask;
6 import org.pcap4j.core.NotOpenException;
7 import org.pcap4j.core.PcapHandle;
8 import org.pcap4j.core.PcapNativeException;
9 import org.pcap4j.core.PcapPacket;
10 import org.pcap4j.packet.DnsPacket;
11 import org.pcap4j.packet.IpV4Packet;
12 import org.pcap4j.packet.TcpPacket;
13
14 import java.io.EOFException;
15 import java.net.UnknownHostException;
16 import java.util.*;
17 import java.util.concurrent.*;
18
19
20 /**
21  * <p>Provides functionality for searching for the presence of a {@link FlowPattern} in a PCAP trace.</p>
22  *
23  * <p>
24  * The (entire) PCAP trace is traversed and parsed on one thread (specifically, the thread that calls
25  * {@link #findFlowPattern()}). This thread builds a {@link DnsMap} using the DNS packets present in the trace and uses
26  * that {@code DnsMap} to reassemble {@link Conversation}s that <em>potentially</em> match the provided
27  * {@link FlowPattern} (in that one end/party of said conversations matches the hostname(s) specified by the given
28  * {@code FlowPattern}).
29  * These potential matches are then examined on background worker thread(s) to determine if they are indeed a (complete)
30  * matches of the provided {@code FlowPattern}.
31  * </p>
32  *
33  * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
34  * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
35  */
36 public class FlowPatternFinder {
37
38     /* Begin class properties */
39     /**
40      * {@link ExecutorService} responsible for parallelizing pattern searches.
41      * Declared as static to allow for reuse of threads across different instances of {@code FlowPatternFinder} and to
42      * avoid the overhead of initializing a new thread pool for each {@code FlowPatternFinder} instance.
43      */
44     private static final ExecutorService EXECUTOR_SERVICE = Executors.newCachedThreadPool();
45     /* End class properties */
46
47     /* Begin instance properties */
48     /**
49      * Holds a set of {@link Conversation}s that <em>potentially</em> match {@link #mPattern} since each individual
50      * {@code Conversation} is communication with the hostname identified by {@code mPattern.getHostname()}.
51      * Note that due to limitations of the {@link Set} interface (specifically, there is no {@code get(T t)} method),
52      * we have to resort to a {@link Map} (in which keys map to themselves) to "mimic" a set with {@code get(T t)}
53      * functionality.
54      *
55      * @see <a href="https://stackoverflow.com/questions/7283338/getting-an-element-from-a-set">this question on StackOverflow.com</a>
56      */
57     private final Map<Conversation, Conversation> mConversations;
58
59     private final DnsMap mDnsMap;
60     private final PcapHandle mPcap;
61     private final FlowPattern mPattern;
62
63     private final List<Future<CompleteMatchPatternComparisonResult>> mPendingComparisons = new ArrayList<>();
64     /* End instance properties */
65
66     /**
67      * Constructs a new {@code FlowPatternFinder}.
68      * @param pcap an <em>open</em> {@link PcapHandle} that provides access to the trace that is to be examined.
69      * @param pattern the {@link FlowPattern} to search for.
70      */
71     public FlowPatternFinder(PcapHandle pcap, FlowPattern pattern) {
72         this.mConversations = new HashMap<>();
73         this.mDnsMap = new DnsMap();
74         this.mPcap = Objects.requireNonNull(pcap,
75                 String.format("Argument of type '%s' cannot be null", PcapHandle.class.getSimpleName()));
76         this.mPattern = Objects.requireNonNull(pattern,
77                 String.format("Argument of type '%s' cannot be null", FlowPattern.class.getSimpleName()));
78     }
79
80     /**
81      * Starts the pattern search.
82      */
83     public void start() {
84         findFlowPattern();
85     }
86
87     /**
88      * Find patterns based on the FlowPattern object (run by a thread)
89      */
90     private void findFlowPattern() {
91         try {
92             PcapPacket packet;
93 //            TODO: The new comparison method is pending
94 //            TODO: For now, just compare using one hostname and one list per FlowPattern
95 //            List<String> hostnameList = mPattern.getHostnameList();
96 //            int hostIndex = 0;
97             while ((packet = mPcap.getNextPacketEx()) != null) {
98                 // Let DnsMap handle DNS packets.
99                 if (packet.get(DnsPacket.class) != null) {
100                     // Check if this is a valid DNS packet
101                     mDnsMap.validateAndAddNewEntry(packet);
102                     continue;
103                 }
104                 // For now, we only work support pattern search in TCP over IPv4.
105                 final IpV4Packet ipPacket = packet.get(IpV4Packet.class);
106                 final TcpPacket tcpPacket = packet.get(TcpPacket.class);
107                 if (ipPacket == null || tcpPacket == null) {
108                     continue;
109                 }
110
111                 String srcAddress = ipPacket.getHeader().getSrcAddr().getHostAddress();
112                 String dstAddress = ipPacket.getHeader().getDstAddr().getHostAddress();
113                 int srcPort = tcpPacket.getHeader().getSrcPort().valueAsInt();
114                 int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();
115                 // Is this packet related to the pattern; i.e. is it going to (or coming from) the cloud server?
116                 boolean fromServer = mDnsMap.isRelatedToCloudServer(srcAddress, mPattern.getHostname());
117                 boolean fromClient = mDnsMap.isRelatedToCloudServer(dstAddress, mPattern.getHostname());
118 //                String currentHostname = hostnameList.get(hostIndex);
119 //                boolean fromServer = mDnsMap.isRelatedToCloudServer(srcAddress, currentHostname);
120 //                boolean fromClient = mDnsMap.isRelatedToCloudServer(dstAddress, currentHostname);
121                 if (!fromServer && !fromClient) {
122                     // Packet not related to pattern, skip it.
123                     continue;
124                 }
125
126                 // Conversations (connections/sessions) are identified by the four-tuple
127                 // (clientIp, clientPort, serverIp, serverPort) (see Conversation Javadoc).
128                 // Create "dummy" conversation for looking up an existing entry.
129                 Conversation conversation = fromClient ? new Conversation(srcAddress, srcPort, dstAddress, dstPort) :
130                         new Conversation(dstAddress, dstPort, srcAddress, srcPort);
131                 // Add the packet so that the "dummy" conversation can be immediately added to the map if no entry
132                 // exists for the conversation that the current packet belongs to.
133                 if (tcpPacket.getHeader().getFin()) {
134                     // Record FIN packets.
135                     conversation.addFinPacket(packet);
136                 }
137                 if (tcpPacket.getPayload() != null) {
138                     // Record regular payload packets.
139                     conversation.addPacket(packet, true);
140                 }
141                 // Note: does not make sense to call attemptAcknowledgementOfFin here as the new packet has no FINs
142                 // in its list, so if this packet is an ACK, it would not be added anyway.
143
144                 // Need to retain a final reference to get access to the packet in the lambda below.
145                 final PcapPacket finalPacket = packet;
146                 // Add the new conversation to the map if an equal entry is not already present.
147                 // If an existing entry is already present, the current packet is simply added to that conversation.
148                 mConversations.merge(conversation, conversation, (existingEntry, toMerge) -> {
149                     // toMerge may not have any payload packets if the current packet is a FIN packet.
150                     if (toMerge.getPackets().size() > 0) {
151                         existingEntry.addPacket(toMerge.getPackets().get(0), true);
152                     }
153                     if (toMerge.getFinAckPairs().size() > 0) {
154                         // Add the FIN packet to the existing entry.
155                         existingEntry.addFinPacket(toMerge.getFinAckPairs().get(0).getFinPacket());
156                     }
157                     if (finalPacket.get(TcpPacket.class).getHeader().getAck()) {
158                         existingEntry.attemptAcknowledgementOfFin(finalPacket);
159                     }
160                     return existingEntry;
161                 });
162                 // Refresh reference to point to entry in map (in case packet was added to existing entry).
163                 conversation = mConversations.get(conversation);
164                 if (conversation.isGracefullyShutdown()) {
165                     // Conversation terminated gracefully, so we can now start analyzing it.
166                     // Remove the Conversation from the map and start the analysis.
167                     // Any future packets identified by the same four tuple will be tied to a new Conversation instance.
168                     mConversations.remove(conversation);
169                     // Create comparison task and send to executor service.
170                     PatternComparisonTask<CompleteMatchPatternComparisonResult> comparisonTask =
171                             new PatternComparisonTask<>(conversation, mPattern, ComparisonFunctions.COMPLETE_MATCH);
172                     mPendingComparisons.add(EXECUTOR_SERVICE.submit(comparisonTask));
173                     // Increment hostIndex to find the next
174                     
175                 }
176             }
177         } catch (EOFException eofe) {
178             // TODO should check for leftover conversations in map here and fire tasks for those.
179             // TODO [cont'd] such tasks may be present if connections did not terminate gracefully or if there are longlived connections.
180             System.out.println("[ findFlowPattern ] Finished processing entire PCAP stream!");
181             System.out.println("[ findFlowPattern ] Now waiting for comparisons to finish...");
182             // Wait for all comparisons to finish, then output their results to std.out.
183             for(Future<CompleteMatchPatternComparisonResult> comparisonTask : mPendingComparisons) {
184                 try {
185                     // Blocks until result is ready.
186                     CompleteMatchPatternComparisonResult comparisonResult = comparisonTask.get();
187                     if (comparisonResult.getResult()) {
188                         System.out.println(comparisonResult.getTextualDescription());
189                     }
190                 } catch (InterruptedException|ExecutionException e) {
191                     e.printStackTrace();
192                 }
193             }
194         } catch (UnknownHostException |
195                  PcapNativeException  |
196                  NotOpenException     |
197                  TimeoutException ex) {
198             ex.printStackTrace();
199         }
200     }
201
202 }