#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Copyright(C) 2007 INL
Written by Damien Boucard <damien.boucard AT inl.fr>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 2 of the License.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
---
graphdesc is a program which generate a graphical representation of
an XML description file (typically called 'desc.xml').
"""

__revision__ = '0.1'
__author__ = 'Damien Boucard'
__copyright__ = 'Copyright 2007, INL'

import sys
from optparse import OptionParser
from pygraphviz import AGraph
from checkdesc import descxml, descmodels

OUTPUT_FORMATS = ['canon','dot','xdot','cmap','dia','fig','gd','gd2',
'gif','hpgl','imap','cmapx','ismap','jpg','jpeg',
'mif','mp','pcl','pic','plain','plain-ext','png','ps',
'ps2','svg','svgz','vrml','vtx','wbmp']
DEFAULT_FORMAT = 'png'
DEFAULT_PROG = 'twopi'

class Dumper:
    
    def __init__(self, desc, output):
        self.desc = desc
        self.output = output
        self.format = DEFAULT_FORMAT
        self.prog = DEFAULT_PROG
        self.graph = AGraph()
    
    def firewall_node(self, name):
        node = self.graph.get_node(name)
        node.attr['shape']='box'
        node.attr['fillcolor']='#FF4040'
        node.attr['color']='#FF8000'
        node.attr['style']='filled,setlinewidth(4)'
    
    def interface_node(self, name):
        node = self.graph.get_node(name)
        node.attr['shape']='circle'
        node.attr['fillcolor']='#FFB0B0'
    
    def address_node(self, name):
        node = self.graph.get_node(name)
        node.attr['shape']='invtriangle'
        node.attr['fillcolor']='#FFFF80'
    
    def network_node(self, name):
        node = self.graph.get_node(name)
        node.attr['shape']='ellipse'
        node.attr['fillcolor']='#80B0FF'
    
    def build(self):
        graph = self.graph
        graph.graph_attr['size'] = '16,12'
        graph.graph_attr['ratio'] = 'fill'
        graph.node_attr['style'] = 'filled'
        for firewall in desc.firewalls:
            graph.add_node(firewall.name)
            self.firewall_node(firewall.name)
            for interface in firewall.interfaces:
                graph.add_edge(firewall.name, interface.name)
                edge = graph.get_edge(firewall.name, interface.name)
                self.interface_node(interface.name)
                #for address in interface:
                #    graph.add_edge(interface.name, str(address.addr))
                #    self.address_node(str(address.addr))
        for network in desc.networks:
            graph.add_node("%s %s" %(network.name, network.addr))
            self.network_node("%s %s" %(network.name, network.addr))
            for connection in network:
                if isinstance(connection, descmodels.DirectConnection):
                    for address in connection.interface:
                        if address.addr in network.addr:
                            graph.add_edge("%s %s" %(network.name, network.addr), connection.interface.name)
                            edge = graph.get_edge("%s %s" %(network.name, network.addr), connection.interface.name)
                            edge.attr['label'] = str(address.addr)
                if isinstance(connection, descmodels.InternetConnection):
                    graph.add_edge("%s %s" %(network.name, network.addr), "@")
                    self.network_node("@")
                    edge = graph.get_edge("%s %s" %(network.name, network.addr), "@")
                    edge.attr['label'] = str(connection.default_gateway)
                if isinstance(connection, descmodels.RoutedConnection):
                    for network2 in desc.networks:
                        if connection.gateway in network2.addr:
                            graph.add_edge("%s %s" %(network.name, network.addr), "%s %s" %(network2.name, network2.addr))
                            edge = graph.get_edge("%s %s" %(network.name, network.addr), "%s %s" %(network2.name, network2.addr))
                            edge.attr['label'] = str(connection.gateway)
        graph.draw(self.output, self.format, self.prog)
            

def dump(desc, file):
    """
    Generate an XML file from a given desc model.

    @param desc: data to export into XML file.
    @type desc: descmodels.Desc
    @param file: XML file-like object to export to.
    @type file: file
    """
    Dumper(desc, file).build()

def parse_command_line():
    usage = "usage: %prog [options] desc_file.xml"
    parser = OptionParser(usage, version = '%prog '+ __revision__)

    # defining expected options
    parser.add_option('-f', '--format', help = "Output format of the graphical representation. If not set, try to guess with filename extension.", metavar = 'FORMAT')
    parser.add_option('-o', '--output', help = "Filename where writes the graphical representation ('<desc_filename>.png' by default).", metavar = 'FORMAT')
    parser.add_option('-l', '--list-formats', dest = "list_formats", help = "Display a list of supported output format and quit.", action = 'store_true')
    ##parser.set_defaults(rescue = False, forward = '', nat_rules = '', auth_ext = False)

    # parsing command line
    (options, args) = parser.parse_args(sys.argv)

    # checking options
    if options.list_formats :
        print "Available output format are:\n  *",\
              "\n  * ".join(OUTPUT_FORMATS)
        sys.exit(0)
    
    # checking number of arguments
    if len(args) != 2:
        print >>sys.stderr, "Bad number of arguments."
        parser.print_help()
        sys.exit(1)
    
    # guessing output filename
    if options.output is None:
        format = options.format
        if format is None:
            format = DEFAULT_FORMAT
        point_index = args[1].rfind('.')
        if point_index == -1:
            options.output = "%s.%s" %(args[1], format)
        else:
            options.output = "%s.%s" %(args[1][:point_index], format)

    # checking arguments
    try:
        desc_xml = open(args[1], 'r')
    except IOError, e:
        print >>sys.stderr, "%s: '%s'" %(e.strerror, e.filename)
        parser.print_help()
        sys.exit(1)

    return options, desc_xml

if __name__ == "__main__":
    options, desc_xml = parse_command_line()
    
    loader = descxml.Loader(desc_xml)
    desc = loader.build()
    
    output_file = open(options.output, "w")
    dumper = Dumper(desc_xml, output_file)
    if options.format is not None:
        dumper.format = options.format
    dumper.build()
    output_file.close()