#!/usr/bin/env python
#
# ypfw: IPFW/Dummynet front-end
#
# 2004 Copyright (c) CAIA
# Centre for Advanced Internet Architectures
# Swinburne Univeristy of Technology
# http://caia.swin.edu.au
#
# author: Claudio Favi
# email: claudio.favi@epfl.ch
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILIgabpaul@melbpc.org.auTY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import curses
import curses.textpad
import curses.panel
import curses.wrapper
import os
import re
from os import popen, popen3

def change_wnd_colors(win, fg, bg):
    if curses.has_colors():
        if bg == curses.COLOR_BLACK:
            cp = 1
        else:
            cp = bg #color pair nb
        if fg == curses.COLOR_WHITE and bg == curses.COLOR_BLACK:
            win.bkgdset(ord(' '), curses.color_pair(0))
        else:
            curses.init_pair(cp, fg, bg)
            win.bkgdset(ord(' '), curses.color_pair(cp))
    else:
        win.bkgdset(ord(' '))#, curses.A_BOLD)

class Window:
    def __init__(self, h, w, y, x, title=None, bgcolor=None, fgcolor=None):
        self.title = title
        self.win = curses.newwin(h,w,y,x)
        self.panel = curses.panel.new_panel(self.win)
        self.panel.set_userptr(self)
        self._setdimensions(h,w,y,x)
        self.statusline = ""
        self.shortcuts = ""
        self.bgcolor = bgcolor
        self.fgcolor = fgcolor
        self.setcolors()
        
        
    def _setdimensions(self, h, w, y, x):
        self.h = h
        self.w = w
        self.y = y
        self.x = x

    def setcolors(self):
        if self.bgcolor != None:
            if self.fgcolor != None:
                fg = self.fgcolor
            else:
                if self.bgcolor == curses.COLOR_BLUE:
                    fg = curses.COLOR_WHITE
                else:
                    fg = curses.COLOR_BLACK
            change_wnd_colors(self.win, self.fgcolor, self.bgcolor)
            
    def resize(self, h, w, y, x):
        self.win.resize(h,w)
        self.win.mvwin(y, x)
        self._setdimensions(h,w,y,x)
        #self.setcolors()
        self.win.redrawwin()
        self.win.refresh()

    def updatewnd(self):
        #print self.title
        if self.h > 2:
            self.win.addnstr(self.h-2, 0,
                             " " + self.title + " " + self.shortcuts + " "*self.w, self.w,
                             curses.A_STANDOUT)
            self.win.move(self.h-1,0)
            self.win.clrtoeol()
            self.win.addnstr(self.h-1, 0, self.statusline, self.w)

    def update(self, key):
        raise StandardError, "should override update()" 


    def handle_input(self, key):
        raise StandardError, "should override handle_input()" 
    


class RulesWindow(Window):
    def __init__(self, h, w, y, x):
        Window.__init__(self, h, w, y, x,
                        title="Rules",
                        bgcolor=curses.COLOR_BLACK,
                        fgcolor=curses.COLOR_GREEN)
        self.ruleswin = curses.newpad(300, w)#self.win.subwin(h-2, w, y, x)

        change_wnd_colors(self.ruleswin, curses.COLOR_WHITE, curses.COLOR_BLUE)
        self.shortcuts = "(a)dd (d)elete (z)ero counters (u)pdate display"
        self.rules = []
        self.active = 0
        self.updaterules()
        self.old_add_entry = ""

    def update(self):
        updaterules(self)
        updatewnd(self)

    def updaterules(self):
        del self.rules[:]
#       pin, pout, perr = popen3("ipfw -a list", 'r')
        pout = popen("ipfw -a list", 'r')
        self.rules = [ l.strip('\n') for l in pout.readlines() ]
#        self.statusline = perr.readline().strip('\n')
        self.active = 0
#        pin.close(); pout.close(); perr.close()

    def resize(self, h, w, y, x):
        Window.resize(self, h, w, y, x)
        self.ruleswin.resize(h-3, w)
        self.ruleswin.mvwin(y+1,x)
        
    
    def updatewnd(self):
        Window.updatewnd(self)
        self.ruleswin.move(0,0) #move cursor
        self.ruleswin.clrtoeol()
        #for i in range(min(len(self.rules), self.h-3)):
        for i in range( len(self.rules) ):
            if i == self.active:
                mode = curses.A_STANDOUT
            else:
                mode = curses.A_NORMAL
            
            self.ruleswin.addnstr(i,0,self.rules[i]+" "*self.w, self.w , mode)

        self.ruleswin.clrtobot()
        self.win.noutrefresh()

        s = max(0, self.active - self.h + 3)
        self.ruleswin.noutrefresh(s,0, self.y, self.x, self.h -2, self.w + self.x)


    def addrule(self):
        # should : open rule add dialog and call ipfw add
        adddiag = curses.newwin(3, self.w-2,
                                self.h/2 - 2 +self.y,
                                1 + self.x)
        change_wnd_colors(adddiag, curses.COLOR_BLACK, curses.COLOR_GREEN)
        adddiag.border()
        adddiag.addstr(0,2,"Add Rule")
        adddiag.addstr(1,1, "ipfw add ")
        textwin = adddiag.subwin(1, self.w-13, self.h/2 - 1 + self.y, 11+self.x)
        change_wnd_colors(textwin, curses.COLOR_WHITE, curses.COLOR_BLACK)
        tpad = curses.textpad.Textbox(textwin)
        textwin.addstr(0,0,self.old_add_entry)
        adddiag.refresh()
        textwin.refresh()
        tpad.edit()
        cmd = tpad.gather()
 #       pin, pout, perr = popen3("ipfw add " + cmd)
        perr = popen("ipfw add " + cmd)
        self.old_add_entry = cmd
        self.statusline = perr.readline().strip('\n')
