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)