Working scripts and plots for 4 devices (smart plugs)
[pingpong.git] / base_gexf_generator.py
1 #!/usr/bin/python
2
3 """
4 Script that constructs a graph in which hosts are nodes.
5 An edge between two hosts indicate that the hosts communicate.
6 Hosts are labeled and identified by their IPs.
7 The graph is written to a file in Graph Exchange XML format for later import and visual inspection in Gephi.
8
9 Update per February 2, 2018:
10 Extension of base_gefx_generator.py.
11 This script constructs a bipartite graph with IoT devices on one side and Internet hosts on the other side.
12 As a result, this graph does NOT show inter IoT device communication.
13
14 The input to this script is the JSON output by extract_from_tshark.py by Anastasia Shuba.
15
16 This script is a simplification of Milad Asgari's parser_data_to_gephi.py script.
17 It serves as a baseline for future scripts that want to include more information in the graph.
18 """
19
20 import socket
21 import json
22 import tldextract
23 import networkx as nx
24 import sys
25 import csv
26 import re
27 import parser.parse_dns
28 from decimal import *
29 from networkx.algorithms import bipartite
30
31 # List of devices
32 DEVICE_MAC_LIST = "devicelist.dat"
33 EXCLUSION_MAC_LIST = "exclusion.dat"
34 COLUMN_MAC = "MAC_address"
35 COLUMN_DEVICE_NAME = "device_name"
36 # Fields
37 JSON_KEY_SOURCE = "_source"
38 JSON_KEY_LAYERS = "layers"
39 JSON_KEY_FRAME = "frame"
40 JSON_KEY_FRAME_PROTOCOLS = "frame.protocols"
41 JSON_KEY_FRAME_TIME_EPOCH = "frame.time_epoch"
42 JSON_KEY_FRAME_LENGTH = "frame.len"
43 JSON_KEY_ETH = "eth"
44 JSON_KEY_ETH_SRC = "eth.src"
45 JSON_KEY_ETH_DST = "eth.dst"
46 JSON_KEY_IPV6 = "ipv6"
47 JSON_KEY_IP = "ip"
48 JSON_KEY_IP_SRC = "ip.src"
49 JSON_KEY_IP_DST = "ip.dst"
50 # Checked protocols
51 JSON_KEY_UDP = "udp"
52 JSON_KEY_TCP = "tcp"
53 # List of checked protocols
54 listchkprot = [ "arp",
55                 "bootp",
56                 "dhcpv6",
57                 "dns",
58                 "llmnr",
59                 "mdns",
60                 "ssdp" ]
61
62 # Switch to generate graph that only shows local communication
63 ONLY_INCLUDE_LOCAL_COMMUNICATION = False
64
65
66 def create_device_list(dev_list_file):
67     """ Create list for smart home devices from a CSV file
68         Args:
69             dev_list_file: CSV file path that contains list of device MAC addresses
70     """
71     # Open the device MAC list file
72     with open(dev_list_file) as csvfile:
73         mac_list = csv.DictReader(csvfile, (COLUMN_MAC, COLUMN_DEVICE_NAME))
74         crude_list = list()
75         for item in mac_list:
76             crude_list.append(item)
77     # Create key-value dictionary
78     dev_list = dict()
79     for item in crude_list:
80         dev_list[item[COLUMN_MAC]] = item[COLUMN_DEVICE_NAME]
81         #print item["MAC_address"] + " => " + item["device_name"]
82     #for key, value in devlist.iteritems():
83     #    print key + " => " + value
84
85     return dev_list
86
87
88 def traverse_and_merge_nodes(G, dev_list_file):
89     """ Merge nodes that have similar properties, e.g. same protocols
90         But, we only do this for leaves (outer nodes), and not for
91         nodes that are in the middle/have many neighbors.
92         The pre-condition is that the node:
93         (1) only has one neighbor, and
94         (2) not a smarthome device.
95         then we compare the edges, whether they use the same protocols
96         or not. If yes, then we collapse that node and we attach
97         it to the very first node that uses that set of protocols.
98         Args:
99             G: a complete networkx graph
100             dev_list_file: CSV file path that contains list of device MAC addresses
101     """
102     nodes = G.nodes()
103     #print "Nodes: ", nodes
104     node_to_merge = dict()
105     # Create list of smarthome devices
106     dev_list = create_device_list(DEVICE_MAC_LIST)
107     # Traverse every node
108     # Check that the node is not a smarthome device
109     for node in nodes:
110         neighbors = G[node] #G.neighbors(node)
111         #print "Neighbors: ", neighbors, "\n"
112         # Skip if the node is a smarthome device
113         if node in dev_list:
114             continue
115         # Skip if the node has many neighbors (non-leaf) or no neighbor at all
116         if len(neighbors) is not 1:
117             continue
118         #print "Node: ", node
119         neighbor = neighbors.keys()[0] #neighbors[0]
120         #print "Neighbor: ", neighbors
121         protocols = G[node][neighbor]['Protocol']
122         #print "Protocol: ", protocols
123         # Store neighbor-protocol as key in dictionary
124         neigh_proto = neighbor + "-" + protocols
125         if neigh_proto not in node_to_merge:
126             node_to_merge[neigh_proto] = node
127         else:
128         # Merge this node if there is already an entry
129             # First delete
130             G.remove_node(node)
131             node_to_merge_with = node_to_merge[neigh_proto]
132             merged_nodes = G.node[node_to_merge_with]['Merged']
133             # Check if this is the first node
134             if merged_nodes is '':
135                 merged_nodes = node
136             else:
137             # Put comma if there is already one or more nodes
138                 merged_nodes += ", " + node
139             # Then attach as attribute
140             G.node[node_to_merge_with]['Merged'] = merged_nodes
141
142     return G
143
144
145 def place_in_graph(G, eth_src, eth_dst, device_dns_mappings, dev_list, layers, 
146         edge_to_prot, edge_to_vol):
147     """ Place nodes and edges on the graph
148         Args:
149             G: the complete graph
150             eth_src: MAC address of source
151             eth_dst: MAC address of destination
152             device_dns_mappings: device to DNS mappings (data structure)
153             dev_list: list of existing smarthome devices
154             layers: layers of JSON file structure
155             edge_to_prot: edge to protocols mappings
156             edge_to_vol: edge to traffic volume mappings
157     """
158     # Get timestamp of packet (router's timestamp)
159     timestamp = Decimal(layers[JSON_KEY_FRAME][JSON_KEY_FRAME_TIME_EPOCH])
160     # Get packet length
161     packet_len = Decimal(layers[JSON_KEY_FRAME][JSON_KEY_FRAME_LENGTH])
162     # Get the protocol and strip just the name of it
163     long_protocol = layers[JSON_KEY_FRAME][JSON_KEY_FRAME_PROTOCOLS]
164     # Split once starting from the end of the string and get it
165     split_protocol = long_protocol.split(':')
166     protocol = None
167     if len(split_protocol) < 5:
168         last_index = len(split_protocol) - 1
169         protocol = split_protocol[last_index]
170     else:
171         protocol = split_protocol[3] + ":" + split_protocol[4]
172     #print "timestamp: ", timestamp, " - new protocol added: ", protocol, "\n"
173     # And source and destination IPs
174     ip_src = layers[JSON_KEY_IP][JSON_KEY_IP_SRC]
175     ip_dst = layers[JSON_KEY_IP][JSON_KEY_IP_DST]
176     # Categorize source and destination IP addresses: local vs. non-local
177     ip_re = re.compile(r'\b192.168.[0-9.]+')
178     src_is_local = ip_re.search(ip_src) 
179     dst_is_local = ip_re.search(ip_dst)
180     # Store protocol into the set (source)
181     protocols = None
182     # Key to search in the dictionary is <src-mac-address>-<dst-mac_address>
183     dict_key = ip_src + "-" + ip_dst
184     #print "Key: ", dict_key
185     if dict_key not in edge_to_prot:
186         edge_to_prot[dict_key] = set()
187     protocols = edge_to_prot[dict_key]
188     protocols.add(protocol)
189     protocols_str = ', '.join(protocols)
190     #print "protocols: ", protocols_str, "\n"
191     # Check packet length and accumulate to get traffic volume
192     if dict_key not in edge_to_vol:
193         edge_to_vol[dict_key] = 0;
194     edge_to_vol[dict_key] = edge_to_vol[dict_key] + packet_len
195     volume = str(edge_to_vol[dict_key])
196
197     # Skip device to cloud communication if we are interested in the local graph.
198     # TODO should this go before the protocol dict is changed?
199     if ONLY_INCLUDE_LOCAL_COMMUNICATION and not (src_is_local and dst_is_local):
200         return
201
202     #print "ip.src =", ip_src, "ip.dst =", ip_dst, "\n"
203     # Place nodes and edges
204     src_node = None
205     dst_node = None
206     # Integer values used for tagging nodes, indicating to Gephi if they are local IoT devices or web servers.
207     remote_node = 0
208     local_node = 1
209     # Values for the 'bipartite' attribute of a node when constructing the bipartite graph
210     bipartite_iot = 0
211     bipartite_web_server = 1
212     if src_is_local:
213         G.add_node(eth_src, Name=dev_list[eth_src], islocal=local_node, bipartite=bipartite_iot)
214         src_node = eth_src
215     else:
216         hostname = None
217         # Check first if the key (eth_dst) exists in the dictionary
218         if eth_dst in device_dns_mappings:
219             # If the source is not local, then it's inbound traffic, and hence the eth_dst is the MAC of the IoT device.
220             hostname = device_dns_mappings[eth_dst].hostname_for_ip_at_time(ip_src, timestamp)                   
221         if hostname is None:
222             # Use IP if no hostname mapping
223             hostname = ip_src
224         # Non-smarthome devices can be merged later
225         G.add_node(hostname, Merged='', islocal=remote_node, bipartite=bipartite_web_server)
226         src_node = hostname
227
228     if dst_is_local:
229         G.add_node(eth_dst, Name=dev_list[eth_dst], islocal=local_node, bipartite=bipartite_iot)
230         dst_node = eth_dst
231     else:
232         hostname = None
233         # Check first if the key (eth_dst) exists in the dictionary
234         if eth_src in device_dns_mappings:
235             # If the destination is not local, then it's outbound traffic, and hence the eth_src is the MAC of the IoT device.
236             hostname = device_dns_mappings[eth_src].hostname_for_ip_at_time(ip_dst, timestamp)
237         if hostname is None:
238             # Use IP if no hostname mapping
239             hostname = ip_dst
240         # Non-smarthome devices can be merged later
241         G.add_node(hostname, Merged='', islocal=remote_node, bipartite=bipartite_web_server)
242         dst_node = hostname
243     G.add_edge(src_node, dst_node, Protocol=protocols_str, Volume=volume)
244
245
246 def parse_json(file_path):
247     """ Parse JSON file and create graph
248         Args:
249             file_path: path to the JSON file
250     """
251     # Create a smart home device list
252     dev_list = create_device_list(DEVICE_MAC_LIST)
253     # Create an exclusion list
254     exc_list = create_device_list(EXCLUSION_MAC_LIST)
255     # First parse the file once, constructing a map that contains information about individual devices' DNS resolutions.
256     device_dns_mappings = parser.parse_dns.parse_json_dns(file_path) # "./json/eth1.dump.json"
257     # Init empty graph
258     G = nx.DiGraph()
259     # Mapping from edge to a set of protocols
260     edge_to_prot = dict()
261     # Mapping from edge to traffic volume
262     edge_to_vol = dict()
263     # Parse file again, this time constructing a graph of device<->server and device<->device communication.
264     with open(file_path) as jf:
265         # Read JSON; data becomes reference to root JSON object (or in our case json array)
266         data = json.load(jf)
267         # Loop through json objects (packets) in data
268         for p in data:
269             # p is a JSON object, not an index - drill down to object containing data from the different layers
270             layers = p[JSON_KEY_SOURCE][JSON_KEY_LAYERS]
271
272             iscontinue = False
273             for prot in listchkprot:
274                 if prot in layers:
275                     iscontinue = True
276             if iscontinue:
277                 continue            
278
279             # Skip any non udp/non tcp traffic
280             if JSON_KEY_UDP not in layers and JSON_KEY_TCP not in layers:
281                 continue
282
283             # Fetch source and destination MACs
284             eth = layers.get(JSON_KEY_ETH, None)
285             if eth is None:
286                 print "[ WARNING: eth data not found ]"
287                 continue
288             eth_src = eth.get(JSON_KEY_ETH_SRC, None)
289             eth_dst = eth.get(JSON_KEY_ETH_DST, None)
290             # Exclude devices in the exclusion list
291             if eth_src in exc_list:
292                 print "[ WARNING: Source ", eth_src, " is excluded from graph! ]"
293                 continue
294             if eth_dst in exc_list:
295                 print "[ WARNING: Destination ", eth_dst, " is excluded from graph! ]"
296                 continue
297             # Exclude if IP does not exist in layers - this means IPv6
298             if JSON_KEY_IP not in layers and JSON_KEY_IPV6 in layers:
299                 continue
300             
301             # Place nodes and edges in graph
302             place_in_graph(G, eth_src, eth_dst, device_dns_mappings, dev_list, layers, 
303                 edge_to_prot, edge_to_vol)
304
305     # Print DNS mapping for reference
306         #for mac in device_dns_mappings:
307         #       ddm = device_dns_mappings[mac]
308         #       ddm.print_mappings()
309     
310     return G
311
312
313 # ------------------------------------------------------
314 # Not currently used.
315 # Might be useful later on if we wish to resolve IPs.
316 def get_domain(host):
317     ext_result = tldextract.extract(str(host))
318     # Be consistent with ReCon and keep suffix
319     domain = ext_result.domain + "." + ext_result.suffix
320     return domain
321
322 def is_IP(addr):
323     try:
324         socket.inet_aton(addr)
325         return True
326     except socket.error:
327         return False
328 # ------------------------------------------------------
329
330
331 if __name__ == '__main__':
332     if len(sys.argv) < 3:
333         print "Usage:", sys.argv[0], "input_file output_file"
334         print "outfile_file should end in .gexf"
335         sys.exit(0)
336     # Input file: Path to JSON file generated from tshark JSON output using Anastasia's script (extract_from_tshark.py).
337     input_file = sys.argv[1]
338     print "[ input_file  =", input_file, "]"
339     # Output file: Path to file where the Gephi XML should be written.
340     output_file = sys.argv[2]
341     print "[ output_file =", output_file, "]"
342     # Construct graph from JSON
343     G = parse_json(input_file)
344     # Contract nodes that have the same properties, i.e. same protocols
345     G = traverse_and_merge_nodes(G, DEVICE_MAC_LIST)
346     # Write Graph in Graph Exchange XML format
347     nx.write_gexf(G, output_file)