convert nfdump csv output to gephi csv

flow_to_gephi.rb

#!/usr/bin/env ruby

if ARGV.empty?
  puts "Usage: #{$0} <nfdump_csv_file>" 
  exit
end

require 'csv'

csv_file = ARGV[0]

# read flow csv file
all_nodes = []
CSV.foreach(csv_file) do |row|
	all_nodes << row[3] << row[4]
end

# add id to each ip address
id_nodes = []
all_nodes.uniq.each do |ip|
	id_nodes << [ip, ip]
end

# write to node.csv
CSV.open("#{csv_file}.node.csv","wb", {:write_headers => true, :headers => ['Id', 'Label']}) do |csv|
	id_nodes.each do |node|
		csv << node
	end
end

# {ip address => node id}
#dict =  Hash[*id_nodes.flatten].invert

# create flows with id
all_flows = []
CSV.foreach(csv_file) do |row|
	all_flows << [row[3], row[4]]
end

# calculate width of edge from flow count
flow_count = {}
flow_group = all_flows.group_by{|flow| flow}
flow_group.each do |k, v|
	flow_count[k] = v.size
end

# write to edge.csv
CSV.open("#{csv_file}.edge.csv","wb", {:write_headers => true, :headers => ['Source', 'Target', 'Weight']}) do |csv|
	all_flows.uniq.each do |flow|
		flow << flow_count[flow]
		csv << flow 
	end  
end