#! /usr/bin/env python
# Copyright 2012 Cumulus Networks LLC, all rights reserved
# Copyright 2013,2015,2016,2017 Cumulus Networks, Inc.  All rights reserved.

#####################################################################
#
# netstat -i is a great tool for summarizing network statistics;
# however, it doesn't clear. This script saves away the
# readings at a point in time and subtracts them from the current
# readings to produce a delta. The readings are saved under the
# "/tmp/cl-netstat-<uid>" directory under a filename that is the userid
# by default. However, the user can choose to save a file with a tag
# and use that tag in future invocations to reference those readings.
#
#####################################################################

# Assume that netstat -i output looks as follows.
# Kernel Interface table
# Iface   MTU     RX-OK RX-ERR RX-DRP RX-OVR    TX-OK TX-ERR TX-DRP TX-OVR Flg
# eth0       1500     0      0      0 0             0      0      0      0 BMU
# lo        16436 246427      0      0 0        246427      0      0      0 LRU
# wlan0      1500 311162      0      0 0        196776      0      0      0 BMRU

import argparse
import subprocess
import getopt
import json
import sys
import os.path
import cPickle as pickle
import re
from collections import namedtuple, OrderedDict
from tabulate import tabulate
from network_docopt import sort_ifaces

NStats = namedtuple("NStats", "mtu, rx_ok, rx_err, rx_drop, rx_ovr, tx_ok,\
                    tx_err, tx_drop, tx_ovr, flags")
header = ['Iface', 'MTU', 'RX_OK', 'RX_ERR',
          'RX_DRP', 'RX_OVR', 'TX_OK', 'TX_ERR',
          'TX_DRP', 'TX_OVR', 'Flg']


def cnstat_create_element(idx, netstats, procstats):
    cntr = []
    fields = [netstats[i] for i in [1]] + \
             [procstats[i] for i in [1, 2, 3, 4, 9, 10, 11, 12]] + \
             [netstats[10]]
    cntr = NStats._make(fields)
    return dict({netstats[0]: cntr})

def convert_int(s):
    try:
        val = int(s)
    except ValueError as v:
        val = s
    return val

def table_as_json(table):
    output = {}

    for line in table:
        if_name = line[0]

        # Build a dictionary where the if_name is the key and the value is
        # a dictionary that holds MTU, TX_DRP, etc
        output[if_name] = {
            header[1] : convert_int(line[1]),
            header[2] : convert_int(line[2]),
            header[3] : convert_int(line[3]),
            header[4] : convert_int(line[4]),
            header[5] : convert_int(line[5]),
            header[6] : convert_int(line[6]),
            header[7] : convert_int(line[7]),
            header[8] : convert_int(line[8]),
            header[9] : convert_int(line[9]),
            header[10] : line[10]
            }
    return json.dumps(output, indent=4, sort_keys=True)


def cnstat_print(cnstat_dict, use_json):
    table = []

    for if_name in sort_ifaces(cnstat_dict.keys()):
        data = cnstat_dict.get(if_name)
        table.append((if_name, data.mtu,
                      data.rx_ok, data.rx_err,
                      data.rx_drop, data.rx_ovr,
                      data.tx_ok, data.tx_err,
                      data.tx_drop, data.tx_err,
                      data.flags))

    if use_json:
        print table_as_json(table)

    else:
        print '\n' + netstat_lines[0]
        print tabulate(table, header, tablefmt='simple') + '\n'


def ns_diff(newstr, oldstr):
    new, old = int(newstr), int(oldstr)

    if new >= old:
        return str((new - old))
    else:
        return 'Unknown'


def cnstat_diff_print(cnstat_new_dict, cnstat_old_dict, use_json):
    table = []

    for if_name, cntr in cnstat_new_dict.iteritems():
        old_cntr = None
        if if_name in cnstat_old_dict:
            old_cntr = cnstat_old_dict.get(if_name)

        if old_cntr is not None:
            table.append((if_name, cntr.mtu,
                          ns_diff(cntr.rx_ok, old_cntr.rx_ok),
                          ns_diff(cntr.rx_err, old_cntr.rx_err),
                          ns_diff(cntr.rx_drop, old_cntr.rx_drop),
                          ns_diff(cntr.rx_ovr, old_cntr.rx_ovr),
                          ns_diff(cntr.tx_ok, old_cntr.tx_ok),
                          ns_diff(cntr.tx_err, old_cntr.tx_err),
                          ns_diff(cntr.tx_drop, old_cntr.tx_drop),
                          ns_diff(cntr.tx_ovr, old_cntr.tx_ovr),
                          cntr.flags))
        else:
            table.append((if_name, cntr.mtu,
                          cntr.rx_ok,
                          cntr.rx_err,
                          cntr.rx_drop,
                          cntr.rx_ovr,
                          cntr.tx_ok,
                          cntr.tx_err,
                          cntr.tx_drop,
                          cntr.tx_err,
                          cntr.flags))

    if use_json:
        print table_as_json(table)
    else:
        print '\n' + netstat_lines[0]
        print tabulate(table, header, tablefmt='simple') + '\n'

