deltahammer

#!/usr/bin/python2.6
import curses
import time
import optparse
import sys
import re
import signal
import readers
import templates
from math import ceil
from string import Template

description="""Description: %prog is a tool to monitor the delta of data and statistics from various static commands.
The program requires either the interface option (-i INTERFACE) or the statistics option (-s TARGET).
Values not in parentheses are deltas, values in parentheses are absolutes."""


# Arguments:
def sigwinch_handler(n, frame):
    curses.endwin()
    curses.initscr()


# Arguments: two dictionaries
# Return: one dictionary
def compare(original, current):
    delta = {}
    for data in original.keys():
        delta[data] = str(int(current[data])\
                            -int(original[data]))
    return delta


# Arguments: two nested dictionaries
# Return: one nested dictionary
def compare_nested(original, current):
    delta = {}
    for name in original.keys():
        delta[name] = {}
        for data in original[name].keys():
            delta[name][data] = str(int(current[name][data])\
                                -int(original[name][data]))
    return delta


# Decides what arguments to call readers.ip_s_link with depending on what argument was given
# Argument: a dictionary containing interface names as the keys and a list of all IPs assigned 
#           to that interface as the value.
# Return: a dictionary for every interface nested in a parent dictionary
def get_link_data(interfaces):
    if opts.interface == 'all':
        data = readers.ip_s_link(interfaces.keys())
    else:
        try: interfaces[opts.interface]
        except KeyError: sys.exit("error: %s is not a valid interface (try 'all')" % opts.interface)
        data = readers.ip_s_link([opts.interface])
    return data


# Argument: a list of sections from 'netstat -s'
# Return: a dictionary containing all the data needed from 'netstat -s' with SYN_RECV count added
def get_statistics_ip_data(sect_list):
    data = readers.netstat_s(sect_list)
    data['syn_recv_count'] = readers.get_output('/usr/sbin/ss -t -a state syn-recv | wc -l').strip()
    return data


# Tries to get the next key pressed
# Argument: a curses window
# Return: a char containing the next key pressed or None
def get_key(stdscr):
    stdscr.nodelay(True) # make getkey not wait
    try: key = stdscr.getkey()
    except: key = None
    else: return key


# Arguments: body - a multiline string
#            page - an int of the page wanted
#            height - screen height
# Returns: string containing the requested section of the body
#          int the page number returned
#          int the last page number
def get_stats_page(stdscr, template, original_data, current_data, requested_page=0):
    (height, width) = stdscr.getmaxyx()
    current_time=time.strftime("%X", time.localtime())
    delta_data = compare(original_data, current_data)
    # Substitute the delta values into the template
    body = Template(template).safe_substitute(delta_data)
    # Substitute the absolute values in
    body = body.replace('_abso', '')
    body = Template(body).safe_substitute(current_data)
    # Remove any lines with unfilled in data
    body = re.sub('\n[^\n]+\$[^\n]+', '', body)
    current_page = 1
    # Get current page from body if given
    if requested_page > 0:
        header_len = len(re.findall('\n', templates.header))
        footer_len = len(re.findall('\n', templates.footer)) + 1
        body_len = len(re.findall('\n', body)) + 1
        usable_height = height - header_len - footer_len
        body_split = body.split('\n')
        page_dic = {1:''}
        for index in range(len(body_split)):
            # If current page is filled to the usable_height, increment current page
            # and start new page
            if int(ceil((index + 1) / float(usable_height))) > current_page:
                current_page += 1
                page_dic[current_page] = ''
            page_dic[current_page] += body_split[index] + '\n'
        # If requested_page is greater than the last page, set to last page
        if requested_page > current_page:
            requested_page = current_page
        body = page_dic[requested_page]

    header = Template(templates.header).safe_substitute(command='netstat -s', delta=opts.delay, 
                                                        start_time=opts.start_time, current_time=current_time,
                                                        current_page=requested_page, last_page=current_page)
    output = header + body

    return (output, requested_page)


# Updates screen, EXITS if string fails to post correctly
# Arguments: stdscr - a curses window
#            output - string to post
def post(stdscr, output):
        stdscr.clear()
        try: stdscr.addstr(0, 0, output)
        except: sys.exit('error: screen creation failed. Your terminal size is small and you should feel small')
        stdscr.refresh()


