Code/Projects/PacketLevelSignatureExtractor/src/main/java/edu/uci/iotproject/DnsMap.java

   1 package edu.uci.iotproject;
   2
   3 import org.pcap4j.core.PacketListener;
   4 import org.pcap4j.core.PcapPacket;
   5 import org.pcap4j.packet.Packet;
   6 import org.pcap4j.packet.DnsPacket;
   7 import org.pcap4j.packet.DnsResourceRecord;
   8 import org.pcap4j.packet.namednumber.DnsResourceRecordType;
   9
  10
  11 import java.net.Inet4Address;
  12 import java.net.UnknownHostException;
  13 import java.util.*;
  14
  15
  16 /**
  17  * This is a class that does DNS mapping.
  18  * Basically an IP address is mapped to its
  19  * respective DNS hostnames.
  20  *
  21  * @author Rahmadi Trimananda (rtrimana@uci.edu)
  22  * @version 0.1
  23  */
  24 public class DnsMap implements PacketListener {
  25
  26     /* Class properties */
  27     private Map<String, Set<String>> ipToHostnameMap;
  28
  29     /* Class constants */
  30     private static final Set<String> EMPTY_SET = Collections.unmodifiableSet(new HashSet<>());
  31
  32
  33     /* Constructor */
  34     public DnsMap() {
  35         ipToHostnameMap = new HashMap<>();
  36     }
  37
  38     @Override
  39     public void gotPacket(PcapPacket packet) {
  40         try {
  41             validateAndAddNewEntry(packet);
  42         } catch (UnknownHostException e) {
  43             e.printStackTrace();
  44         }
  45     }
  46
  47     /**
  48      * Gets a packet and determine if this is a DNS packet
  49      *
  50      * @param   packet  Packet object
  51      * @return          DnsPacket object or null
  52      */
  53     private DnsPacket getDnsPacket(Packet packet) {
  54         DnsPacket dnsPacket = packet.get(DnsPacket.class);
  55         return dnsPacket;
  56     }
  57
  58     /**
  59      * Checks DNS packet and build the map data structure that
  60      * maps IP addresses to DNS hostnames
  61      *
  62      * @param   packet  PcapPacket object
  63      */
  64     public void validateAndAddNewEntry(PcapPacket packet) throws UnknownHostException {
  65         // Make sure that this is a DNS packet
  66         DnsPacket dnsPacket = getDnsPacket(packet);
  67         if (dnsPacket != null) {
  68             // We only care about DNS answers
  69             if (dnsPacket.getHeader().getAnswers().size() != 0) {
  70                 String hostname = dnsPacket.getHeader().getQuestions().get(0).getQName().getName();
  71                 for(DnsResourceRecord answer : dnsPacket.getHeader().getAnswers()) {
  72                     // We only care about type A records
  73                     if (!answer.getDataType().equals(DnsResourceRecordType.A))
  74                         continue;
  75                     // Sanity check. For some reason the hostname appears to be the empty string in the answer .
  76                     // We hence have to assume that all answers correspond to a single question that holds the hostname as part of its object tree.
  77                     // Therefore, if there are more questions in one query-reply exchange, we are in trouble.
  78                     if (!answer.getName().getName().equals("") && !answer.getName().getName().equals(hostname))
  79                         throw new RuntimeException("[DNS parser] mismatch between hostname in question and hostname in answer");
  80                     // The IP in byte representation.
  81                     byte[] ipBytes = answer.getRData().getRawData();
  82                     // Convert to string representation.
  83                     String ip = Inet4Address.getByAddress(ipBytes).getHostAddress();
  84                     Set<String> hostnameSet = new HashSet<>();
  85                     hostnameSet.add(hostname);
  86                     // Update or insert depending on presence of key:
  87                     // Concat the existing set and the new set if ip already present as key,
  88                     // otherwise add an entry for ip pointing to new set.
  89                     ipToHostnameMap.merge(ip, hostnameSet, (v1, v2) -> { v1.addAll(v2); return v1; });
  90                 }
  91             }
  92         }
  93     }
  94
  95
  96     /**
  97      * Checks DNS packet and build the map data structure that
  98      * maps IP addresses to DNS hostnames
  99      *
 100      * @param   address     Address to check
 101      * @param   hostname    Hostname to check
 102      */
 103     public boolean isRelatedToCloudServer(String address, String hostname) {
 104         return ipToHostnameMap.getOrDefault(address, EMPTY_SET).contains(hostname);
 105     }
 106
 107     public Set<String> getHostnamesForIp(String ip) {
 108         Set<String> hostnames = ipToHostnameMap.get(ip);
 109         return hostnames != null ? Collections.unmodifiableSet(hostnames) : null;
 110     }
 111 }