Initial experimentation with simple pipeline setup: Call parse_json_dns from parse_dn...
[pingpong.git] / base_gefx_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 The input to this script is the JSON output by extract_from_tshark.py by Anastasia Shuba.
10
11 This script is a simplification of Milad Asgari's parser_data_to_gephi.py script.
12 It serves as a baseline for future scripts that want to include more information in the graph.
13 """
14
15 import socket
16 import json
17 import tldextract
18 import networkx as nx
19 import sys
20
21 import parse_dns
22
23 def parse_json(file_path):
24
25     maps_tuple = parse_dns.parse_json_dns("./dns.json")
26     hn_ip_map = maps_tuple[0]
27     ip_hn_map = maps_tuple[1]
28
29     # Init empty graph
30     G = nx.DiGraph() 
31     with open(file_path) as jf:
32         # Read JSON.
33         # data becomes reference to root JSON object (or in our case json array)
34         data = json.load(jf)
35         # Loop through json objects in data
36         for k in data:
37             #print "k is:",k
38             # Fetch source and destination IPs.
39             # Each of these become a Node in the Graph.
40             src_ip = data[k]["src_ip"]
41             dst_ip = data[k]["dst_ip"]
42
43             if dst_ip in ip_hn_map:
44                 # hack to get first element in set
45                 for e in ip_hn_map[dst_ip]:
46                     break
47                 dst_ip = e
48
49             ''' Graph construction '''
50             # No need to check if the Nodes and/or Edges we add already exist:
51             # NetworkX won't add already existing nodes/edges (except in the case of a MultiGraph or MultiDiGraph (see NetworkX doc)).
52             # Add a node for each host.
53             G.add_node(src_ip)
54             G.add_node(dst_ip)
55             # Connect these two nodes.
56             G.add_edge(src_ip, dst_ip)
57     return G
58
59 # ------------------------------------------------------
60 # Not currently used.
61 # Might be useful later on if we wish to resolve IPs.
62 def get_domain(host):
63     ext_result = tldextract.extract(str(host))
64     # Be consistent with ReCon and keep suffix
65     domain = ext_result.domain + "." + ext_result.suffix
66     return domain
67
68 def is_IP(addr):
69     try:
70         socket.inet_aton(addr)
71         return True
72     except socket.error:
73         return False
74 # ------------------------------------------------------
75
76 if __name__ == '__main__':
77     if len(sys.argv) < 3:
78         print "Usage:", sys.argv[0], "input_file output_file"
79         print "outfile_file should end in .gexf"
80         sys.exit(0)
81     # Input file: Path to JSON file generated from tshark JSON output using Anastasia's script (extract_from_tshark.py).
82     input_file = sys.argv[1]
83     print "[ input_file  =", input_file, "]"
84     # Output file: Path to file where the Gephi XML should be written.
85     output_file = sys.argv[2]
86     print "[ output_file =", output_file, "]"
87     # Construct graph from JSON
88     G = parse_json(input_file)
89     # Write Graph in Graph Exchange XML format
90     nx.write_gexf(G, output_file)