#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Copyright(C) 2008,2013 Eric Leblond

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 3 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
"""

import getopt
import sys
import os
import pg
from string import join
from configobj import ConfigObj
from validate import Validator
import pkg_resources
import visual
from time import sleep

from nf3d.infos import VERSION
from nf3d.connobj import packet, connection, connections, filters_list

def main_loop(connlists):
    visual.rate(50)
# Drag and drop loop
    connlist = connlists[0]
    while 1:
        sleep(0.05)
        if visual.scene.mouse.events:
            c = visual.scene.mouse.getevent()
            if c.pick and hasattr(c.pick,"icolor"):   # pick up the object
                for connl in connlists:
                    if c.pick in connl.conns + connl.packets:
                        connlist = connl
                        break
                if not c.shift:
                    connlist.normalize()
                if (hasattr(c.pick, "label")):
                    connlist.select(c.pick) 
                else:
                    connlist.normalize()
        if visual.scene.kb.keys: # is there an event waiting to be processed?
            s = visual.scene.kb.getkey() # obtain keyboard information
            if (len(s) == 1):
                if (s in filters_list):
                    connlist.highlight(filters_list[s])
                elif (s == 'r'):
                    connlist.refresh()
                elif (s == 'c'):
                    for connsl in connlists:
                        connsl.normalize()
                elif (s == 'l'):
                    for connsl in connlists:
                        connsl.toggle_label()
                elif (s == 'F'):
                    connlist.apply_filter()
                elif (s == 'O'):
                    connlist.switch_order()
                elif (s == 'R'):
                    connlist.reset_filter()
                elif (s == 'w'):
                    connlist.highlight_write()
                elif (s == 'a'):
                    connlist.toggle_adaptative()
                elif (s == 'q'):
		    print "'q' key pressed, exiting"
		    visual.scene.visible = 0
		    break
                elif (s == 'z'):
                    connlist.duration = connlist.duration/2
                    connlist.refresh()
                elif (s == 'Z'):
                    connlist.duration = connlist.duration*2
                    connlist.refresh()
                elif (s == 'C'):
                    for connsl in connlists:
                        connsl.set_level(connlist.length()+15*connsl.config['display']['radius'])
                    newconns = connections(connlist.starttime, connlist.endtime, connlist.duration, connsl.config)
                    newconns.build(conn = connlist.pgconn)
                    newconns.plate()
                    connlists.append(newconns)
                elif (s == 'D'):
                    if (len(connlists) > 1):
                        connlist.clear()
                        for connsl in connlists:
                            if (connsl == connlist):
                                break
                            connsl.set_level(-(connlist.length()+100*self.config['display']['radius']))
                        connlists.remove(connlist)
                        connlist = connlists[0]
                    else:
                        print "Will not destroy to avoid null display"
                elif (s == '?'):
                    help()
                elif (s == ':'):
                    interactive_filter(connlist)
            elif (s in ('up', 'down')):
                connlist.move_select(s)
            elif (s in ('left', 'right')):
                connlist.move_time(s)
            else:
                print "Key pressed: %s" % (s)
    print "Thanks for using nf3d"
    sys.exit(1)

def interactive_filter(connlist):
    print "Please give filter (key=value)"
    cmd = sys.stdin.readline()
    try:
        (k, v) = cmd.split('=')
    except:
        print "Syntax is 'key=value'"
        return
    filter_elt = {}
    if (filters_list.has_key(k) and v):
        connlist.highlight_filter = {}
        if (v.isdigit()):
            connlist.highlight_filter[filters_list[k]] = v
        else:
            connlist.highlight_filter[filters_list[k]] = v.strip()
        try:
            connlist.apply_filter()
        except:
            print "Invalid filter command: %s" % (cmd)
    else:
        print "Unknown key in command: %s" % (cmd)

def usage():
    filters = ""
    for k,v in filters_list.iteritems():
        filters += "[-%s %s] " % (k,v.upper())
    print "nf3d [-hV] [-S START] [-E END] [-D DURATION] %s" % filters

def version():
    print "nf3d (%s)" % (VERSION)

def help():
    print """
Keyboard shortcuts:\n
'c': switch highlighted items to normal
'l': toggle label display on selected items
'r': refresh datas
'w': output information about object to sdtout
down arrow: highlight next item
up arrow: highlight prev item
'd': highlight item with same original destination IP
's': highlight item with same original source IP
'p': highlight item with same original destination port
'P': highlight item with same original source port
right arrow: move time window right
left arrow: move time window left
'F': only display highlighted connections
'R': reset filter and display all connections in the time window
'C': duplicate the selected connections list
'D': delete the selected connections list
'O': order display based on last selected item
'?': display this help message
':': interactive filter
"""

def main():
    start = 0
    end = 0
    duration = 0
    configfile = "/etc/nf3d.conf"
    fullscreen = 0
    filter = {}
    try:
        options = "VhS:E:D:c:f" + ":".join(filters_list.keys()) + ":"  
        opts, args = getopt.getopt(sys.argv[1:], options, ["version", "help", "start=", "end=", "duration=", "config=", "fullscreen"])
    except getopt.GetoptError, err:
        # print help information and exit:
        print str(err) # will print something like "option -a not recognized"
        usage()
        sys.exit(2)
    output = None
    verbose = False
    for o, a in opts:
        if o == "-v":
            verbose = True
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-V", "--version"):
            version()
            sys.exit()
        elif o in ("-S", "--start"):
            start = int(a)
        elif o in ("-E", "--end"):
            end = int(a)
        elif o in ("-D", "--duration"):
            duration = int(a)
        elif o in ("-c", "--config"):
            configfile = a
	elif o in ("-f", "--fullscreen"):
	    fullscreen = 1
        elif o.strip('-') in filters_list.keys():
            filter[filters_list[o.strip('-')]] = a
        else:
            assert False, "unhandled option"

    if not os.path.exists(configfile):
        print "Unable to read config file %s" % (configfile)
        return

    vdt = Validator()
    cfgspec_filename = pkg_resources.resource_filename('nf3d', 'nf3dspec.conf')
    nf3dconfig = ConfigObj(configfile, configspec = cfgspec_filename)
    nf3dconfig.validate(vdt)

    visual.scene.title = "nf3d: 3D Netfilter visualization"
    visual.scene.width = int(nf3dconfig['display']['width'])
    visual.scene.height =  int(nf3dconfig['display']['height'])
    visual.scene.forward = (0.25,-10.25,-10)
    visual.scene.fullscreen = fullscreen
    visual.scene.autocenter = 1

    # Init connections list
    pgcnx = pg.connect(nf3dconfig['database']['database'], nf3dconfig['database']['host'], \
            int(nf3dconfig['database']['port']), None, None, \
            nf3dconfig['database']['login'], nf3dconfig['database']['password'])
    connlist = connections(start, end, duration, nf3dconfig)
    connlist.filter = filter
    connlist.build(conn = pgcnx)
    connlist.plate()
    connlists = []
    connlists.append(connlist)
    main_loop(connlists)

if __name__ == "__main__":
    main()

