Refactoring and restructuring - Adding DnsMap class
[pingpong.git] / Code / Projects / SmartPlugDetector / src / main / java / edu / uci / iotproject / FlowPatternFinder.java
1 package edu.uci.iotproject;
2
3 import org.pcap4j.core.NotOpenException;
4 import org.pcap4j.core.PcapHandle;
5 import org.pcap4j.core.PcapNativeException;
6 import org.pcap4j.core.PcapPacket;
7 import org.pcap4j.packet.IpV4Packet;
8 import org.pcap4j.packet.Packet;
9 import org.pcap4j.packet.TcpPacket;
10 import org.pcap4j.packet.DnsPacket;
11
12 import java.io.EOFException;
13 import java.net.UnknownHostException;
14 import java.time.Instant;
15 import java.util.*;
16 import java.util.concurrent.TimeoutException;
17
18 /**
19  * Provides functionality for searching for the presence of a {@link FlowPattern} in a PCAP trace.
20  *
21  * @author Janus Varmarken
22  */
23 public class FlowPatternFinder {
24
25     /* Class properties */
26     private final Map<Conversation, List<PcapPacket>> connections = new HashMap<>();
27     
28     private DnsMap dnsMap;
29     
30     public FlowPatternFinder() {
31         this.dnsMap = new DnsMap();
32     }
33
34     // TODO clean up exceptions etc.
35     public void findFlowPattern(PcapHandle pcap, FlowPattern pattern)
36             throws PcapNativeException, NotOpenException, TimeoutException {
37         int counter = 0;
38         try {
39             PcapPacket packet;
40             while ((packet = pcap.getNextPacketEx()) != null) {
41
42                 // Check if this is a valid DNS packet
43                 dnsMap.validateAndAddNewEntry(packet);
44                 // For now, we only work support pattern search in TCP over IPv4.
45                 IpV4Packet ipPacket = packet.get(IpV4Packet.class);
46                 TcpPacket tcpPacket = packet.get(TcpPacket.class);
47                 if (ipPacket == null || tcpPacket == null) {
48                     continue;
49                 }
50                 String srcAddress = ipPacket.getHeader().getSrcAddr().getHostAddress();
51                 String dstAddress = ipPacket.getHeader().getDstAddr().getHostAddress();
52                 int srcPort = tcpPacket.getHeader().getSrcPort().valueAsInt();
53                 int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();
54                 // Is this packet related to the pattern and coming from the cloud server?
55                 boolean fromServer = dnsMap.isRelatedToCloudServer(srcAddress, pattern.getHostname());
56                 // Is this packet related to the pattern and going to the cloud server?
57                 boolean fromClient = dnsMap.isRelatedToCloudServer(dstAddress, pattern.getHostname());
58                 if (!fromServer && !fromClient) {
59                     // Packet not related to pattern, skip it.
60                     continue;
61                 }
62                 if (tcpPacket.getPayload() == null) {
63                     // We skip non-payload control packets as these are less predictable and should therefore not be
64                     // part of a signature (e.g. receiver can choose not to ACK immediately)
65                     continue;
66                 }
67
68                 // Identify conversations (connections/sessions) by the four-tuple (clientIp, clientPort, serverIp, serverPort).
69                 // TODO: this is strictly not sufficient to differentiate one TCP session from another, but should suffice for now.
70                 Conversation conversation = fromClient ? new Conversation(srcAddress, srcPort, dstAddress, dstPort) :
71                         new Conversation(dstAddress, dstPort, srcAddress, srcPort);
72                 List<PcapPacket> listWrappedPacket = new ArrayList<>();               
73                 listWrappedPacket.add(packet);
74                 // Create new conversation entry, or append packet to existing.
75                 connections.merge(conversation, listWrappedPacket, (v1, v2) -> {
76                     // TODO: in theory, this is insufficient to detect retransmissions due to TCP seq.no. rollover.
77                     // TODO: bad for performance, O(n) for each packet added to flow (n being length of the flow).
78                     boolean retransmission = v1.stream().anyMatch(p -> p.get(TcpPacket.class).getHeader().getSequenceNumber() == v2.get(0).get(TcpPacket.class).getHeader().getSequenceNumber());
79                     if (!retransmission) {
80                         // Do not add if retransmission -> avoid duplicate packets in flow
81                         v1.addAll(v2);
82                     }
83                     return v1;
84                 });
85             }
86         } catch (EOFException eofe) {
87             System.out.println("findFlowPattern: finished processing entire file");
88             find(pattern);
89         } catch (UnknownHostException ex) {
90             System.out.println(); 
91             ex.printStackTrace();
92         }
93     }
94
95     private void find(FlowPattern pattern) {
96         for (Conversation con : connections.keySet()) {
97             List<PcapPacket> packets = connections.get(con);
98             if (packets.size() != pattern.getPacketOrder().size()) {
99                 // Not a complete match if different number of packets.
100                 continue;
101             }
102             boolean completeMatch = true;
103             for (int i = 0; i < packets.size(); i++) {
104                 TcpPacket tcpPacket = packets.get(i).get(TcpPacket.class);
105                 if (tcpPacket.getPayload().length() != pattern.getPacketOrder().get(i)) {
106                     completeMatch = false;
107                     break;
108                 }
109             }
110             if (completeMatch) {
111                 PcapPacket firstPacketInFlow = packets.get(0);
112                 System.out.println(
113                         String.format("[ detected a complete match of pattern '%s' at %s]",
114                                 pattern.getPatternId(), firstPacketInFlow.getTimestamp().toString()));
115             }
116         }
117     }
118
119     /**
120      * Immutable class used for identifying a conversation/connection/session/flow (packet's belonging to the same
121      * session between a client and a server).
122      */
123     private static class Conversation {
124
125         private final String clientIp;
126         private final int clientPort;
127         private final String serverIp;
128         private final int serverPort;
129
130         public Conversation(String clientIp, int clientPort, String serverIp, int serverPort) {
131             this.clientIp = clientIp;
132             this.clientPort = clientPort;
133             this.serverIp = serverIp;
134             this.serverPort = serverPort;
135         }
136
137
138         // =========================================================================================================
139         // We simply reuse equals and hashCode methods of String.class to be able to use this immutable class as a key
140         // in a Map.
141         @Override
142         public boolean equals(Object obj) {
143             return obj instanceof Conversation && this.toString().equals(obj.toString());
144         }
145         @Override
146         public int hashCode() {
147             return toString().hashCode();
148         }
149         // =========================================================================================================
150
151         @Override
152         public String toString() {
153             return String.format("%s:%d %s:%d", clientIp, clientPort, serverIp, serverPort);
154         }
155     }
156
157 }