#!/usr/bin/python -u
import sys
import time
import re
import threading
import argparse
import os
import subprocess
import traceback
from distutils.spawn import find_executable

sync_start_tollerance = 0.001
killed = threading.Event()
buckets = []
lock = threading.RLock()


def printit():

    global buckets

    # sync all threads
    while (time.time() % 1.0) > sync_start_tollerance:
        pass

    start_time = time.time()

    while not killed.is_set():

        with lock:
            now = time.time()
            buckets_copy = buckets
            buckets = [0]*len(args.ips)

        elapsed_seconds = now - start_time
        bps = [(B/elapsed_seconds)*8 for B in buckets_copy]
        start_time = now

        print "%f,%s" % (now, ",".join(map(str, bps)))

        # sleep for a reporting interval with jitter compensation
        time.sleep(args.rep_seconds - time.time() % args.rep_seconds)


def main():

    global buckets

    # compile regex expression
    regex = re.compile("length (\d+): (\d+\.\d+\.\d+\.\d+)\.\d+")
    # build fast access dictionary to associate ips to list index
    ips_to_idx = {args.ips[i]: i for i in range(len(args.ips))}
    # get tcpdump command
    tcpdump_path = find_executable('tcpdump')
    if not tcpdump_path:
        print >> sys.stderr, 'error: unable to locate tcpdump'
        sys.exit(-1)
    filter_str = 'or '.join(['src '+ip for ip in args.ips])
    cmd = "%s -i %s -l -e -n %s" % (tcpdump_path, args.iface, filter_str)
    # initialize buckets
    buckets = [0]*len(args.ips)
    # start print thread
    threading.Thread(target=printit).start()

    # start tcp dump, parse stdout
    try:
        p = subprocess.Popen(
            cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

        for line in iter(p.stdout.readline, b''):
            m = regex.search(line.rstrip())
            if m:
                with lock:
                    buckets[ips_to_idx[m.group(2)]] += int(m.group(1))

    except (KeyboardInterrupt, SystemExit):
        pass

    except:
        print >> sys.stderr, 'error: unable to parse tcpdump output'
        traceback.print_exc(file=sys.stderr)

    finally:
        p.kill()
        killed.set()
        sys.exit()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='Print rx bitrate by IP hosts as timestamped CSV report')

    parser.add_argument(
        'ips', metavar='ip_addr', type=str,
        nargs='+', help='ip addresses')
    parser.add_argument(
        '-i', metavar='iface', required=True, type=str,
        help='interface to monitor', dest='iface')
    parser.add_argument(
        '-r', metavar='seconds', default='1', type=float,
        help='report interval in seconds', dest='rep_seconds')

    if os.getuid() != 0:
        print >> sys.stderr, 'error: need sudo privileges to run'
        sys.exit(-1)

    args = parser.parse_args()

    main()