Add parse_dns.py: reads DNS traffic from a file (JSON formatted by tshark) and constr...
[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 def parse_json(file_path):
22     # Init empty graph
23     G = nx.DiGraph() 
24     with open(file_path) as jf:
25         # Read JSON.
26         # data becomes reference to root JSON object (or in our case json array)
27         data = json.load(jf)
28         # Loop through json objects in data
29         for k in data:
30             # Fetch source and destination IPs.
31             # Each of these become a Node in the Graph.
32             src_ip = data[k]["src_ip"]
33             dst_ip = data[k]["dst_ip"]
34             ''' Graph construction '''
35             # No need to check if the Nodes and/or Edges we add already exist:
36             # NetworkX won't add already existing nodes/edges (except in the case of a MultiGraph or MultiDiGraph (see NetworkX doc)).
37             # Add a node for each host.
38             G.add_node(src_ip)
39             G.add_node(dst_ip)
40             # Connect these two nodes.
41             G.add_edge(src_ip, dst_ip)
42     return G
43
44 # ------------------------------------------------------
45 # Not currently used.
46 # Might be useful later on if we wish to resolve IPs.
47 def get_domain(host):
48     ext_result = tldextract.extract(str(host))
49     # Be consistent with ReCon and keep suffix
50     domain = ext_result.domain + "." + ext_result.suffix
51     return domain
52
53 def is_IP(addr):
54     try:
55         socket.inet_aton(addr)
56         return True
57     except socket.error:
58         return False
59 # ------------------------------------------------------
60
61 if __name__ == '__main__':
62     if len(sys.argv) < 3:
63         print "Usage:", sys.argv[0], "input_file output_file"
64         print "outfile_file should end in .gexf"
65         sys.exit(0)
66     # Input file: Path to JSON file generated from tshark JSON output using Anastasia's script (extract_from_tshark.py).
67     input_file = sys.argv[1]
68     print "[ input_file  =", input_file, "]"
69     # Output file: Path to file where the Gephi XML should be written.
70     output_file = sys.argv[2]
71     print "[ output_file =", output_file, "]"
72     # Construct graph from JSON
73     G = parse_json(input_file)
74     # Write Graph in Graph Exchange XML format
75     nx.write_gexf(G, output_file)