#        pin.close(); pout.close(); perr.close()
        del tpad, textwin, adddiag
        self.updaterules()

    def deleterule(self):
        if len(self.rules) == 0 or not 0<=self.active<len(self.rules) :
            return
        #call ipfw delete
        rulenb = self.rules[self.active].split()[0]
#        pin, pout, perr = popen3("ipfw delete " + rulenb)
        perr = popen("ipfw delete " + rulenb)
        self.statusline = perr.readline().strip('\n')
#        pin.close(); pout.close(); perr.close()
        self.updaterules()
        #del self.rules[self.active]
        #if len(self.rules) == 0:
        #    self.active = 0
        #else:
        #    self.active = (self.active-1) % len(self.rules)

    def zerocounters(self):
        perr = popen("ipfw zero")
        self.statusline = perr.readline().strip('\n')
        self.updaterules()
    
    def handle_input(self, key):
        if 0<key<256:
            c=chr(key)
            if c in 'uU':
                self.updaterules()
            elif c in 'aA':
                self.addrule()
            elif c in 'dD':
                self.deleterule()
            elif c in 'zZ':
                self.zerocounters()
        else:
            if key == curses.KEY_UP and len(self.rules)!=0:
                self.active = (self.active - 1)%len(self.rules)
                self.statusline = ""
            elif key == curses.KEY_DOWN and len(self.rules)!=0:
                self.active = (self.active + 1)%len(self.rules)
                self.statusline = ""
        self.updatewnd()

class PipesWindow(Window):
    def __init__(self, h, w, y, x):
        Window.__init__(self, h, w, y, x, title="Pipes")
        self.pipeswin = curses.newpad(300, w)#self.win.subwin(h-2, w, y, x)

        change_wnd_colors(self.pipeswin, curses.COLOR_WHITE, curses.COLOR_BLUE)

        self.shortcuts = "(a)dd (d)elete (c)onfig (u)pdate display"
        self.pipes = []
        self.updatepipes()
        self.active = 0
        self.old_add_entry = ""
        self.old_config_entry = ""

    def update(self):
        updatepipes(self)
        updatewnd(self)
        
    def updatepipes(self):
        del self.pipes[:]

        pout = popen("ipfw pipe list", 'r')
        lines = [ l.strip('\n') for l in pout ]
        i=0
        while i < len(lines):
            if re.match("^\d+:", lines[i]):
                self.pipes.append(lines[i])
                self.pipes.append(lines[i+1])
                i += 1
            i += 1
            
        pout.close()

    def addpipe(self):
        adddiag = curses.newwin(3, self.w-2,
                                self.h/2 - 2 +self.y,
                                1 + self.x)
        change_wnd_colors(adddiag, curses.COLOR_BLACK, curses.COLOR_GREEN)
        adddiag.border()
        adddiag.addstr(0,2,"Add Pipe")
        adddiag.addstr(1,1, "Pipe nb: ")
        textwin = adddiag.subwin(1, self.w-13, self.h/2 - 1 + self.y, 11+self.x)
        change_wnd_colors(textwin, curses.COLOR_WHITE, curses.COLOR_BLACK)
        tpad = curses.textpad.Textbox(textwin)
        textwin.addstr(0,0,self.old_add_entry)
        adddiag.refresh()
        textwin.refresh()
        tpad.edit()
        cmd = tpad.gather()
 #       pin, pout, perr = popen3("ipfw add " + cmd)
        perr = popen("ipfw pipe " + cmd + " config 0")
        self.old_add_entry = cmd
        self.statusline = perr.readline().strip('\n')
#        pin.close(); pout.close(); perr.close()
        del tpad, textwin, adddiag
        self.updatepipes()

    def configpipe(self):
        pipenb = self.pipes[self.active*2].split(': ')[0]
        adddiag = curses.newwin(3, self.w-2,
                                self.h/2 - 2 +self.y,
                                1 + self.x)
        change_wnd_colors(adddiag, curses.COLOR_BLACK, curses.COLOR_GREEN)
        adddiag.border()
        adddiag.addstr(0,2,"Config Pipe %d (ipfw pipe %d config...)"%(int(pipenb), int(pipenb)))
        textwin = adddiag.subwin(1, self.w-4, self.h/2 - 1 + self.y, 2+self.x)
        change_wnd_colors(textwin, curses.COLOR_WHITE, curses.COLOR_BLACK)
        tpad = curses.textpad.Textbox(textwin)
        textwin.addstr(0,0,self.old_config_entry)
        adddiag.refresh()
        textwin.refresh()
        tpad.edit()
        cmd = tpad.gather()
 #       pin, pout, perr = popen3("ipfw add " + cmd)
        perr = popen("ipfw pipe " + pipenb + " config " + cmd)
        self.old_config_entry = cmd
        self.statusline = perr.readline().strip('\n')