# Main function for the interface page
# Argument: a curses window
# Return: a string of the last output
def link(stdscr):
    curses.use_default_colors()
    interfaces = readers.ip_a()
    original_data = get_link_data(interfaces)
    while True:
        current_data = get_link_data(interfaces)
        delta_data = compare_nested(original_data, current_data)
        current_time=time.strftime("%X", time.localtime())
        output = Template(templates.header).safe_substitute(command='ip -s link', delta=opts.delay, 
                                                            start_time=opts.start_time, current_time=current_time,
                                                            current_page='1', last_page='1')
        for name in delta_data.keys():
            ips = ', '.join(interfaces[name]['ipv4'])
            # Substitute the delta values into the template
            section = Template(templates.link).safe_substitute(delta_data[name], ips=ips, name=name)
            # Substitute the absolute values in
            section = section.replace('_abso', '')
            section = Template(section).safe_substitute(current_data[name])
            output += section + "\n\n"
        # Post output to screen
        output += templates.footer
        post(stdscr, output)

        # While delay, get keys and deal with it them
        stop = time.time() + int(opts.delay)
        while time.time() < stop:
            key = get_key(stdscr)
            if key == 'q':
                return output 
            elif key == 'r':
                opts.start_time=time.strftime("%X", time.localtime())
                original_data = get_link_data(interfaces)
            elif key == '+':
                opts.delay += 1
            elif key == '-':
                if opts.delay > 1: opts.delay -= 1




# Main function for the statistics IP page
# Argument: a curses window
# Return: a string of the last output
def statistics_ip(stdscr):
    curses.use_default_colors()
    sect_list = ['Ip','IpExt']
    original_data = get_statistics_ip_data(sect_list)
    current_page = 1
    while True:
        current_data = get_statistics_ip_data(sect_list)
        # Compiles body and then caculates the current section that can fit on the screen
        (output, current_page) = get_stats_page(stdscr, templates.ip, original_data, current_data, current_page)
        output += templates.footer
        post(stdscr, output)

        # While delay, get keys and deal with it them
        stop = time.time() + int(opts.delay)
        while time.time() < stop:
            key = get_key(stdscr)
            if key == 'q':
                (output, current_page) = get_stats_page(stdscr, templates.ip, original_data, current_data)
                return output
            elif key == 'r':
                opts.start_time=time.strftime("%X", time.localtime())
                original_data = get_statistics_ip_data(sect_list)
            elif key == '+':
                opts.delay += 1
            elif key == '-':
                if opts.delay > 1: opts.delay -= 1
            elif key == 'KEY_PPAGE':
                if current_page > 1: current_page -= 1
            # Check if there is a next page later
            elif key == 'KEY_NPAGE':
                current_page += 1
            
        
# Main function for the statistics IP page
# Argument: a curses window
# Return: a string of the last output
def statistics_tcpv(stdscr):
    curses.use_default_colors()
    sect_list = ['Tcp','TcpExt']
    original_data = readers.netstat_s(sect_list)
    current_page = 1
    while True:
        current_data = readers.netstat_s(sect_list)
        # Compiles body and then caculates the current section that can fit on the screen
        (output, current_page) = get_stats_page(stdscr, templates.tcpv, original_data, current_data, current_page)
        output += templates.footer
        post(stdscr, output)

        # While delay, get keys and deal with it them
        stop = time.time() + int(opts.delay)
        while time.time() < stop:
            key = get_key(stdscr)
            if key == 'q':
                (output, current_page) = get_stats_page(stdscr, templates.tcpv, original_data, current_data)
                return output
            elif key == 'r':
                opts.start_time=time.strftime("%X", time.localtime())
                original_data = readers.netstat_s(sect_list)
            elif key == '+':
                opts.delay += 1
            elif key == '-':
                if opts.delay > 1: opts.delay -= 1
            elif key == 'KEY_PPAGE':
                if current_page > 1: current_page -= 1
            # Check if there is a next page later
            elif key == 'KEY_NPAGE':
                current_page += 1
            

# Parse settings from commandline options
# TODO: add program description including that must be run as root
# TODO: group require options (
cli_parser = optparse.OptionParser(usage='%prog [OPTIONS]',
                                   description=description)
cli_parser.add_option('-d', '--delay', dest='delay', default=5, type='int', metavar='SECONDS', \
                      help="Delay between updates [default: %default seconds]")
cli_parser.add_option('-i', '--interface', dest='interface', type='string', metavar='INTERFACE', \
                      help="Shows link status of specified interface (or put 'all'). Output similar to 'ifconfig' or 'ip -s link'")
cli_parser.add_option('-s', '--statistics', dest='statistics', type='choice', choices=['ip', 'tcp', 'tcpv'], metavar='TARGET', \
                      help="Shows various statistics from 'netstat -s'.  Must specify a target to watch (options are: ip, tcp, tcpv)")

(opts, args) = cli_parser.parse_args()

# Verify arguments
if opts.interface and opts.statistics:
    sys.exit('error: invalid option combination, see help')
elif not (opts.interface or opts.statistics):
    sys.exit('error: missing required option, see help')
elif opts.delay < 1:
    sys.exit('error: delay time cannot be less than 1')

# Set start time
opts.start_time=time.strftime("%X", time.localtime())

# To deal with terminal resizes:
signal.signal(signal.SIGWINCH, sigwinch_handler)

if opts.interface:
    output = curses.wrapper(link)
elif opts.statistics == 'ip':
    output = curses.wrapper(statistics_ip)
elif opts.statistics == 'tcpv':
    output = curses.wrapper(statistics_tcpv)

print(output)
sys.exit(0)