#!/usr/bin/env python3

import configparser
import optparse
import os
import re
import rrdtool
import subprocess
import time

def main():
    conf = configparser.ConfigParser()

    conflist = []
    if ('HOME' in os.environ):
        conflist.append(os.path.expanduser('~/.speedcheck.ini'))
    conf.read(conflist)

    usage = '''%prog [options]'''
    parser = optparse.OptionParser(usage)
    parser.add_option('-v', '--verbose', action="store_true", default=False,
                    help="Enable debug output")
    parser.add_option('-r', '--rrd', action="store", help="Path to RRD")
    parser.add_option('-g', '--graphdir', action="store", help="Directory for graphs")

    (opts, args) = parser.parse_args()
    if opts.rrd == None:
        if conf.has_option('global', 'rrd'):
            opts.rrd = conf.get('global', 'rrd')
        else:
            parser.error('Path to RRD must be specified in either the ini or on the command line')

    if opts.graphdir == None:
        if conf.has_option('global', 'graphdir'):
            opts.graphdir = conf.get('global', 'graphdir')
        else:
            parser.error('Graph directory must be specified in either the ini or on the command line')

    if opts.verbose:
        print('Fetching stats...')
    stats = fetchstats(conf)
    if opts.verbose:
        print(stats)
    if opts.verbose:
        print('Updating RRD')
    updaterrd(opts.rrd, stats)
    if opts.verbose:
        print('Updating graph')
    graphrrd(opts.rrd, opts.graphdir)

def fetchstats(conf):
    stats = {}
    if conf.has_option('global', 'neardl'):
        stats['neardl'] = testdl(conf.get('global', 'neardl'))
    if conf.has_option('global', 'nearul'):
        stats['nearul'] = testul(conf.get('global', 'nearul'))
    if conf.has_option('global', 'nearping'):
        stats['nearpl'], stats['nearlat'] = testping(conf.get('global', 'nearping'))
    if conf.has_option('global', 'fardl'):
        stats['fardl'] = testdl(conf.get('global', 'fardl'))
    if conf.has_option('global', 'farul'):
        stats['farul'] = testul(conf.get('global', 'farul'))
    if conf.has_option('global', 'farping'):
        stats['farpl'], stats['farlat'] = testping(conf.get('global', 'farping'))

    return stats

def testdl(url):
    p = subprocess.Popen(['curl', '-w', '%{speed_download}', '-so', '/dev/null', url], stdout = subprocess.PIPE)
    speed, xxx = p.communicate()
    if p.returncode != 0:
        print('Error %d fetching \'%s\'' % (p.returncode, url))
        return None
    return float(speed) * 8.0 / 1024.0 # convert to kbit/sec

def testping(host):
    p = subprocess.Popen(['ping', '-c', '5', '-t', '8', '-q', host], stdout = subprocess.PIPE)
    stdout, stderr = p.communicate()
    l = stdout.decode('ascii', 'ignore').split('\n')
    if len(l) != 6:
        print('Unable to parse ping line:', l)
    xx, xx, xx, plossline, latline, xx = l
    ploss = float(re.match('.* received, ([0-9.]+)% packet loss', plossline).groups()[0])
    latency = float(re.match('.*stddev = [0-9.]+/([0-9.]+)/.* ms', latline).groups()[0])
    return ploss, latency

def createrrd(rrdname):
    # Create RRD for upstream/downstream speed, packet loss and
    # latency for near and far site
    # Do a test every half and hour
    # Average 2 for hourly stats (keep 168 - a weeks worth)
    # Average 48 for hourly stats (keep 1825 - 5 years worth)
    # Detemine  minimum & maximum for an hour and keep a weeks worth.
    rrdtool.create(rrdname,
                   '--step', '300',
                   'DS:neardl:GAUGE:3600:0:U',
                   'DS:nearul:GAUGE:3600:0:U',
                   'DS:nearpl:GAUGE:3600:0:100',
                   'DS:nearlat:GAUGE:3600:0:U',
                   'DS:fardl:GAUGE:3600:0:U',
                   'DS:farul:GAUGE:3600:0:U',
                   'DS:farpl:GAUGE:3600:0:100',
                   'DS:farlat:GAUGE:3600:0:U',
                   'RRA:AVERAGE:0.1:2:168',
                   'RRA:AVERAGE:0.1:48:1825',
                   'RRA:MIN:0.1:2:168',
                   'RRA:MAX:0.1:2:168',
                   )