#        pin.close(); pout.close(); perr.close()
        del tpad, textwin, adddiag
        self.updatepipes()

    def deletepipe(self):
        if len(self.pipes) == 0 or not 0<=self.active<len(self.pipes)/2:
            return
        pipenb = self.pipes[self.active*2].split(': ')[0]
        perr = popen('ipfw pipe ' + pipenb + ' delete')
        self.statusline = perr.readline().strip('\n')
        self.updatepipes()

    def resize(self, h, w, y, x):
        Window.resize(self, h, w, y, x)
        self.pipeswin.resize(h-3, w)
        self.pipeswin.mvwin(y+1,x)

    def updatewnd(self):
        Window.updatewnd(self)
        self.pipeswin.move(0,0)
        self.pipeswin.clrtoeol()
        for i in range( len(self.pipes)/2 ):
#        for i in range(min(len(self.pipes)/2, (self.h-4)/2)):
            if i == self.active:
                mode = curses.A_STANDOUT
            else:
                mode = curses.A_NORMAL

            self.pipeswin.addnstr( (i*2), 0, self.pipes[i*2]+" "*self.w, self.w , mode)
            self.pipeswin.addnstr( (i*2)+1, 0, self.pipes[i*2+1]+" "*self.w, self.w , mode)
        self.pipeswin.clrtobot()
        self.win.noutrefresh()

        s = max(0, self.active*2 - self.h + 4)
        self.pipeswin.noutrefresh(s,0, self.y, self.x, self.h - 2, self.w + self.x )
        
    def handle_input(self, key):
        if 0<key<256:
            c=chr(key)
            if c in 'uU':
                self.updatepipes()
            elif c in 'aA':
                self.addpipe()
            elif c in 'cC':
                self.configpipe()
            elif c in 'dD':
                self.deletepipe()
            elif c in 'zZ':
                pass
                #self.zerocounters()
        else:
            if key == curses.KEY_UP and len(self.pipes)!=0:
                self.active = (self.active - 1)%(len(self.pipes)/2)
                self.statusline = ""
            elif key == curses.KEY_DOWN and len(self.pipes)!=0:
                self.active = (self.active + 1)%(len(self.pipes)/2)
                self.statusline = ""
        self.updatewnd()



class Mainscreen:
    def __init__(self, stdscr):
        self.stdscr = stdscr
        self.stdscr.scrollok(0)
        self.scr_h, self.scr_w = self.stdscr.getmaxyx()
        #self.stdscr.border()
        self.createwindows()
#        self.positionwindows()
#        curses.curs_set(0)

    def run(self):
        self.headerwin.addnstr(0,0, "http://caia.swin.edu.au"+ " "*self.scr_w, self.scr_w)
        self.headerwin.addnstr(0,(self.scr_w/2-2), "YpFW" + " "*(self.scr_w/2-18) + "Powered by CAIA ", self.scr_w)
        curses.panel.update_panels()
        self.headerwin.noutrefresh()
        curses.panel.top_panel().userptr().updatewnd()
#        self.ruleswin.win.noutrefresh()
        curses.doupdate()

        while 1:
            key = self.stdscr.getch()
            if key==ord('\t'):
                curses.panel.bottom_panel().top()
                curses.panel.update_panels()
                curses.panel.top_panel().userptr().updatewnd()
            elif key==ord('q') or key==ord('Q'):
                break
            else:
                curses.panel.top_panel().userptr().handle_input(key)

            curses.doupdate()
    
    def createwindows(self):
        self.headerwin = curses.newwin(2, self.scr_w, 0, 0)
        if curses.has_colors():
            fg = curses.COLOR_BLACK
            bg = curses.COLOR_GREEN
            cp = bg #color pair nb
            curses.init_pair(cp, fg, bg)
            self.headerwin.bkgdset(ord(' '), curses.color_pair(cp))
        else:
            self.headerwin.bkgdset(ord(' '))
        
        self.pipeswin = PipesWindow(self.scr_h-1, self.scr_w, 1, 0)
        self.ruleswin = RulesWindow(self.scr_h-1, self.scr_w, 1, 0)
        self.ruleswin.panel.top()

    def positionwindows(self):
        self.pipeswin.resize(self.scr_h-1, self.scr_w, 1, 0)
        self.ruleswin.resize(self.scr_h-1, self.scr_w, 1, 0)
        


def main(stdscr):
    m = Mainscreen(stdscr);
    m.run()

curses.wrapper(main)