Convert tab indented to GML graph

convert-to-gml.py

#!/bin/python
## sudo apt-get install python-networkx python-pydot

# Converts from tab indented lines to graph modelling language (GML)
# Example indented lines:
# 
# Duqu 2.0 : SAND
# 	Spear-phishing : SAND 
# 		Exploit CVE-2014-4148 : SAND
# 			Target downloads and opens word document : SAND
# 				Vulnerability in Windows TrueType Font : SAND
# 				Kernel-mode access : SAND

# Usage python convert-to-gml.py -f file.txt -w output.dot

import csv, argparse

def main(read, write):

	try:
		outFile = open(write, 'w')
	except:
		print e

	opts = "\nrankdir=RL;\nordering=out;\n\n"
	nodes = []
	nodesStr = ""
	edgesStr = ""

	# Read tab indented file
	with open(read, 'rb') as csvfile:
		reader = csv.reader(csvfile, dialect=csv.excel_tab)
		for row in reader:
			# Count the number of indents.
			indent = 0
			for i in row:
				if i == '':
					indent = indent+1

			# Extract the name and operator.
			try: 
				name = row[-1].split(":")[0].strip()
			except: 
				name = ""
			try:
				operator = row[-1].split(":")[1].strip()
			except:
				operator = ""
			# try:
			# 	group = row[-1].split(":")[2].strip()
			# except:
			# 	group = ""

			# Ignore comments
			if name and name[0] != '#':
				nodes.append([name, operator, indent])
			else: 
				if not name == "": 
					print "[-] Skipping comment -", name[1:]

	weightNo = len(nodes)
	for i in range(len(nodes)):
		operator = nodes[i][1]
		shape = "ellipse"
		
		if operator == "JOIN":
			shape = "component"

		nodesStr += "\"{}\" [{}, {}];\r\n".format(nodes[i][0], "fontname=sans", "shape="+shape)

		if i == 0:
			continue

		# Look backwards and find the parent node (indent-1) 
		parent = ""
		for k in range(i):
			if (nodes[i][2]) -1 == nodes[k][2]:
				parent = nodes[k]

		# print nodes[i][0], " child of ", parent[0]
		arrowHead = "none" # Default OR
		colour = "#3EBCE6" # Default OR
		
		if operator == "SAND":
			arrowHead = "dot"
			colour = "#BC3EE6" # Purpleish BC3EE6 C390D4
		elif operator == "AND":
			arrowHead = "normal"
			colour = "#68E63E" # Greenish 68E63E 90C3D4
		elif operator == "OR":
			arrowHead = "none"
			colour = "#3EBCE6" # Blueish 3EBCE6 A1D490
		
		edgesStr += "\"{}\" -> \"{}\" [{}, {}, {}];\r\n".format(nodes[i][0], parent[0], "arrowhead="+arrowHead, "color=\""+colour+"\"", "weight="+str(weightNo))
		
		weightNo = weightNo - 1

	outFile.write("strict digraph \"\" {\n")
	outFile.write(opts)
	outFile.write(nodesStr+"\n")
	outFile.write(edgesStr)
	outFile.write("}")

if __name__ == "__main__":
	parser = argparse.ArgumentParser(description='Tab indented lines to GML')
	
	parser.add_argument('-f', '--file', dest='file', action='store', 
		help='Tab indented file to read.', required=True)

	parser.add_argument('-w', '--write', dest='output', action='store', 
		help='Write GML file here', required=True)

	args = parser.parse_args()
	main(args.file, args.output)