if __name__ == "__main__":
    parser  = argparse.ArgumentParser(description='Wrapper for netstat',
                                      version='1.0.2',
                                      formatter_class=argparse.RawTextHelpFormatter,
                                      epilog="""
Note: Clearing stats does not affect hardware or software values.
      cl-netstat saves the current stats when -c is given, so they can
      be compared with later values.  The -c and -d options are per user (UID)
      by default.  Use the -t TAG option to change this behavior.  You must
      use the same -t TAG value with subsequent commands to get valid results.

Examples:
  cl-netstat -c -t test
  cl-netstat -t test
  cl-netstat -d -t test
  cl-netstat
  cl-netstat -r
""")
    parser.add_argument('-c', '--clear', action='store_true', help='Copy & clear stats per user (tag)')
    parser.add_argument('-d', '--delete', action='store_true', help='Delete saved stats, either the uid or the specified tag')
    parser.add_argument('-D', '--delete-all', action='store_true', help='Delete all saved stats')
    parser.add_argument('-j', '--json', action='store_true', help='Display in JSON format')
    parser.add_argument('-r', '--raw', action='store_true', help='Raw stats (unmodified output of netstat)')
    parser.add_argument('-t', '--tag', type=str, help='Save stats with name TAG', default=None)
    parser.add_argument('--clear-interface', type=str, default=None, help='clear stats for a single interface')
    args = parser.parse_args()

    save_fresh_stats = args.clear
    save_fresh_stats_single_interface = args.clear_interface
    delete_saved_stats = args.delete
    delete_all_stats = args.delete_all
    use_json = args.json
    raw_stats = args.raw
    tag_name = args.tag
    uid = str(os.getuid())

    if tag_name is not None:
        cnstat_file = uid + "-" + tag_name
    else:
        cnstat_file = uid

    cnstat_dir = "/tmp/cl-netstat-" + uid
    cnstat_fqn_file = cnstat_dir + "/" + cnstat_file

    if delete_all_stats:
        if os.path.exists(cnstat_dir):
            for file in os.listdir(cnstat_dir):
                os.remove(cnstat_dir + "/" + file)

            try:
                os.rmdir(cnstat_dir)
            except IOError as e:
                print e.errno, e
                sys.exit(e)
        sys.exit(0)

    if delete_saved_stats:
        try:
            os.remove(cnstat_fqn_file)
        except IOError as e:
            if e.errno != ENOENT:
                print e.errno, e
                sys.exit(1)
        finally:
            if os.path.exists(cnstat_dir) and os.listdir(cnstat_dir) == []:
                os.rmdir(cnstat_dir)
            sys.exit(0)

    try:
        cmd = ['/bin/netstat', '-i']

        # If we are clearing counters get the netstat output for all interfaces
        if save_fresh_stats or save_fresh_stats_single_interface:
            cmd.append('-all')

        netstat_out = subprocess.Popen(cmd,
                                       stdout=subprocess.PIPE,
                                       shell=False).communicate()[0]
    except EnvironmentError as e:
        print e, e.errno
        sys.exit(e.errno)

    netstat_lines = netstat_out.split("\n")

    # Since netstat -i returns some stats as 32-bits, get full 64-bit
    # stats from /prov/net/dev and display only the 64-bit stats.
    try:
        proc_out = subprocess.Popen((['/bin/cat', '/proc/net/dev']),
                                    stdout=subprocess.PIPE,
                                    shell=False).communicate()[0]
    except EnvironmentError as e:
        print e, e.errno
        sys.exit(e.errno)

    proc = {}
    for line in proc_out.split("\n"):
        parsed = re.findall("\s*([^ ]+):(.*)", line)
        if not parsed:
            continue
        iface, stats = parsed[0]
        proc[iface] = stats.split()

    # At this point, either we'll create a file or open an existing one.
    if not os.path.exists(cnstat_dir):
        try:
            os.makedirs(cnstat_dir)
        except IOError as e:
            print e.errno, e
            sys.exit(1)

    # Build a dictionary of the stats
    cnstat_dict = OrderedDict()

    # Populate cnstat_dict with the current state saved to cnstat_fqn_file.
    # We will update cnstat_fqn_file with the current counters for the
    # interface we are clearing.
    if save_fresh_stats_single_interface:
        if os.path.exists(cnstat_fqn_file):
            cnstat_dict = pickle.load(open(cnstat_fqn_file, 'r'))

    # We skip the first 2 lines since they contain no interface information
    for i in range(2, len(netstat_lines) - 1):
        netstats = netstat_lines[i].split()
        if ":" in netstats[0]:
            continue    # skip aliased interfaces

        if save_fresh_stats_single_interface is None or netstats[0] == save_fresh_stats_single_interface:
            procstats = proc.get(netstats[0])
            cnstat_dict.update(cnstat_create_element(i, netstats, procstats))

    # Now decide what information to display
    if raw_stats:
        cnstat_print(cnstat_dict, use_json)
        sys.exit(0)

    if save_fresh_stats or save_fresh_stats_single_interface:
        try:
            pickle.dump(cnstat_dict, open(cnstat_fqn_file, 'w'))
        except IOError as e:
            sys.exit(e.errno)
        else:
            if save_fresh_stats_single_interface:
                print "Cleared counters for %s" % save_fresh_stats_single_interface
            else:
                print "Cleared counters"
            sys.exit(0)

    cnstat_cached_dict = OrderedDict()

    if os.path.isfile(cnstat_fqn_file):
        try:
            cnstat_cached_dict = pickle.load(open(cnstat_fqn_file, 'r'))
            cnstat_diff_print(cnstat_dict, cnstat_cached_dict, use_json)
        except IOError as e:
            print e.errno, e
    else:
        if tag_name:
            print "\nFile '%s' does not exist" % cnstat_fqn_file
            print "Did you run 'cl-netstat -c -t %s' to record the counters via tag %s?\n" % (tag_name, tag_name)
        else:
            cnstat_print(cnstat_dict, use_json)
