#!/usr/bin/env python -O
# Rick van der Zwet, (c) Dec 2006
# <rick@wzoeterwoude.net>

#Simple module to probe the network and draw a picture about it
# Requirements
# -pygraphviz 3.33 <https://networkx.lanl.gov/wiki/pygraphviz>

import sys
import re
import pickle
import getopt

from common import *
from pygraphviz import AGraph




class ip2name:
    ip2name_dict = {}
    def __init__(self,infile):
        self.ip2name_dict = {}
        log("Using %s for conversion" % infile,2)
        try:
            input = open(infile, 'r')
        except:
            return None
        for line in  input.readlines():
            (k, l ) =  line.rstrip("\n").split("|")
            self.ip2name_dict[k] = l
        input.close()

    def output(self,ip):
        try:
            name = self.ip2name_dict[ip]
        except:
            name = ''
        return ("%s [%s]" % (ip, name))



def draw(infile,outfile,names_infile):
    ip2node = ip2name(names_infile)

    G=AGraph(directed=True)
    G.node_attr['style'] = 'filled'
    G.node_attr['shape'] = 'box'
    G.node_attr['color'] = 'green'
    G.graph_attr['splines']='true'
    G.graph_attr['overlap']='false'
    G.graph_attr['fixedsize']='false'

    file = open(infile, 'rb')
    nodes = pickle.load(file)
    file.close()
    list = nodes.keys()
    for k in list:
        k_node = ip2node.output(k)
        log("Node %s" % ip2node.output(k),3)
        for l in list:
            l_node = ip2node.output(l)
            if nodes[k].linked(nodes[l]):
                G.add_edge(l_node,k_node)
                log("-%s" % l_node,3)
    
    G.layout()
    G.draw(outfile)
    log("File saved in '%s'" % outfile,2)



def usage():
    print '''---
%(prog)s --draw [--infile=<file>] [--outfile=<file>] [--ip2name=<file>]
%(prog)s --help''' % { 'prog': sys.argv[0] }


def main():
    global loglevel
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hdi::o::v::', \
        ["help", "draw","outfile=", "infile=","ip2name=","verbose="])
    except getopt.GetoptError:
        usage()
        sys.exit(128)

    action = False
    infile = False
    outfile = False
    ip2name = False

    for o, a in opts:
        log("Command '%s' = '%s'" % (o, a),2)
        if o in ("-v", "--verbose"):
            if a.isdigit():
                loglevel = a
            else:
                loglevel = a.count('v')
            log("Loglevel now at %s" % loglevel,4)
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        if o in ("-o", "--outfile"):
            outfile = a
        if o in ("-i", "--infile"):
             infile = a
        if o in ("--ip2name"):
            ip2name = a
        if o in ("-d", "--draw"):
            action = 'draw'

    if action == 'draw':
        if not infile:
            infile = 'data.pkl'
        if not outfile:
            outfile = 'network.svg'
        draw(infile,outfile,ip2name)
    else:
        usage()
        sys.exit(128)

if __name__ == "__main__":
    main()