Completing flow for time series graph generation

[pingpong.git] / base_gefx_generator.py

diff --git a/base_gefx_generator.py b/base_gefx_generator.py
index 703fe456bbcdc70e5b06546c9f946fa3316eca9b..ca3aba70c4006ef6a9776f8b06fb3b5eb86cec3c 100644 (file)
--- a/base_gefx_generator.py
+++ b/base_gefx_generator.py
@@ -17,16 +17,36 @@ import json
 import tldextract
 import networkx as nx
 import sys
+import csv
 from decimal import *
 
-import parse_dns
+import parser.parse_dns
+
+DEVICE_MAC_LIST = "devicelist.dat"
+COLUMN_MAC = "MAC_address"
+COLUMN_DEVICE_NAME = "device_name"
 
 JSON_KEY_ETH_SRC = "eth.src"
 JSON_KEY_ETH_DST = "eth.dst"
 
 def parse_json(file_path):
 
-    device_dns_mappings = parse_dns.parse_json_dns("./dns.json")
+    # Open the device MAC list file
+    with open(DEVICE_MAC_LIST) as csvfile:
+        maclist = csv.DictReader(csvfile, (COLUMN_MAC, COLUMN_DEVICE_NAME))
+        crudelist = list()
+        for item in maclist:
+            crudelist.append(item)
+            #print(item)
+    # Create key-value dictionary
+    devlist = dict()
+    for item in crudelist:
+        devlist[item[COLUMN_MAC]] = item[COLUMN_DEVICE_NAME]
+        #print item["MAC_address"] + " => " + item["device_name"]
+    #for key, value in devlist.iteritems():
+    #    print key + " => " + value
+
+    device_dns_mappings = parser.parse_dns.parse_json_dns("./json/dns.json")
 
     # Init empty graph
     G = nx.DiGraph() 
@@ -49,12 +69,12 @@ def parse_json(file_path):
             elif eth_dst in device_dns_mappings:
                 iot_device = eth_dst
             else:
-                print "[ WARNING: DNS mapping not found for device with MAC", eth_src, "OR", eth_dst, "]"
+#                print "[ WARNING: DNS mapping not found for device with MAC", eth_src, "OR", eth_dst, "]"
                 # This must be local communication between two IoT devices OR an IoT device talking to a hardcoded IP.
                 # For now let's assume local communication.
                 # Add a node for each device and an edge between them.
-                G.add_node(eth_src)
-                G.add_node(eth_dst)
+                G.add_node(eth_src, Name=devlist[eth_src])
+                G.add_node(eth_dst, Name=devlist[eth_src])
                 G.add_edge(eth_src, eth_dst)
                 # TODO add regex check on src+dst IP to figure out if hardcoded server IP (e.g. check if one of the two are NOT a 192.168.x.y IP)
                 continue
@@ -67,7 +87,7 @@ def parse_json(file_path):
             
             # Add a node for each host.
             # First add node for IoT device.
-            G.add_node(iot_device)
+            G.add_node(iot_device, Name=devlist[eth_src])
             # Then add node for the server.
             # For this we need to distinguish between outbound and inbound traffic so that we look up the proper IP in our DNS map.
             # For outbound traffic, the server's IP is the destination IP.
@@ -79,9 +99,9 @@ def parse_json(file_path):
                 # However, we only get here for the DNS that have not performed any DNS lookups
                 # We should use a regex check early in the loop to see if it is two local devices communicating.
                 # This way we would not have to consider these corner cases later on.
-                print "[ WARNING: no ip-hostname mapping found for ip", server_ip, " -- adding eth.src->eth.dst edge, but note that this may be incorrect if IoT device has hardcoded server IP ]"
-                G.add_node(eth_src)
-                G.add_node(eth_dst)
+#                print "[ WARNING: no ip-hostname mapping found for ip", server_ip, " -- adding eth.src->eth.dst edge, but note that this may be incorrect if IoT device has hardcoded server IP ]"
+                G.add_node(eth_src, Name=devlist[eth_src])
+                G.add_node(eth_dst, Name=devlist[eth_src])
                 G.add_edge(eth_src, eth_dst)
                 continue
             G.add_node(hostname)