#!/usr/bin/env python
#
# Record network device counters every 1 second and write in RRD or plain text.
#
# Copyright 2016 Ludovic Pouzenc <ludovic@pouzenc.fr>
#
# This file is part of if_rrd_fast.
#
# if_rrd_fast 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, either version 3 of the License, or
# (at your option) any later version.
#
# if_rrd_fast 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 if_rrd_fast.  If not, see <http://www.gnu.org/licenses/>.
#
import sys,os,signal,argparse
from time import time,sleep;
from math import ceil;
from pyrrd.rrd import DataSource, RRA, RRD


def signal_handler(signal, frame):
    global exit_flag
    exit_flag=True

def output_text(f_outpout_text, values):
    f_outpout_text.writelines('%i:%s\n'%(ts,v) for (ts,v) in values)

if __name__ == "__main__":
    # Workaround for rrd bug about float parsing and locale on HWPREDICT
    os.environ['LC_NUMERIC']='C'

    # Argument parsing
    parser = argparse.ArgumentParser(description='Record network device counters every 1 second and write in RRD or plain text.')
    parser.add_argument('-v', '--verbose', action='store_true', help='verbose on stdout, showing what the program is doing')
    parser.add_argument('-o', '--output', metavar='RRDFILE', default='./out.rrd', help='put acquired values in RRDFILE')
    parser.add_argument('-O', '--text-output', metavar='TEXTFILE', default='./out.txt', help='put acquired values in TEXTFILE')
    parser.add_argument('-I', '--import-first', metavar='TEXTFILE', help='import values from TEXTFILE before acquiring new data')
    parser.add_argument('iface', help='read counters from this network interface')
    args = parser.parse_args()

    # Signal setup
    exit_flag=False
    signal.signal(signal.SIGINT, signal_handler)

    # Init input files   
    sysfs_dir='/sys/class/net/' + args.iface + '/statistics/'
    if args.verbose: print "Using %s..." % sysfs_dir
    f_rx_bytes = open(sysfs_dir + 'rx_bytes')
    f_tx_bytes = open(sysfs_dir + 'tx_bytes')

    # Guess the start timestamp of RRD from import file
    start=None
    if args.import_first:
        f_import = open(args.import_first)
        try:
            (first_ts,first_data) = f_import.next().rstrip().split(':',1)
            start = int(first_ts)-1
        except:
            print "Skipping import : can't find the first timestamp at begining of file"
        f_import.seek(0)

    # Open / create output files
    if args.verbose: print "Initializing %s..." % args.output
    dataSources = []
    roundRobinArchives = []
    dataSources.append(DataSource(dsName='rx', dsType='COUNTER', heartbeat=2, minval=0, maxval=1e9))
    dataSources.append(DataSource(dsName='tx', dsType='COUNTER', heartbeat=2, minval=0, maxval=1e9))
    roundRobinArchives.append(RRA(cf='AVERAGE', xff=0.5, steps=1, rows=1152))
    roundRobinArchives.append(RRA(cf='AVERAGE', xff=0.5, steps=60, rows=1152))
    roundRobinArchives.append(RRA(cf='AVERAGE', xff=0.5, steps=300, rows=1152))
    roundRobinArchives.append(RRA(cf='AVERAGE', xff=0.5, steps=3600, rows=1152))
    roundRobinArchives.append(RRA(cf='MAX', xff=0.5, steps=1, rows=1152))
    roundRobinArchives.append(RRA(cf='MAX', xff=0.5, steps=60, rows=1152))
    roundRobinArchives.append(RRA(cf='MAX', xff=0.5, steps=300, rows=1152))
    roundRobinArchives.append(RRA(cf='MAX', xff=0.5, steps=3600, rows=1152))
    roundRobinArchives.append(RRA(cf='MIN', xff=0.5, steps=1, rows=1152))
    roundRobinArchives.append(RRA(cf='MIN', xff=0.5, steps=60, rows=1152))
    roundRobinArchives.append(RRA(cf='MIN', xff=0.5, steps=300, rows=1152))
    roundRobinArchives.append(RRA(cf='MIN', xff=0.5, steps=3600, rows=1152))
    myRRD = RRD(filename=args.output, start=start, step=1, ds=dataSources, rra=roundRobinArchives)
    myRRD.create()

    # Import data
    if args.import_first:
        if start:
            print "Importing %s..." % args.import_first
            lineno=0
            for line in f_import:
                lineno = lineno + 1;
                myRRD.bufferValue(line.rstrip())
                if (lineno%300==0): myRRD.update()
            myRRD.update()
        f_import.close()
        print "Imported %i PDP" % lineno
        
    # Open text output file (after import because could be the same file)
    if args.verbose: print "Opening %s..." % args.text_output
    f_outpout_text = open(args.text_output, 'a')

   # Main data acquire/write loop
    while(exit_flag==False):
	# Try to align out self on next second start
	ts=time();
	ts2=ceil(ts);
	sleep(ts2-ts); # FIXME traped signals could make things messy
        # Acquire values
        f_rx_bytes.seek(0)
        f_tx_bytes.seek(0)
        rx=f_rx_bytes.read().rstrip()
        tx=f_tx_bytes.read().rstrip()
        if args.verbose: print "at %i : %s/%s" % (ts,rx,tx)
        # Buffer and periodically write down values in RRD and text files
        myRRD.bufferValue(ts,rx,tx) 
        if (ts2%300==0): 
            output_text(f_outpout_text, myRRD.values)
            myRRD.update()

    # Cleanup for exiting
    if len(myRRD.values):
        output_text(f_outpout_text, myRRD.values)
        myRRD.update()

    f_rx_bytes.close()
    f_tx_bytes.close()
    f_outpout_text.close()

    print "Last timestamp : %d" % ts