def updaterrd(rrdname, stats):
    try:
        os.stat(rrdname)
    except OSError as e:
        if e.errno == 2:
            print('Creating RRD...')
            createrrd(rrdname)
    s = '%d:' % (int(time.time()))
    for a in ['neardl', 'nearul', 'nearpl', 'nearlat', 'fardl', 'farul', 'farpl', 'farlat']:
        if a in stats:
            s += '%f:' % (stats[a])
        else:
            s += 'U:'
    s = s[0:-1]
    rrdtool.update(rrdname, s)

def graphrrd(rrdname, graphdir):
    latencyargs = (
        '-a', 'SVG',
        '--vertical-label', 'milliseconds',

        'DEF:nearlat=%s:nearlat:AVERAGE' % rrdname,
        'DEF:nearlatmin=%s:nearlat:MIN' % rrdname,
        'DEF:nearlatmax=%s:nearlat:MAX' % rrdname,
        'CDEF:nearlatdif=nearlatmax,nearlatmin,-',

        'LINE0.001:nearlatmin#000000:',
        'AREA:nearlatdif#00dc76::STACK',
        'LINE1:nearlatmax#00ff00:Near latency',

        'DEF:farlat=%s:farlat:AVERAGE' % rrdname,
        'DEF:farlatmin=%s:farlat:MIN' % rrdname,
        'DEF:farlatmax=%s:farlat:MAX' % rrdname,
        'CDEF:farlatdif=farlatmax,farlatmin,-',

        'LINE0.001:farlatmin#000000:',
        'AREA:farlatdif#dc0076::STACK',
        'LINE1:farlatmax#ff0000:Far packetloss',

        'DEF:nearpl=%s:nearpl:AVERAGE' % rrdname,
        'LINE1:nearpl#0000ff:Near packet loss (%)',

        'DEF:farpl=%s:farpl:AVERAGE' % rrdname,
        'LINE1:nearpl#ffff00:Far packet loss (%)',
        )
    rrdtool.graph('%s/latency-hour-link.svg' % (graphdir),
                  '--width', '768',
                  '--height', '256',
                  '--start', 'end - 7d',
                  '--end', 'now',
                  *latencyargs)
    rrdtool.graph('%s/latency-daily-link.svg' % (graphdir),
                  '--width', '768',
                  '--height', '256',
                  '--start', 'end - 365d',
                  '--end', 'now',
                  *latencyargs)

    bwargs = (
        '-a', 'SVG',
        '-X', '0',
        '--vertical-label', 'kbit/sec',

        'DEF:neardl=%s:neardl:AVERAGE' % rrdname,
        'DEF:neardlmin=%s:neardl:MIN' % rrdname,
        'DEF:neardlmax=%s:neardl:MAX' % rrdname,
        'CDEF:neardldif=neardlmax,neardlmin,-',

        'LINE0.001:neardlmin#000000:',
        'AREA:neardldif#00dc76::STACK',
        'LINE1:neardlmax#00ff00:Near download',

        'DEF:fardl=%s:fardl:AVERAGE' % rrdname,
        'DEF:fardlmin=%s:fardl:MIN' % rrdname,
        'DEF:fardlmax=%s:fardl:MAX' % rrdname,
        'CDEF:fardldif=fardlmax,fardlmin,-',

        'LINE0.001:fardlmin#000000:',
        'AREA:fardldif#dc0076::STACK',
        'LINE1:fardlmax#ff0000:Far download',
        )
    rrdtool.graph('%s/bw-hour-link.svg' % (graphdir),
                  '--width', '768',
                  '--height', '256',
                  '--start', 'end - 1d',
                  '--end', 'now',
                  *bwargs)
    rrdtool.graph('%s/bw-daily-link.svg' % (graphdir),
                  '--width', '768',
                  '--height', '256',
                  '--start', 'end - 7d',
                  '--end', 'now',
                  *bwargs)

if __name__ == '__main__':
    main()